dreamerx.models module#

Utility models and layers for construction of reinforcement learning agents.

class dreamerx.models.Actor(input_dim: int, actions: list[Action], num_layers: int = 3, hidden_dim: int = 512)#

Bases: Module

Actor network.

Pulls action heads from a latent representation and samples.

__init__(input_dim: int, actions: list[Action], num_layers: int = 3, hidden_dim: int = 512) → None#

Initialize the actor network.

Parameters:

  • input_dim (int) – The dimension of the input vector.

  • actions (list[dreamerx.actions.Action]) – A list of action definitions, can be continuous or discrete.

  • num_layers (int) – The number of blocks in the MLP. (Default: 3)

  • hidden_dim (int) – The dimension of the hidden layers. (Default: 512)

forward(x: Tensor) → tuple[Tensor, Tensor]#

Perform a forward pass through the actor network.

Parameters:

x (torch.Tensor) – The input tensor of shape (batch_dim, input_dim).

Returns:

A tuple containing the sampled actions and their distributions.

Return type:

tuple[torch.Tensor, torch.Tensor]

class dreamerx.models.AttentionDecoder(input_dim: int, output_dims: list[int], num_queries: list[int] | int = 1, hidden_dim: int = 512, num_layers: int = 2, num_heads: int = 8)#

Bases: Module

Attention decoder for reconstructing sequential observations from latent representations.

__init__(input_dim: int, output_dims: list[int], num_queries: list[int] | int = 1, hidden_dim: int = 512, num_layers: int = 2, num_heads: int = 8) → None#

Initialize the attention decoder.

Parameters:

  • input_dim (int) – The dimension of the input latent representation.

  • output_dims (list[int]) – A list of output dimensions for each sequential observation to be reconstructed with attention.

  • num_queries (list[int] | int) – A list of the number of queries to use for each output sequential observation, which determines the maximum number of segments that can be reconstructed for each output. If a single integer is provided, the same number of queries will be used for each output. (Default: 1)

  • hidden_dim (int) – The dimension of the hidden layers and output per query. (Default: 512)

  • num_layers (int) – The number of layers in the attention decoder. (Default: 2)

  • num_heads (int | None) – The number of attention heads to use. If None, defaults to the number of layers. (Default: None)

forward(x: Tensor) → tuple[list[Tensor], list[Tensor]]#

Perform a forward pass through the attention decoder.

Parameters:

x (torch.Tensor) – The input tensor of shape (batch_dim, input_dim).

Returns:

A tuple containing a list of output tensors for each sequential observation in the original dimension and a list of existence logits for each sequential observation.

Return type:

tuple[list[torch.Tensor], list[torch.Tensor]]

class dreamerx.models.AttentionEncoder(input_dims: list[int], hidden_dim: int = 512, num_layers: int = 2, num_heads: int = 8, num_queries: int = 1)#

Bases: Module

Attention encoder for processing sequential observations with attention mechanisms.

__init__(input_dims: list[int], hidden_dim: int = 512, num_layers: int = 2, num_heads: int = 8, num_queries: int = 1) → None#

Initialize the attention encoder.

Parameters:

  • input_dims (list[int]) – A list of input dimensions for each sequential observation to be processed with attention.

  • hidden_dim (int) – The dimension of the hidden layers and output per query. (Default: 512)

  • num_layers (int) – The number of layers in the attention encoder. (Default: 2)

  • num_heads (int | None) – The number of attention heads to use. If None, defaults to the number of layers. (Default: None)

  • num_queries (int) – The number of learned queries to use in the cross attention layer. (Default: 1)

forward(x: list[Tensor]) → Tensor#

Perform a forward pass through the attention encoder.

Parameters:

x (list[torch.Tensor]) – A list of input tensors for each sequential observation, where each tensor has shape (batch_dims, seq_len, input_dim).

Returns:

The output tensor of shape (batch_dims, num_queries * hidden_dim).

Return type:

torch.Tensor

