Math utilities module¶

class
mbrl.util.math.
Normalizer
(in_size: int, device: torch.device, dtype=torch.float32)¶ Bases:
object
Class that keeps a running mean and variance and normalizes data accordingly.
The statistics kept are stored in torch tensors.
 Parameters
in_size (int) – the size of the data that will be normalized.
device (torch.device) – the device in which the data will reside.
dtype (torch.dtype) – the data type to use for the normalizer.

denormalize
(val: Union[float, torch.Tensor, numpy.ndarray]) → torch.Tensor¶ Denormalizes the value according to the stored statistics.
Equivalent to sigma * val + mu, where mu and sigma are the stored mean and standard deviation, respectively.
 Parameters
val (float, np.ndarray or torch.Tensor) – The value to denormalize.
 Returns
The denormalized value.
 Return type
(torch.Tensor)

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

normalize
(val: Union[float, torch.Tensor, numpy.ndarray]) → torch.Tensor¶ Normalizes the value according to the stored statistics.
Equivalent to (val  mu) / sigma, where mu and sigma are the stored mean and standard deviation, respectively.
 Parameters
val (float, np.ndarray or torch.Tensor) – The value to normalize.
 Returns
The normalized value.
 Return type
(torch.Tensor)

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

update_stats
(data: Union[torch.Tensor, numpy.ndarray])¶ Updates the stored statistics using the given data.
Equivalent to self.stats.mean = data.mean(0) and self.stats.std = data.std(0).
 Parameters
data (np.ndarray or torch.Tensor) – The data used to compute the statistics.

mbrl.util.math.
gaussian_nll
(pred_mean: torch.Tensor, pred_logvar: torch.Tensor, target: torch.Tensor, reduce: bool = True) → torch.Tensor¶ Negative loglikelihood for Gaussian distribution
 Parameters
pred_mean (tensor) – the predicted mean.
pred_logvar (tensor) – the predicted log variance.
target (tensor) – the target value.
reduce (bool) – if
False
the loss is returned w/o reducing. Defaults toTrue
.
 Returns
the negative loglikelihood.
 Return type
(tensor)

mbrl.util.math.
powerlaw_psd_gaussian
(exponent: float, size: Union[int, Iterable[int]], device: torch.device, fmin: float = 0)¶ Gaussian (1/f)**beta noise.
Based on the algorithm in: Timmer, J. and Koenig, M.:On generating power law noise. Astron. Astrophys. 300, 707710 (1995)
Normalised to unit variance
 Parameters
exponent (float) – the powerspectrum of the generated noise is proportional to S(f) = (1 / f)**exponent.
size (int or iterable) – the output shape and the desired power spectrum is in the last coordinate.
device (torch.device) – device where computations will be performed.
fmin (float) – lowfrequency cutoff. Default: 0 corresponds to original paper.
 Returns
(torch.Tensor): The samples.

mbrl.util.math.
propagate
(predictions: Tuple[torch.Tensor, …], propagation_method: str = 'expectation', propagation_indices: Optional[torch.Tensor] = None) → Tuple[torch.Tensor, …]¶ Propagates ensemble outputs according to desired method.
Implements propagations options as described in Chua et al., NeurIPS 2018 paper (PETS) https://arxiv.org/pdf/1805.12114.pdf
Valid propagation options are:
“random_model”: equivalent to
propagate_random_model()
. This corresponds to TS1 propagation in the PETS paper.“fixed_model”: equivalent to
propagate_fixed_model()
. This can be used to implement TSinf propagation, described in the PETS paper.“expectation”: equivalent to
propagate_expectation()
.
 Parameters
predictions (tuple of tensors) – the predictions to propagate. Each tensor’s shape must be
E x B x Od
, whereE
,B
, andOd
represent the number of models, batch size, and output dimension, respectively.propagation_method (str) – the propagation method to use.
propagation_indices (tensor, optional) – the model indices to choose (will use the same for all predictions). Only needed if
propagation == "fixed_model"
.
 Returns
