Models module¶

This module provides implementations of common model architectures used in model-based RL, including probabilistic and deterministic ensembles. All models in the library derive from class mbrl.models.Model. We provide a generic ensemble implementation, mbrl.models.BasicEnsemble, that can be used to produce epistemic uncertainty estimates for any subclass of Model. For efficiency considerations, some specific model implementations also provide their own ensemble implementations, without having to rely on BasicEnsemble. One such model is mbrl.models.GaussianMLP, which can be used as a single model or as an ensemble. Additionally, it can be used as a deterministic model trained with MSE loss, or a parameterized Gaussian with mean and log variance outputs, trained with negative log-likelihood.

class mbrl.models.BasicEnsemble(ensemble_size: int, device: Union[str, torch.device], member_cfg: omegaconf.dictconfig.DictConfig, propagation_method: Optional[str] = None)¶

Bases: mbrl.models.model.Ensemble

Implements an ensemble of bootstrapped models.

This model is a basic implementation of the ensemble of bootstrapped models described in the Chua et al., NeurIPS 2018 paper (PETS) https://arxiv.org/pdf/1805.12114.pdf, and includes support for different uncertainty propagation options (see forward()). The underlying model can be any subclass of mbrl.models.Model, and the ensemble forward simply loops over all models during the forward and backward pass (hence the term basic).

All members of the ensemble will be identical, and they must be subclasses of mbrl.models.Model. This method assumes that the models have an attribute model.deterministic that indicates if the model is deterministic or not.

Members can be accessed using ensemble[i], to recover the i-th model in the ensemble. Doing len(ensemble) returns its size, and the ensemble can also be iterated over the models (e.g., calling for i, model in enumerate(ensemble).

Valid propagation options are:

“random_model”: for each output in the batch a model will be chosen at random. This corresponds to TS1 propagation in the PETS paper.

“fixed_model”: for output j-th in the batch, the model will be chosen according to the model index in propagation_indices[j]. This can be used to implement TSinf propagation, described in the PETS paper.

“expectation”: the output for each element in the batch will be the mean across models.

Parameters

ensemble_size (int) – how many models to include in the ensemble.
device (str or torch.device) – the device to use for the model.
member_cfg (omegaconf.DictConfig) – the configuration needed to instantiate the models in the ensemble. They will be instantiated using hydra.utils.instantiate(member_cfg).
propagation_method (str, optional) – the uncertainty propagation method to use (see above). Defaults to None.

eval_score(model_in: torch.Tensor, target: Optional[torch.Tensor] = None) → torch.Tensor¶

Computes the average score over all members given input/target.

The input and target tensors are replicated once for each model in the ensemble.

Parameters

model_in (tensor) – the inputs to the models.
target (tensor) – the expected output for the given inputs.

Returns

the average score over all models.

Return type

(tensor)

forward(x: torch.Tensor, rng: Optional[torch._C.Generator] = None) → Tuple[torch.Tensor, torch.Tensor]¶

Computes the output of the ensemble.

The forward pass for the ensemble computes forward passes for of its models, and aggregates the prediction in different ways, according to the desired epistemic uncertainty propagation method.

If no propagation is desired (i.e., self.propagation_method is None), then the outputs of the model are stacked into single tensors (one for mean, one for logvar). The shape of each output tensor will then be E x B x D, where E, B and D represent ensemble size, batch size, and output dimension, respectively.

For all other propagation options, the output is of size B x D.

Parameters

x (tensor) – the input to the models (shape B x D). The input will be evaluated over all models, then aggregated according to propagation, as explained above.
rng (torch.Generator, optional) – random number generator to use for “random_model” propagation.

Returns

one for aggregated mean predictions, and one for aggregated log variance prediction (or None if the ensemble members don’t predict variance).

Return type

(tuple of two tensors)

loss(model_ins: Sequence[torch.Tensor], targets: Optional[Sequence[torch.Tensor]] = None) → torch.Tensor¶

Computes average loss over the losses of all members of the ensemble.

Parameters

model_ins (sequence of tensors) – one input for each model in the ensemble.
targets (sequence of tensors) – one target for each model in the ensemble.

Returns

the average loss over all members.

Return type

(tensor)

reset(x: torch.Tensor, rng: Optional[torch._C.Generator] = None) → torch.Tensor¶

Initializes any internal dependent state when using the model for simulation.

Initializes model indices for “fixed_model” propagation method a bootstrapped ensemble with TSinf propagation).