property output_dim: int#

Output dimension of the attention encoder.

Type:

int

class dreamerx.models.BlockLinear(input_dim: int, output_dim: int, num_blocks: int = 8, bias: bool = True)#

Bases: Module

Linear layer that splits the input into blocks and processes each block separately.

__init__(input_dim: int, output_dim: int, num_blocks: int = 8, bias: bool = True) → None#

Initialize the BlockLinear layer.

Parameters:

  • input_dim (int) – The size of the input tensor.

  • output_dim (int) – The size of the output tensor.

  • num_blocks (int) – The number of blocks for splitting the input and output. (Default: 8)

  • bias (bool) – Whether to include a bias term. (Default: True)

forward(x: Tensor) → Tensor#

Perform a forward pass through the BlockLinear layer.

Parameters:

x (torch.Tensor) – The input tensor of shape (batch_dim, input_dim).

Returns:

The output tensor of shape (batch_dim, output_dim).

Return type:

torch.Tensor

class dreamerx.models.BlockRecurrentModel(stoch_dim: int, act_dim: int, hidden_dim: int, deter_dim: int, num_blocks: int = 8)#

Bases: Module

Block recurrent model for processing sequences of latent representations.

Uses a block GRU to process sequences of latent representations, with layer normalization and SiLU activations on the output.

__init__(stoch_dim: int, act_dim: int, hidden_dim: int, deter_dim: int, num_blocks: int = 8) → None#

Initialize the recurrent model.

Parameters:

  • stoch_dim (int) – The size of the stochastic state.

  • act_dim (int) – The size of the action tensor.

  • hidden_dim (int) – The size of the hidden dimension.

  • deter_dim (int) – The size of the deterministic state.

  • num_blocks (int) – The number of blocks to split the deterministic state. Must divide deter_dim. (Default: 8)

forward(s: Tensor, a: Tensor, d: Tensor | None = None) → Tensor#

Perform a forward pass through the recurrent model.

Parameters:

  • s (torch.Tensor) – The stochastic tensor of shape (batch_dim, seq_len, latent_dim).

  • a (torch.Tensor) – The action tensor of shape (batch_dim, seq_len, action_dim).

  • d (torch.Tensor | None) – The initial deterministic state of shape (1, batch_dim, deter_dim), or None to use zeros. (Default: None)

Returns:

The final deterministic state of shape (batch_dim, deter_dim).

Return type:

torch.Tensor

property deter_dim: int#

Deterministic size of the recurrent model.

Type:

int

class dreamerx.models.CNNDecoder(output_channels: int, input_dim: int, image_dim: tuple[int, int] = (64, 64), num_blocks: int = 4)#

Bases: Module

CNN decoder for reconstructing image observations from latent representations.

Uses transposed convolutions to upsample to original resolution, with interleaved layer normalization and SiLU activations.

__init__(output_channels: int, input_dim: int, image_dim: tuple[int, int] = (64, 64), num_blocks: int = 4) → None#

Initialize the CNN decoder.

Parameters:

  • output_channels (int) – The number of output channels in the reconstructed image.

  • input_dim (int) – The size of the latent representation.

  • image_dim (tuple[int, int]) – The size of the output image. (Default: (64, 64))

  • num_blocks (int) – The number of blocks in the CNN. (Default: 4)

forward(x: Tensor) → Tensor#

Perform a forward pass through the CNN decoder.

Parameters:

x (torch.Tensor) – The input tensor of shape (batch_dim, latent_dim).

Returns:

The output tensor of shape (batch_dim, channels, height, width).

Return type:

torch.Tensor

class dreamerx.models.CNNEncoder(input_channels: int, image_dim: tuple[int, int] = (64, 64), num_blocks: int = 4)#

Bases: Module

CNN encoder for processing image observations.

Processes a (generally) 64x64 image to 4x4 using [2, 4, 8, 16] channels and interleaved layer normalization and SiLU activations.