the propagated predictions.
 Return type
(tuple of tensors)

mbrl.util.math.
propagate_expectation
(predictions: Tuple[torch.Tensor, …]) → Tuple[torch.Tensor, …]¶ Propagates ensemble outputs by taking expectation over model predictions.
 Parameters
predictions (tuple of tensors) – the predictions to propagate. Each tensor’s shape must be
E x B x Od
, whereE
,B
, andOd
represent the number of models, batch size, and output dimension, respectively. Returns
 the chosen predictions, so that
output[k][i, :] = predictions[k].mean(dim=0)
 Return type
(tuple of tensors)

mbrl.util.math.
propagate_fixed_model
(predictions: Tuple[torch.Tensor, …], propagation_indices: torch.Tensor) → Tuple[torch.Tensor, …]¶ Propagates ensemble outputs by taking expectation over model predictions.
 Parameters
predictions (tuple of tensors) – the predictions to propagate. Each tensor’s shape must be
E x B x Od
, whereE
,B
, andOd
represent the number of models, batch size, and output dimension, respectively.propagation_indices (tensor) – the model indices to choose (will use the same for all predictions).
 Returns
 the chosen predictions, so that
output[k][i, :] = predictions[k].mean(dim=0)
 Return type
(tuple of tensors)

mbrl.util.math.
propagate_from_indices
(predicted_tensor: torch.Tensor, indices: torch.Tensor) → torch.Tensor¶ Propagates ensemble outputs using the given indices.
 Parameters
predicted_tensor (tensor) – the prediction to propagate. Shape must be
E x B x Od
, whereE
,B
, andOd
represent the number of models, batch size, and output dimension, respectively.indices (tensor) – the model indices to choose.
 Returns
 the chosen prediction, so that
output[i, :] = predicted_tensor[indices[i], i, :].
 Return type
(tensor)

mbrl.util.math.
propagate_random_model
(predictions: Tuple[torch.Tensor, …]) → Tuple[torch.Tensor, …]¶ Propagates ensemble outputs by choosing a random model.
 Parameters
predictions (tuple of tensors) – the predictions to propagate. Each tensor’s shape must be
E x B x Od
, whereE
,B
, andOd
represent the number of models, batch size, and output dimension, respectively. Returns
 the chosen predictions, so that
output[k][i, :] = predictions[k][random_choice, i, :].
 Return type
(tuple of tensors)

mbrl.util.math.
quantize_obs
(obs: numpy.ndarray, bit_depth: int, original_bit_depth: int = 8, add_noise: bool = False)¶ Quantizes an array of pixel observations to the desired bit depth.
 Parameters
obs (np.ndarray) – the array to quantize.
bit_depth (int) – the desired bit depth.
original_bit_depth (int, optional) – the original bit depth, defaults to 8.
add_noise (bool, optional) – if
True
, uniform noise in the range (0, 2 ** (8  bit_depth)) will be added. Defaults toFalse
.`
 Returns
the quantized version of the array.
 Return type
(np.ndarray)

mbrl.util.math.
truncated_linear
(min_x: float, max_x: float, min_y: float, max_y: float, x: float) → float¶ Truncated linear function.
 Implements the following function:
f1(x) = min_y + (x  min_x) / (max_x  min_x) * (max_y  min_y) f(x) = min(max_y, max(min_y, f1(x)))
If max_x  min_x < 1e10, then it behaves as the constant f(x) = max_y

mbrl.util.math.
truncated_normal_
(tensor: torch.Tensor, mean: float = 0, std: float = 1)¶ Samples from a truncated normal distribution inplace.
 Parameters
tensor (tensor) – the tensor in which sampled values will be stored.
mean (float) – the desired mean (default = 0).
std (float) – the desired standard deviation (default = 1).
 Returns
 the tensor with the stored values. Note that this modifies the input tensor
in place, so this is just a pointer to the same object.
 Return type
(tensor)