Parameters

x (tensor) – the input to the model.
rng (random number generator) – a rng to use for sampling the model indices.

Returns

forwards the same input.

Return type

(tensor)

set_elite(elite_models: Sequence[int])¶: For ensemble models, indicates if some models should be considered elite.

class mbrl.models.Ensemble(num_members: int, device: Union[str, torch.device], propagation_method: str, deterministic: bool = False, *args, **kwargs)¶

Bases: mbrl.models.model.Model, abc.ABC

Base abstract class for all ensemble of bootstrapped models.

Implements an ensemble of bootstrapped models described in the Chua et al., NeurIPS 2018 paper (PETS) https://arxiv.org/pdf/1805.12114.pdf,

Uncertainty propagation methods are available that can be used to aggregate the outputs of the different models in the ensemble. Valid propagation options are:

“random_model”: for each output in the batch a model will be chosen at random. This corresponds to TS1 propagation in the PETS paper.

“fixed_model”: for output j-th in the batch, the model will be chosen according to the model index in propagation_indices[j]. This can be used to implement TSinf propagation, described in the PETS paper.

“expectation”: the output for each element in the batch will be the mean across models.

The default value of None indicates that no uncertainty propagation, and the forward method returns all outputs of all models.

Subclasses of Ensemble are responsible for implementing the above functionality.

Parameters

num_members (int) – how many models in the ensemble.
device (str or torch.device) – device to use for the model.
propagation_method (str, optional) – the uncertainty propagation method to use (see above). Defaults to None.
deterministic (bool) – if True, the model will be trained using MSE loss and no logvar prediction will be done. Defaults to False.

abstract eval_score(model_in: Union[torch.Tensor, mbrl.types.TransitionBatch], target: Optional[torch.Tensor] = None) → torch.Tensor¶

Computes an evaluation score for the model over the given input/target.

This method should compute a non-reduced score for the model, intended mostly for logging/debugging purposes (so, it should not keep gradient information). For example, the following could be a valid implementation of eval_score:

with torch.no_grad():
    return torch.functional.mse_loss(model(model_in), target, reduction="none")

Parameters

model_in (tensor or batch of transitions) – the inputs to the model.
target (tensor or sequence of tensors) – the expected output for the given inputs, if it cannot be computed from model_in.

Returns

a non-reduced tensor score.

Return type

(tensor)

forward(x: Union[torch.Tensor, mbrl.types.TransitionBatch], **kwargs) → Tuple[torch.Tensor, …]¶

Computes the output of the dynamics model.

Parameters: x (tensor or batch of transitions) – the input to the model.
Returns: all tensors predicted by the model (e.g., .mean and logvar).
Return type: (tuple of tensors)

abstract loss(model_in: Union[torch.Tensor, mbrl.types.TransitionBatch], target: Optional[torch.Tensor] = None) → torch.Tensor¶

Computes a loss that can be used to update the model using backpropagation.

Parameters

model_in (tensor or batch of transitions) – the inputs to the model.
target (tensor, optional) – the expected output for the given inputs, if it cannot be computed from model_in.

Returns

a loss tensor.

Return type

(tensor)

sample(x: Union[torch.Tensor, mbrl.types.TransitionBatch], deterministic: bool = False, rng: Optional[torch._C.Generator] = None, **kwargs) → Tuple[torch.Tensor, …]¶

Samples an output of the dynamics model from the modeled Gaussian.

Parameters

x (tensor) – the input to the model.
deterministic (bool) – if True, the model returns a deterministic “sample” (e.g., the mean prediction). Defaults to False.
rng (random number generator) – a rng to use for sampling.

Returns

the sampled output.

Return type

(tensor)

set_elite(elite_models: Sequence[int])¶: For ensemble models, indicates if some models should be considered elite.

class mbrl.models.EnsembleLinearLayer(num_members: int, in_size: int, out_size: int, bias: bool = True)¶

Bases: torch.nn.modules.module.Module

Efficient linear layer for ensemble models.

extra_repr() → str¶

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x)¶

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class mbrl.models.GaussianMLP(in_size: int, out_size: int, device: Union[str, torch.device], num_layers: int = 4, ensemble_size: int = 1, hid_size: int = 200, use_silu: bool = False, deterministic: bool = False, propagation_method: Optional[str] = None, learn_logvar_bounds: bool = False)¶