__init__(input_channels: int, image_dim: tuple[int, int] = (64, 64), num_blocks: int = 4) → None#

Initialize the CNN encoder.

Parameters:

  • input_channels (int) – The number of input channels in the image.

  • image_dim (tuple[int, int]) – The size of the input image. (Default: (64, 64))

  • num_blocks (int) – The number of blocks in the CNN. (Default: 4)

forward(x: Tensor) → Tensor#

Perform a forward pass through the CNN encoder.

Parameters:

x (torch.Tensor) – The input tensor of shape (batch_dim, channels, height, width).

Returns:

The output tensor of shape (batch_dim, output_dim, height, width).

Return type:

torch.Tensor

property output_dim: int#

Output dimension of the CNN encoder.

Type:

int

class dreamerx.models.ChannelNorm(*args: list[Any], **kwargs: dict[str, Any])#

Bases: LayerNorm

Layer normalization across the channel dimension.

__init__(*args: list[Any], **kwargs: dict[str, Any]) → None#

Initialize the ChannelNorm layer.

Parameters:

  • args (list[Any]) – Positional arguments for the LayerNorm.

  • kwargs (dict[str, Any]) – Keyword arguments for the LayerNorm.

forward(x: Tensor) → Tensor#

Apply layer normalization across the channel dimension.

Swaps the channel dimension to the end, applies layer normalization, then swaps it back.

Parameters:

x (torch.Tensor) – The input tensor of shape (batch_dim, channels, height, width).

Returns:

The normalized tensor of the same shape.

Return type:

torch.Tensor

class dreamerx.models.CompoundDecoder(*decoder_specs: list[dict[str, Any]], input_dim: int = 512, num_blocks: int = 4, num_layers: int = 3, num_att_layers: int = 2, num_heads: int = 8, hidden_dim: int = 512)#

Bases: Module

Compound decoder for reconstructing MLP, CNN, and attention observations from latent representations.

__init__(*decoder_specs: list[dict[str, Any]], input_dim: int = 512, num_blocks: int = 4, num_layers: int = 3, num_att_layers: int = 2, num_heads: int = 8, hidden_dim: int = 512) → None#

Initialize the compound decoder.

Parameters:

  • decoder_specs (list[dict[str, Any]]) – A list of specifications for each decoder, where each specification is a dictionary containing the type and output specifications for the decoder. These options include type and outputs, where type is one of MLP, CNN, or ATTENTION, and outputs is a list of output specifications for the decoder. Specifically, for MLP, outputs should be a list of dictionaries containing the key range with a list of (start, end) indices for the output range. For CNN, outputs should be a list of dictionaries containing the key image_dim with a tuple of (channels, height, width) for the output image and an optional offset containing the index of the first value for the flattened image. For ATTENTION, outputs should be a list of dictionaries containing the key range with a list of (start, end) indices for the output range, and an optional key segment_len with the length of each segment to use as an output from the attention decoder.

  • input_dim (int) – The input dimension of each decoder. (Default: 512)

  • num_blocks (int) – The number of blocks in each CNN decoder. (Default: 4)

  • num_layers (int) – The number of layers in each block of the MLP decoders. (Default: 3)

  • num_att_layers (int) – The number of layers in the attention decoders. (Default: 2)

  • num_heads (int) – The number of attention heads in each attention decoder. (Default: 8)

  • hidden_dim (int) – The hidden dimension of each decoder. (Default: 512)

forward(x: Tensor) → list[Tensor | tuple[list[Tensor], list[Tensor]]]#

Perform a forward pass through the compound decoder, combining the outputs of each individual decoder.

Parameters:

  • x (torch.Tensor) – The input tensor of shape (batch_dim, input_dim).

  • reconstruct (bool) – Whether to reconstruct the output in its original format.

Returns:

A list of all reconstructed outputs, where MLP and CNN decoders return tensors, while attention decoders return a tuple containing a list of output tensors and a list of existence logits for each sequential observation.

Return type:

