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¶

De-normalizes 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 de-normalize.
Returns: The de-normalized 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 log-likelihood 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 to True.

Returns

the negative log-likelihood.

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, 707-710 (1995)

Normalised to unit variance

Parameters

exponent (float) – the power-spectrum 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) – low-frequency 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, where E, B, and Od 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, where E, B, and Od 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, where E, B, and Od 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, where E, B, and Od 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, where E, B, and Od 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 to False.`

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 < 1e-10, 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 in-place.

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)