Bases: mbrl.models.model.Ensemble

Implements an ensemble of multi-layer perceptrons each modeling a Gaussian distribution.

This model corresponds to a Probabilistic Ensemble in the Chua et al., NeurIPS 2018 paper (PETS) https://arxiv.org/pdf/1805.12114.pdf

It predicts per output mean and log variance, and its weights are updated using a Gaussian negative log likelihood loss. The log variance is bounded between learned min_log_var and max_log_var parameters, trained as explained in Appendix A.1 of the paper.

This class can also be used to build an ensemble of GaussianMLP models, by setting ensemble_size > 1 in the constructor. Then, a single forward pass can be used to evaluate multiple independent MLPs at the same time. When this mode is active, the constructor will set self.num_members = ensemble_size.

For the ensemble variant, uncertainty propagation methods are available that can be used to aggregate the outputs of the different models in the ensemble. Valid propagation options are:

“random_model”: for each output in the batch a model will be chosen at random. This corresponds to TS1 propagation in the PETS paper.

“fixed_model”: for output j-th in the batch, the model will be chosen according to the model index in propagation_indices[j]. This can be used to implement TSinf propagation, described in the PETS paper.

“expectation”: the output for each element in the batch will be the mean across models.

The default value of None indicates that no uncertainty propagation, and the forward method returns all outputs of all models.

Parameters

in_size (int) – size of model input.
out_size (int) – size of model output.
device (str or torch.device) – the device to use for the model.
num_layers (int) – the number of layers in the model (e.g., if num_layers == 3, then model graph looks like input -h1-> -h2-> -l3-> output).
ensemble_size (int) – the number of members in the ensemble. Defaults to 1.
hid_size (int) – the size of the hidden layers (e.g., size of h1 and h2 in the graph above).
use_silu (bool) – if True, hidden layers will use SiLU activations, otherwise ReLU activations will be used. Defaults to False.
deterministic (bool) – if True, the model will be trained using MSE loss and no logvar prediction will be done. Defaults to False.
propagation_method (str, optional) – the uncertainty propagation method to use (see above). Defaults to None.
learn_logvar_bounds (bool) – if True, the logvar bounds will be learned, otherwise they will be constant. Defaults to False.

eval_score(model_in: torch.Tensor, target: Optional[torch.Tensor] = None) → torch.Tensor¶

Computes the squared error for the model over the given input/target.

When model is not an ensemble, this is equivalent to F.mse_loss(model(model_in, target), reduction=”none”). If the model is ensemble, then return is batched over the model dimension.

Parameters