list[torch.Tensor | tuple[list[torch.Tensor], list[torch.Tensor]]]

class dreamerx.models.CompoundEncoder(*encoder_specs: list[dict[str, Any]], output_dim: int = 512, num_blocks: int = 4, num_layers: int = 3, num_att_layers: int = 2, num_heads: int = 8, hidden_dim: int = 512)#

Bases: Module

Compound encoder for processing MLP, CNN, and attention observations.

__init__(*encoder_specs: list[dict[str, Any]], output_dim: int = 512, num_blocks: int = 4, num_layers: int = 3, num_att_layers: int = 2, num_heads: int = 8, hidden_dim: int = 512) → None#

Initialize the compound encoder.

Parameters:

  • encoder_specs (list[dict[str, Any]]) – A list of specifications for each encoder, where each specification is a dictionary containing the type and segment specifications for the encoder. These options include type and segments, where type is one of MLP, CNN, or ATTENTION, and segments is a list of segment specifications for the encoder. Specifically, for MLP, segments should be a list of dictionaries containing the key range with a list of (start, end) indices for the input range. For CNN, segments should be a list of dictionaries containing the key image_dim with a tuple of (channels, height, width) for the input image and an optional offset containing the index of the first value for the flattened image. For ATTENTION, segments should be a list of dictionaries containing the key range with a list of (start, end) indices for the input range, and an optional key segment_len with the length of each segment to use as an input to the attention encoder.

  • output_dim (int) – The output dimension of each encoder. (Default: 512)

  • num_blocks (int) – The number of blocks in each CNN encoder. (Default: 4)

  • num_layers (int) – The number of layers in each block of the MLP encoders. (Default: 3)

  • num_att_layers (int) – The number of layers in the attention encoders. (Default: 2)

  • num_heads (int) – The number of attention heads in each attention encoder. (Default: 8)

  • hidden_dim (int) – The hidden dimension of each encoder. (Default: 512)

forward(representation: list[Tensor | list[Tensor]]) → Tensor#

Perform a forward pass through the compound encoder, combining the outputs of each individual encoder.

Parameters:

representation (list[torch.Tensor | list[torch.Tensor]]) – A list of input representations for each encoder, normally produced using the extract_representation function.

Returns:

The output tensor of shape (batch_dim, output_dim).

Return type:

torch.Tensor

property output_dim: int#

Output dimension of the compound encoder.

Type:

int

class dreamerx.models.LayerNormGRU(input_dim: int, hidden_dim: int)#

Bases: Module

GRU layer with internal layer normalization.

__init__(input_dim: int, hidden_dim: int) → None#

Initialize the LayerNormGRU.

Parameters:

  • input_dim (int) – The size of the input tensor.

  • hidden_dim (int) – The size of the hidden state.

forward(x: Tensor, h: Tensor | None = None) → Tensor#

Perform a forward pass through the LayerNormGRU.

Applies layer normalization to the hidden state after the GRU computation.

Parameters:

  • x (torch.Tensor) – The input tensor of shape (batch_dim, input_dim).

  • h (torch.Tensor | None) – The initial hidden state of shape (batch_dim, hidden_dim), or None to use zeros. (Default: None)

Returns:

The final hidden state of shape (batch_dim, hidden_dim).

Return type:

torch.Tensor

class dreamerx.models.MLP(input_dim: int, output_dim: int | None = None, hidden_dims: list[int] = [])#

Bases: Module

MLP for processing vector inputs.

__init__(input_dim: int, output_dim: int | None = None, hidden_dims: list[int] = []) → None#

Initialize the MLP.

Parameters:

  • input_dim (int) – The dimension of the input vector.

  • output_dim (int | None) – The dimension of the output vector. If provided, a final hidden layer will be added with no normalization or activation

  • hidden_dims (list[int]) – The dimensions of the hidden layers.

forward(x: Tensor) → Tensor#

Perform a forward pass through the MLP.

Parameters:

