dreamerx.losses module#

Utility loss functions for reinforcement learning agents.

dreamerx.losses.attention_reconstruction_loss(entities_pred: Tensor, entities_target: Tensor, existence_logits: Tensor, existence_target: Tensor = None, reconstruction_loss_type: str = 'positional') → tuple[Tensor, Tensor]#

Compute the attention reconstruction loss between the predicted and target distributions.

Parameters:

  • entities_pred (torch.Tensor) – The predicted entity features, of shape (batch_dim, entities, features).

  • entities_target (torch.Tensor) – The target entity features, of shape (batch_dim, entities, features).

  • existence_logits (torch.Tensor) – The predicted entity existence logits, of shape (batch_dim, entities).

  • existence_target (torch.Tensor) – The target entity existence probabilities (0 or 1), of shape (batch_dim, entities). If not provided, it will be computed as 1 for entities with any non-zero features.

Returns:

A tuple containing the attention reconstruction loss and the existence loss.

Return type:

tuple[torch.Tensor, torch.Tensor]

dreamerx.losses.hungarian_loss(pred: Tensor, target: Tensor, dims: int = 2, clip: float = 1e-08) → Tensor#

Compute the Hungarian loss between the predicted and target distributions on the final dimension.

This loss function is not vectorized and is computationally expensive.

Parameters:

  • pred (torch.Tensor) – The predicted distribution, of shape (batch_dim, entities, features).

  • target (torch.Tensor) – The target distribution, of shape (batch_dim, entities, features).

  • dims (int) – The number of final dimensions to compute the loss over. (Default: 1)

  • clip (float) – The minimal value of the MSE loss, to prevent NaNs. (Default: 1e-8)

Returns:

The Hungarian loss.

Return type:

torch.Tensor

dreamerx.losses.mse_loss(pred: Tensor, target: Tensor, dims: int = 1) → Tensor#

Compute the mean squared error between the predicted and target distributions on the final dimension.

Parameters:

  • pred (torch.Tensor) – The predicted distribution.

  • target (torch.Tensor) – The target distribution.

  • dims (int) – The number of final dimensions to compute the loss over. (Default: 1)

Returns:

The mean squared error loss.

Return type:

torch.Tensor