model_in (tensor) – input tensor. The shape must be B x Id, where B` and Id batch size, and input dimension, respectively.
target (tensor) – target tensor. The shape must be B x Od, where B and Od represent batch size, and output dimension, respectively.

Returns

a tensor with the squared error per output dimension, batched over model.

Return type

(tensor)

forward(x: torch.Tensor, rng: Optional[torch._C.Generator] = None, use_propagation: bool = True) → Tuple[torch.Tensor, torch.Tensor]¶

Computes mean and logvar predictions for the given input.

When self.num_members > 1, the model supports uncertainty propagation options that can be used to aggregate the outputs of the different models in the ensemble. Valid propagation options are:

“random_model”: for each output in the batch a model will be chosen at random. This corresponds to TS1 propagation in the PETS paper.

“fixed_model”: for output j-th in the batch, the model will be chosen according to the model index in propagation_indices[j]. This can be used to implement TSinf propagation, described in the PETS paper.

“expectation”: the output for each element in the batch will be the mean across models.

If a set of elite models has been indicated (via set_elite()), then all propagation methods will operate with only on the elite set. This has no effect when propagation is None, in which case the forward pass will return one output for each model.

Parameters

x (tensor) –
the input to the model. When self.propagation is None, the shape must be E x B x Id or B x Id, where E, B and Id represent ensemble size, batch size, and input dimension, respectively. In this case, each model in the ensemble will get one slice from the first dimension (e.g., the i-th ensemble member gets x[i]).

For other values of self.propagation (and use_propagation=True), the shape must be B x Id.
rng (torch.Generator, optional) – random number generator to use for “random_model” propagation.
use_propagation (bool) – if False, the propagation method will be ignored and the method will return outputs for all models. Defaults to True.

Returns

the predicted mean and log variance of the output. If propagation is not None, the output will be 2-D (batch size, and output dimension). Otherwise, the outputs will have shape E x B x Od, where Od represents output dimension.

Return type

(tuple of two tensors)

Note

For efficiency considerations, the propagation method used by this class is an approximate version of that described by Chua et al. In particular, instead of sampling models independently for each input in the batch, we ensure that each model gets exactly the same number of samples (which are assigned randomly with equal probability), resulting in a smaller batch size which we use for the forward pass. If this is a concern, consider using propagation=None, and passing the output to mbrl.util.math.propagate().

loss(model_in: torch.Tensor, target: Optional[torch.Tensor] = None) → torch.Tensor¶

Computes Gaussian NLL loss.

It also includes terms for max_logvar and min_logvar with small weights, with positive and negative signs, respectively.

Parameters

model_in (tensor) – input tensor. The shape must be E x B x Id, or B x Id where E, B and Id represent ensemble size, batch size, and input dimension, respectively.
target (tensor) – target tensor. The shape must be E x B x Id, or B x Od where E, B and Od represent ensemble size, batch size, and output dimension, respectively.

Returns

a loss tensor representing the Gaussian negative log-likelihood of the model over the given input/target. If the model is an ensemble, returns the average over all models.

Return type

(tensor)

reset(x: torch.Tensor, rng: Optional[torch._C.Generator] = None) → torch.Tensor¶

Initializes any internal dependent state when using the model for simulation.

Initializes model indices for “fixed_model” propagation method a bootstrapped ensemble with TSinf propagation).

Parameters

x (tensor) – the input to the model.
rng (random number generator) – a rng to use for sampling the model indices.

Returns

forwards the same input.

Return type

(tensor)

set_elite(elite_indices: Sequence[int])¶: For ensemble models, indicates if some models should be considered elite.

class mbrl.models.Model(*args, **kwargs)¶

Bases: torch.nn.modules.module.Module, abc.ABC

Base abstract class for all dynamics models.

All classes derived from Model must implement the following methods:

forward: computes the model output.

loss: computes a loss tensor that can be used for backpropagation.

eval_score: computes a non-reduced tensor that gives an evaluation score for the model on the input data (e.g., squared error per element).

save: saves the model to a given path.

load: loads the model from a given path.

Subclasses may also want to overrides sample() and reset().

Parameters: device (str or torch.device) – device to use for the model.

abstract eval_score(model_in: Union[torch.Tensor, mbrl.types.TransitionBatch], target: Optional[torch.Tensor] = None) → torch.Tensor¶

Computes an evaluation score for the model over the given input/target.

This method should compute a non-reduced score for the model, intended mostly for logging/debugging purposes (so, it should not keep gradient information). For example, the following could be a valid implementation of eval_score:

with torch.no_grad():
    return torch.functional.mse_loss(model(model_in), target, reduction="none")

Parameters

model_in (tensor or batch of transitions) – the inputs to the model.
target (tensor or sequence of tensors) – the expected output for the given inputs, if it cannot be computed from model_in.

Returns

a non-reduced tensor score.

Return type

(tensor)

forward(x: Union[torch.Tensor, mbrl.types.TransitionBatch], **kwargs) → Tuple[torch.Tensor, …]¶

Computes the output of the dynamics model.

Parameters: x (tensor or batch of transitions) – the input to the model.
Returns: all tensors predicted by the model (e.g., .mean and logvar).
Return type: (tuple of tensors)

load(path: Union[str, pathlib.Path])¶: Loads the model from the given path.

abstract loss(model_in: Union[torch.Tensor, mbrl.types.TransitionBatch], target: Optional[torch.Tensor] = None) → torch.Tensor¶

Computes a loss that can be used to update the model using backpropagation.

Parameters

model_in (tensor or batch of transitions) – the inputs to the model.
target (tensor, optional) – the expected output for the given inputs, if it cannot be computed from model_in.

Returns

a loss tensor.

Return type

(tensor)

reset(x: Union[torch.Tensor, mbrl.types.TransitionBatch], **kwargs) → torch.Tensor¶

Initializes any internal dependent state when using the model for simulation.

For most models this just returns the same input that is given as input. However, for some models this method can be used to initialize data that should be kept constant during a simulated trajectory (for example model indices when using a bootstrapped ensemble with TSinf propagation). It can also be used to return latent states computed by the model.

Parameters: x (tensor or batch of transitions) – the input to the model.
Returns: the initial state sampled for the model.
Return type: (tensor)

sample(x: Union[torch.Tensor, mbrl.types.TransitionBatch], deterministic: bool = False, **kwargs) → Tuple[torch.Tensor, …]¶

Samples an output of the dynamics model.

The default implementation for all models is equivalent to self.forward(x)[0]. This method will be used by ModelEnv to simulate a step with the model.

Parameters

x (tensor or batch of transitions) – the input to the model.
deterministic (bool) – if True, the model returns a deterministic “sample” (e.g., the mean prediction). Defaults to False.

Returns

any number of tensors that can be sampled from: the model (e.g., observations, rewards, terminations).

Return type

(tuple of tensor)

save(path: Union[str, pathlib.Path])¶: Saves the model to the given path.

update(model_in: Union[torch.Tensor, mbrl.types.TransitionBatch], optimizer: torch.optim.optimizer.Optimizer, target: Optional[torch.Tensor] = None) → float¶

Updates the model using backpropagation with given input and target tensors.

Provides a basic update function, following the steps below:

optimizer.zero_grad()
loss = self.loss(model_in, target)
loss.backward()
optimizer.step()

Parameters

model_in (tensor or batch of transitions) – the inputs to the model.
optimizer (torch.optimizer) – the optimizer to use for the model.
target (tensor or sequence of tensors) – the expected output for the given inputs, if it cannot be computed from model_in.

Returns

the numeric value of the computed loss.

Return type

(float)

class mbrl.models.ModelEnv(env: gym.core.Env, model: mbrl.models.one_dim_tr_model.OneDTransitionRewardModel, termination_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], reward_fn: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, generator: Optional[torch._C.Generator] = None)¶

Bases: object

Wraps a dynamics model into a gym-like environment.

This class can wrap a dynamics model to be used as an environment. The only requirement to use this class is for the model to use this wrapper is to have a method called predict() with signature next_observs, rewards = model.predict(obs,actions, sample=, rng=)

Parameters

env (gym.Env) – the original gym environment for which the model was trained.
model (mbrl.models.Model) – the model to wrap.
termination_fn (callable) – a function that receives actions and observations, and returns a boolean flag indicating whether the episode should end or not.
reward_fn (callable, optional) – a function that receives actions and observations and returns the value of the resulting reward in the environment. Defaults to None, in which case predicted rewards will be used.
generator (torch.Generator, optional) – a torch random number generator (must be in the same device as the given model). If None (default value), a new generator will be created using the default torch seed.

evaluate_action_sequences(action_sequences: torch.Tensor, initial_state: numpy.ndarray, num_particles: int) → torch.Tensor¶

Evaluates a batch of action sequences on the model.

Parameters

action_sequences (torch.Tensor) – a batch of action sequences to evaluate. Shape must be B x H x A, where B, H, and A represent batch size, horizon, and action dimension, respectively.
initial_state (np.ndarray) – the initial state for the trajectories.
num_particles (int) – number of times each action sequence is replicated. The final value of the sequence will be the average over its particles values.

Returns

the accumulated reward for each action sequence, averaged over its particles.

Return type

(torch.Tensor)

reset(initial_obs_batch: numpy.ndarray, return_as_np: bool = True) → Union[torch.Tensor, numpy.ndarray]¶

Resets the model environment.

Parameters

initial_obs_batch (np.ndarray) – a batch of initial observations. One episode for each observation will be run in parallel. Shape must be B x D, where B is batch size, and D is the observation dimension.
return_as_np (bool) – if True, this method and step() will return numpy arrays, otherwise it returns torch tensors in the same device as the model. Defaults to True.

Returns

the initial observation in the type indicated by return_as_np.

Return type

(torch.Tensor or np.ndarray)

step(actions: Union[torch.Tensor, numpy.ndarray], sample: bool = False) → Tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray], numpy.ndarray, Dict]¶

Steps the model environment with the given batch of actions.

Parameters

actions (torch.Tensor or np.ndarray) – the actions for each “episode” to rollout. Shape must be B x A, where B is the batch size (i.e., number of episodes), and A is the action dimension. Note that B must correspond to the batch size used when calling reset(). If a np.ndarray is given, it’s converted to a torch.Tensor and sent to the model device.
sample (bool) – if True model predictions are stochastic. Defaults to False.

Returns

contains the predicted next observation, reward, done flag and metadata. The done flag is computed using the termination_fn passed in the constructor.

Return type

(tuple)

class mbrl.models.ModelTrainer(model: mbrl.models.model.Model, optim_lr: float = 0.0001, weight_decay: float = 1e-05, logger: Optional[mbrl.util.logger.Logger] = None)¶

Bases: object

Trainer for dynamics models.

Parameters

model (mbrl.models.Model) – a model to train.
optim_lr (float) – the learning rate for the optimizer (using Adam).
weight_decay (float) – the weight decay to use.
logger (mbrl.util.Logger, optional) – the logger to use.

evaluate(dataset: mbrl.util.replay_buffer.TransitionIterator) → torch.Tensor¶

Evaluates the model on the validation dataset.

Iterates over the dataset, one batch at a time, and calls mbrl.models.Model.eval_score() to compute the model score over the batch. The method returns the average score over the whole dataset.

Parameters

dataset (bool) – the transition iterator to use.

Returns

The average score of the model over the dataset (and for ensembles, per: ensemble member).

Return type

(tensor)

maybe_get_best_weights(best_val_score: torch.Tensor, val_score: torch.Tensor, threshold: float = 0.01) → Optional[Dict]¶

Return the current model state dict if the validation score improves.

For ensembles, this checks the validation for each ensemble member separately.

Parameters

best_val_score (tensor) – the current best validation losses per model.
val_score (tensor) – the new validation loss per model.
threshold (float) – the threshold for relative improvement.

Returns

if the validation score’s relative improvement over the best validation score is higher than the threshold, returns the state dictionary of the stored model, otherwise returns None.

Return type

(dict, optional)

train(dataset_train: mbrl.util.replay_buffer.TransitionIterator, dataset_val: Optional[mbrl.util.replay_buffer.TransitionIterator] = None, num_epochs: Optional[int] = None, patience: Optional[int] = 1, callback: Optional[Callable] = None) → Tuple[List[float], List[float]]¶

Trains the model for some number of epochs.

This method iterates over the stored train dataset, one batch of transitions at a time, updates the model.

If a validation dataset is provided in the constructor, this method will also evaluate the model over the validation data once per training epoch. The method will keep track of the weights with the best validation score, and after training the weights of the model will be set to the best weights. If no validation dataset is provided, the method will keep the model with the best loss over training data.

Parameters

dataset_train (mbrl.util.TransitionIterator) – the iterator to use for the training data.
dataset_val (mbrl.util.TransitionIterator, optional) – an iterator to use for the validation data.
num_epochs (int, optional) – if provided, the maximum number of epochs to train for. Default is None, which indicates there is no limit.
patience (int, optional) – if provided, the patience to use for training. That is, training will stop after patience number of epochs without improvement.
callback (callable, optional) –
if provided, this function will be called after every training epoch with the following positional arguments:
- the model that’s being trained
- total number of calls made to trainer.train()
- current epoch
- training loss
- validation score (for ensembles, factored per member)
- best validation score so far

Returns

the history of training losses and validation losses.

Return type

(tuple of two list(float))

class mbrl.models.OneDTransitionRewardModel(model: mbrl.models.model.Model, target_is_delta: bool = True, normalize: bool = False, learned_rewards: bool = True, obs_process_fn: Optional[Callable[[numpy.ndarray], numpy.ndarray]] = None, no_delta_list: Optional[List[int]] = None, num_elites: Optional[int] = None)¶

Bases: mbrl.models.model.Model

Wrapper class for 1-D dynamics models.

This model functions as a wrapper for another model to convert transition batches into 1-D transition reward models. It also provides data manipulations that are common when using dynamics models with 1-D observations and actions, so that users don’t have to manipulate the underlying model’s inputs and outputs directly (e.g., predicting delta observations, input normalization).

The wrapper assumes that the wrapped model inputs/outputs will be consistent with

[pred_obs_{t+1}, pred_rewards_{t+1} (optional)] = model([obs_t, action_t]).

Parameters

model (mbrl.model.Model) – the model to wrap.
target_is_delta (bool) – if True, the predicted observations will represent the difference respect to the input observations. That is, ignoring rewards, pred_obs_{t + 1} = obs_t + model([obs_t, act_t]). Defaults to True. Can be deactivated per dimension using no_delta_list.
normalize (bool) – if true, the wrapper will create a normalizer for model inputs, which will be used every time the model is called using the methods in this class. To update the normalizer statistics, the user needs to call update_normalizer() before using the model. Defaults to False.
learned_rewards (bool) – if True, the wrapper considers the last output of the model to correspond to rewards predictions, and will use it to construct training targets for the model and when returning model predictions. Defaults to True.
obs_process_fn (callable, optional) – if provided, observations will be passed through this function before being given to the model (and before the normalizer also). The processed observations should have the same dimensions as the original. Defaults to None.
no_delta_list (list(int), optional) – if provided, represents a list of dimensions over which the model predicts the actual observation and not just a delta.
num_elites (int, optional) – if provided, only the best num_elites models according to validation score are used when calling predict(). Defaults to None which means that all models will always be included in the elite set.

eval_score(batch: mbrl.types.TransitionBatch, target: Optional[torch.Tensor] = None) → torch.Tensor¶

Evaluates the model score over a batch of transitions.

This method constructs input and targets from the information in the batch, then calls self.model.eval_score() on them and returns the value.

Parameters: batch (transition batch) – a batch of transition to train the model.
Returns: as returned by model.eval_score().
Return type: (tensor)

forward(x: torch.Tensor, **kwargs) → Tuple[torch.Tensor, …]¶: Calls forward method of base model with the given input and args.

get_output_and_targets(batch: mbrl.types.TransitionBatch) → Tuple[Tuple[torch.Tensor, …], torch.Tensor]¶

Returns the model output and the target tensors given a batch of transitions.

This method constructs input and targets from the information in the batch, then calls self.model.forward() on them and returns the value. No gradient information will be kept.

Parameters: batch (transition batch) – a batch of transition to train the model.
Returns: the model outputs and the target for this batch.
Return type: (tuple(tensor), tensor)

load(load_dir: Union[str, pathlib.Path])¶: Loads the model from the given path.

loss(batch: mbrl.types.TransitionBatch, target: Optional[torch.Tensor] = None) → torch.Tensor¶

Evaluates the model score over a batch of transitions.

This method constructs input and targets from the information in the batch, then calls self.model.eval_score() on them and returns the value.

Parameters: batch (transition batch) – a batch of transition to train the model.
Returns: as returned by model.eval_score().
Return type: (tensor)

reset(x: mbrl.types.TransitionBatch, rng: Optional[torch._C.Generator] = None) → torch.Tensor¶

Calls reset on the underlying model.

Parameters

x (tensor) – the input to the model.
rng (random number generator) – a rng to use for sampling the model indices.

Returns

the output of the underlying model.

Return type

(tensor)

sample(x: mbrl.types.TransitionBatch, deterministic: bool = False, rng: Optional[torch._C.Generator] = None, **kwargs) → Tuple[torch.Tensor, torch.Tensor]¶

Samples next observations and rewards from the underlying model.

This wrapper assumes that the underlying model’s sample method returns a tuple with just one tensor, which concatenates next_observation and reward.

Parameters

x (transition) – a batch of transitions.
deterministic (bool) – if True, the model returns a deterministic “sample” (e.g., the mean prediction). Defaults to False.
rng (random number generator) – a rng to use for sampling.

Returns

predicted next_observation (o_{t+1}) and rewards (r_{t+1}).

Return type

(tuple of two tensors)

save(save_dir: Union[str, pathlib.Path])¶: Saves the model to the given path.

update(batch: mbrl.types.TransitionBatch, optimizer: torch.optim.optimizer.Optimizer, target: Optional[torch.Tensor] = None) → float¶

Updates the model given a batch of transitions and an optimizer.

Parameters

batch (transition batch) – a batch of transition to train the model.
optimizer (torch optimizer) – the optimizer to use to update the model.

update_normalizer(batch: mbrl.types.TransitionBatch)¶

Updates the normalizer statistics using the batch of transition data.

The normalizer will compute mean and standard deviation the obs and action in the transition. If an observation processing function has been provided, it will be called on obs before updating the normalizer.

Parameters: batch (mbrl.types.TransitionBatch) – The batch of transition data. Only obs and action will be used, since these are the inputs to the model.

mbrl.models.truncated_normal_init(m: torch.nn.modules.module.Module)¶: Initializes the weights of the given module using a truncated normal distribution.