x (torch.Tensor) – The input tensor of shape (batch_dim, input_dim).

Returns:

The output tensor of shape (batch_dim, hidden_dims[-1]).

Return type:

torch.Tensor

class dreamerx.models.MLPDecoder(input_dim: int, output_dim: int, num_layers: int = 3, hidden_dim: int = 512)#

Bases: Module

MLP decoder for reconstructing vector observations from latent representations.

__init__(input_dim: int, output_dim: int, num_layers: int = 3, hidden_dim: int = 512) → None#

Initialize the MLP decoder.

Parameters:

  • input_dim (int) – The dimension of the input tensor.

  • output_dim (int) – The dimension of the output vector observations.

  • num_blocks (int) – The number of blocks in the MLP. (Default: 3)

  • hidden_dim (int) – The dimension of the hidden layers. (Default: 512)

forward(x: Tensor) → Tensor#

Perform a forward pass through the MLP decoder.

Parameters:

x (torch.Tensor) – The input tensor of shape (batch_dim, hidden_dim).

Returns:

The output tensor of shape (batch_dim, output_dim).

Return type:

torch.Tensor

class dreamerx.models.MLPEncoder(input_dim: int, output_dim: int, num_layers: int = 3, hidden_dim: int = 512, use_symlog: bool = True)#

Bases: Module

MLP encoder for processing vector observations.

__init__(input_dim: int, output_dim: int, num_layers: int = 3, hidden_dim: int = 512, use_symlog: bool = True) → None#

Initialize the MLP encoder.

Parameters:

  • input_dim (int) – The dimension of the input vector observation.

  • output_dim (int) – The dimension of the output vector.

  • num_blocks (int) – The number of blocks in the MLP. (Default: 3)

  • hidden_dim (int) – The dimension of the hidden layers. (Default: 512)

  • use_symlog (bool) – Whether to apply symlog transformation to the input. (Default: True)

forward(x: Tensor) → Tensor#

Perform a forward pass through the MLP encoder.

Parameters:

x (torch.Tensor) – The input tensor of shape (batch_dim, input_dim).

Returns:

The output tensor of shape (batch_dim, hidden_dim).

Return type:

torch.Tensor

property output_dim: int#

Output dimension of the MLP decoder.

Type:

int

class dreamerx.models.PositionalEncoding(d_model: int, dropout: float = 0.1, max_len: int = 5000)#

Bases: Module

Module for adding positional encoding to input tensors.

From https://pytorch-tutorials-preview.netlify.app/beginner/transformer_tutorial.html.

__init__(d_model: int, dropout: float = 0.1, max_len: int = 5000) → None#

Initialize the positional encoding.

Parameters:

  • d_model (int) – The dimension of the model and positional encoding.

  • dropout (float) – The dropout rate. (Default: 0.1)

  • max_len (int) – The maximum length of the input sequences. (Default: 5000)

forward(x: Tensor) → Tensor#

Add positional encoding to the input tensor.

Parameters:

x (torch.Tensor) – The input tensor

Returns:

The output tensor with positional encoding added.

Return type:

torch.Tensor

class dreamerx.models.RSSM(recurrent_model: Module, representation_model: Module, transition_model: Module, bins: int = 32, learnable_initial_state: bool = False)#

Bases: Module

Recurrent state-space model for modeling environment dynamics.

The RSSM consists of a recurrent model for tracking the hidden state, a representation model for inferring the stochastic state from the hidden state and observation, and a transition model for predicting the stochastic state from the hidden state.

__init__(recurrent_model: Module, representation_model: Module, transition_model: Module, bins: int = 32, learnable_initial_state: bool = False) → None#

Initialize the RSSM.

Parameters:

  • recurrent_model (nn.Module) – The recurrent model for keeping track of the environment dynamics.

  • representation_model (nn.Module) – The model for inferring the stochastic state from the hidden state and observation.

  • transition_model (nn.Module) – The model for predicting the stochastic state from the hidden state.

  • bins (int) – The number of bins for the stochastic state. (Default: 32)

forward(action: Tensor | None = None, posterior: Tensor | None = None, hidden_state: Tensor | None = None, embedded_obs: Tensor | None = None, initialize: Tensor | None = None, batch_dim: int | None = None, compute_prior: bool = True) → dict[str, Tensor]#

Perform one step of the RSSM.

Will compute the hidden state using action, posterior, and previous hidden state. If not available, will instead use the initial hidden state. If embedded_obs is not provided, imagines the next hidden state.

Parameters:

  • action (torch.Tensor | None) – The action taken, of shape (batch_dim, action_dim).

  • posterior (torch.Tensor | None) – The posterior stochastic state from the previous step, of shape (batch_dim, latent_dim).

  • hidden_state (torch.Tensor | None) – The initial hidden state of the recurrent model, of shape (batch_dim, hidden_dim).

  • embedded_obs (torch.Tensor | None) – The embedded observation, of shape (batch_dim, obs_dim), or None if imagining.

  • initialize (torch.Tensor | None) – Boolean tensor of hidden states to initialize of shape (batch_dim). (Default: None)

  • batch_dim (int | None) – The batch dimension, inferred if not provided. (Default: None)

  • compute_prior (bool) – Whether to compute the prior distribution. (Default: True)

Returns:

A dictionary containing the prior and posterior logits and samples, and the updated hidden state.

Return type:

dict[str, torch.Tensor]

infer_stochastic(hidden_state: Tensor, embedded_obs: Tensor | None = None) → tuple[Tensor, Tensor]#

Infer the stochastic state.

Infer the stochastic state from the hidden state (using the transition model) or from the hidden state and the embedded observation (using the representation model).

Parameters:

  • hidden_state (torch.Tensor) – The hidden state from the recurrent model, of shape (batch_dim, hidden_dim).

  • embedded_obs (torch.Tensor | None) – The embedded observation, of shape (batch_dim, obs_dim). (Default: None)

Returns:

The logits and sampled stochastic state.

Return type:

tuple[torch.Tensor, torch.Tensor]

property initial_hidden_state: Tensor#

Trainable initial hidden state of the recurrent model.

Type:

torch.Tensor

class dreamerx.models.SingleRecurrentModel(input_dim: int, hidden_dim: int, deter_dim: int)#

Bases: Module

Recurrent model for processing sequences of latent representations.

Uses a GRU to process sequences of latent representations, with layer normalization and SiLU activations on the output.

__init__(input_dim: int, hidden_dim: int, deter_dim: int) → None#

Initialize the recurrent model.

Parameters:

  • input_dim (int) – The size of the input tensor.

  • hidden_dim (int) – The size of the hidden dimensions.

  • deter_dim (int) – The size of the deterministic state output by the recurrent model.

forward(x: Tensor, h: Tensor | None = None) → Tensor#

Perform a forward pass through the recurrent model.

Parameters:

  • x (torch.Tensor) – The input tensor of shape (batch_dim, seq_len, latent_dim).

  • h (torch.Tensor | None) – The initial deterministic state of shape (1, batch_dim, deter_dim), or None to use zeros. (Default: None)

Returns:

The final deterministic state of shape (batch_dim, deter_dim).

Return type:

torch.Tensor

property deter_dim: int#

Deterministic size of the recurrent model.

Type:

int

dreamerx.models.extract_representation(x: Tensor, specs: list[dict[str, Any]]) → list[Tensor | list[Tensor]]#

Extract representations from the input tensor based on the provided specifications.

Parameters:

  • x (torch.Tensor) – The input tensor of shape (batch_dim, input_dim).

  • specs (list[dict[str, Any]]) – A list of specifications for each representation to extract. More details can be found in the documentation for the CompoundEncoder class.

Returns:

A list of extracted representations, where each representation is either a tensor or a list of tensors depending on the specification.

Return type:

list[torch.Tensor | list[torch.Tensor]]