representations

cnn

class xuance.mindspore.representations.cnn.AC_CNN_Atari(*args: Any, **kwargs: Any)[source]

Bases: Cell

construct(observations: mindspore.Tensor)[source]
class xuance.mindspore.representations.cnn.Basic_CNN(*args: Any, **kwargs: Any)[source]

Bases: Cell

construct(observations: mindspore.Tensor)[source]

mlp

class xuance.mindspore.representations.mlp.Basic_Identical(*args: Any, **kwargs: Any)[source]

Bases: Cell

construct(observations: mindspore.Tensor)[source]
class xuance.mindspore.representations.mlp.Basic_MLP(*args: Any, **kwargs: Any)[source]

Bases: Cell

construct(observations: mindspore.Tensor)[source]

rnn

class xuance.mindspore.representations.rnn.Basic_RNN(*args: Any, **kwargs: Any)[source]

Bases: Cell

construct(x: mindspore.Tensor, h: mindspore.Tensor, c: mindspore.Tensor = None)[source]
get_hidden_item(i, *rnn_hidden)[source]
init_hidden(batch)[source]
init_hidden_item(indexes: list, *rnn_hidden)[source]

vit

class xuance.mindspore.representations.vit.Attention(*args: Any, **kwargs: Any)[source]

Bases: Module

forward(x)[source]
class xuance.mindspore.representations.vit.Basic_ViT(*args: Any, **kwargs: Any)[source]

Bases: Module

forward(observations: numpy.ndarray)[source]
class xuance.mindspore.representations.vit.FeedForward(*args: Any, **kwargs: Any)[source]

Bases: Module

forward(x)[source]
class xuance.mindspore.representations.vit.Transformer(*args: Any, **kwargs: Any)[source]

Bases: Module

forward(x)[source]
class xuance.mindspore.representations.vit.ViT(*args: Any, **kwargs: Any)[source]

Bases: Module

forward(video)[source]
xuance.mindspore.representations.vit.pair(t)[source]

world_model

class xuance.mindspore.representations.world_model.Actor(*args: Any, **kwargs: Any)[source]

Bases: Module

The wrapper class of the Dreamer_v2 Actor model.

Parameters:

  • latent_state_size (int) – the dimension of the latent state (stochastic size + recurrent_state_size).

  • actions_dim (Sequence[int]) – the dimension in output of the actor. The number of actions if continuous, the dimension of the action if discrete.

  • is_continuous (bool) – whether the actions are continuous.

  • distribution_config (Dict[str, Any]) – The configs of the distributions.

  • init_std (float) – the amount to sum to the standard deviation. Default to 0.0.

  • min_std (float) – the minimum standard deviation for the actions. Default to 1.0.

  • max_std (float) – the maximum standard deviation for the actions. Default to 1.0.

  • dense_units (int) – the dimension of the hidden dense layers. Default to 1024.

  • activation (int) – the activation function to apply after the dense layers. Default to nn.SiLU.

  • mlp_layers (int) – the number of dense layers. Default to 5.

  • layer_norm_cls (Callable[..., nn.Module]) – the layer norm to apply after the input projection. Defaults to LayerNorm.

  • layer_norm_kw (Dict[str, Any]) – the kwargs of the layer norm. Default to {“eps”: 1e-3}.

  • unimix – (float, optional): the percentage of uniform distribution to inject into the categorical distribution over actions, i.e. given some logits l and probabilities p = softmax(l), then p = (1 - self.unimix) * p + self.unimix * unif, where unif = `1 / self.discrete. Defaults to 0.01.

  • action_clip (float) – the action clip parameter. Default to 1.0.

forward(state: torch.Tensor, greedy: bool = False, mask: Dict[str, torch.Tensor] | None = None) → Tuple[Sequence[torch.Tensor], Sequence[torch.distributions.Distribution]][source]

Call the forward method of the actor model and reorganizes the result with shape (batch_size, *, num_actions), where * means any number of dimensions including None.

Parameters:

  • state (Tensor) – the current state of shape (batch_size, *, stochastic_size + recurrent_state_size).

  • greedy (bool) – whether or not to sample the actions. Default to False.

  • mask (Dict[str, Tensor], optional) – the mask to use on the actions. Default to None.

Returns:

The tensor of the actions taken by the agent with shape (batch_size, *, num_actions). The distribution of the actions

class xuance.mindspore.representations.world_model.CNNDecoder(*args: Any, **kwargs: Any)[source]

Bases: Module

The exact inverse of the CNNEncoder class. It assumes an initial resolution of 4x4, and in 4 stages reconstructs the observation image to 64x64. If multiple images are to be reconstructed, then it will create a dictionary with an entry for every reconstructed image. No bias is used if a nn.LayerNorm is used after the nn.Conv2dTranspose layer.

Parameters:

  • output_channels (Sequence[int]) – the output channels, one for every image observation.

  • channels_multiplier (int) – the channels multiplier, same for the encoder network.

  • latent_state_size (int) – the size of the latent state. Before applying the decoder, a nn.Linear layer is used to project the latent state to a feature vector of dimension [8 * channels_multiplier, 4, 4].

  • cnn_encoder_output_dim (int) – the output of the image encoder. It should be equal to 8 * channels_multiplier * 4 * 4.

  • image_size (Tuple[int, int]) – the final image size.

  • activation (nn.Module, optional) – the activation function. Defaults to nn.SiLU.

  • layer_norm_cls (Callable[..., nn.Module]) – the layer norm to apply after the input projection. Defaults to LayerNormChannelLast.

  • layer_norm_kw (Dict[str, Any]) – the kwargs of the layer norm. Default to {“eps”: 1e-3}.

  • stages (int) – how many stages in the CNN decoder.

forward(latent_states: torch.Tensor) → List[torch.Tensor][source]
class xuance.mindspore.representations.world_model.CNNEncoder(*args: Any, **kwargs: Any)[source]

Bases: Module

The Dreamer-V3 image encoder. This is composed of 4 nn.Conv2d with kernel_size=3, stride=2 and padding=1. No bias is used if a nn.LayerNorm is used after the convolution. This 4-stages model assumes that the image is a 64x64 and it ends with a resolution of 4x4. If more than one image is to be encoded, then those will be concatenated on the channel dimension and fed to the encoder.

Parameters:

  • input_channels (Sequence[int]) – the input channels, one for each image observation to encode.

  • image_size (Tuple[int, int]) – the image size as (Height,Width).

  • channels_multiplier (int) – the multiplier for the output channels. Given the 4 stages, the 4 output channels will be [1, 2, 4, 8] * channels_multiplier.

  • layer_norm_cls (Callable[..., nn.Module]) – the layer norm to apply after the input projection. Defaults to LayerNormChannelLast.

  • layer_norm_kw (Dict[str, Any]) – the kwargs of the layer norm. Default to {“eps”: 1e-3}.

  • activation (ModuleType, optional) – the activation function. Defaults to nn.SiLU.

  • stages (int, optional) – how many stages for the CNN.

forward(obs: torch.Tensor) → torch.Tensor[source]
class xuance.mindspore.representations.world_model.DreamerV3WorldModel(*args: Any, **kwargs: Any)[source]

Bases: Module

obs_space

world_model, actor, critic, target_critic for agent: player (link to policy.world_model.~ & policy.actor)

Type:

for policy

class xuance.mindspore.representations.world_model.MLPDecoder(*args: Any, **kwargs: Any)[source]

Bases: Module

The exact inverse of the MLPEncoder. This is composed of N nn.Linear layers, where N is specified by mlp_layers. No bias is used if a nn.LayerNorm is used after the linear layer. If more than one vector is to be decoded, then it will create a dictionary with an entry for every reconstructed vector.

Parameters:

  • keys (Sequence[str]) – the keys representing the vector observations to decode.

  • output_dims (Sequence[int]) – the dimensions of every vector to decode.

  • latent_state_size (int) – the dimension of the latent state.

  • mlp_layers (int, optional) – how many mlp layers. Defaults to 4.

  • dense_units (int, optional) – the dimension of every mlp. Defaults to 512.

  • layer_norm_cls (Callable[..., nn.Module]) – the layer norm to apply after the input projection. Defaults to LayerNorm.

  • layer_norm_kw (Dict[str, Any]) – the kwargs of the layer norm. Default to {“eps”: 1e-3}.

  • activation (ModuleType, optional) – the activation function after every layer. Defaults to nn.SiLU.

forward(latent_states: torch.Tensor) → Dict[str, torch.Tensor][source]
class xuance.mindspore.representations.world_model.MLPEncoder(*args: Any, **kwargs: Any)[source]

Bases: Module

The Dreamer-V3 vector encoder. This is composed of N nn.Linear layers, where N is specified by mlp_layers. No bias is used if a nn.LayerNorm is used after the linear layer. If more than one vector is to be encoded, then those will concatenated on the last dimension before being fed to the encoder.

Parameters:

  • input_dims (Sequence[int]) – the dimensions of every vector to encode.

  • mlp_layers (int, optional) – how many mlp layers. Defaults to 4.

  • dense_units (int, optional) – the dimension of every mlp. Defaults to 512.

  • layer_norm_cls (Callable[..., nn.Module]) – the layer norm to apply after the input projection. Defaults to LayerNorm.

  • layer_norm_kw (Dict[str, Any]) – the kwargs of the layer norm. Default to {“eps”: 1e-3}.

  • activation (ModuleType, optional) – the activation function after every layer. Defaults to nn.SiLU.

  • sym_log_inputs (bool, optional) – whether to squash the input with the sym_log function. Defaults to True.

forward(obs: torch.Tensor) → torch.Tensor[source]
class xuance.mindspore.representations.world_model.PlayerDV3(*args: Any, **kwargs: Any)[source]

Bases: Module

The model of the Dreamer_v3 player.

Parameters:

  • encoder (MultiEncoder) – the encoder.

  • rssm (RSSM) – the RSSM model.

  • actor (Module) – the actor.

  • actions_dim (Sequence[int]) – the dimension of the actions.

  • num_envs (int) – the number of environments.

  • stochastic_size (int) – the size of the stochastic state.

  • recurrent_state_size (int) – the size of the recurrent state.

  • transition_model (Module) – the transition model.

  • discrete_size (int) – the dimension of a single Categorical variable in the stochastic state (prior or posterior). Defaults to 32.

  • actor_type (str, optional) –

    which actor the player is using (‘task’ or ‘exploration’).

    Default to None.

    (bool, optional): whether to use the DecoupledRSSM model.

get_actions(obs: Dict[str, torch.Tensor], greedy: bool = False, mask: Dict[str, torch.Tensor] | None = None) → Sequence[torch.Tensor][source]

Return the greedy actions.

Parameters:

  • obs (Dict[str, Tensor]) – the current observations.

  • greedy (bool) – whether or not to sample the actions. Default to False.

  • mask (Optional[Dict[str, Tensor]]) – action mask

Returns:

The actions the agent has to perform.

init_states(reset_envs: Sequence[int] | None = None, num_envs: int | None = None) → None

Initialize the states and the actions for the ended environments.

Parameters:

  • reset_envs (Optional[Sequence[int]], optional) – which environments’ states to reset. If None, then all environments’ states are reset. Defaults to None.

  • num_envs (Optional[int]) – the number of environments. If None, then it will be self.num_envs # prop added to deal with xuance test

class xuance.mindspore.representations.world_model.RSSM(*args: Any, **kwargs: Any)[source]

Bases: Module

RSSM model for the model-base Dreamer agent.

Parameters:

  • recurrent_model (nn.Module) – the recurrent model of the RSSM model described in [https://arxiv.org/abs/1811.04551](https://arxiv.org/abs/1811.04551).

  • representation_model (nn.Module) – the representation model composed by a multi-layer perceptron to compute the stochastic part of the latent state. For more information see [https://arxiv.org/abs/2010.02193](https://arxiv.org/abs/2010.02193).

  • transition_model (nn.Module) – the transition model described in [https://arxiv.org/abs/2010.02193](https://arxiv.org/abs/2010.02193). The model is composed by a multi-layer perceptron to predict the stochastic part of the latent state.

  • distribution_config (Dict[str, Any]) – the configs of the distributions.

  • discrete (int, optional) – the size of the Categorical variables. Defaults to 32.

  • unimix – (float, optional): the percentage of uniform distribution to inject into the categorical distribution over states, i.e. given some logits l and probabilities p = softmax(l), then p = (1 - self.unimix) * p + self.unimix * unif, where unif = `1 / self.discrete. Defaults to 0.01.

dynamic(posterior: torch.Tensor, recurrent_state: torch.Tensor, action: torch.Tensor, embedded_obs: torch.Tensor, is_first: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor][source]
Perform one step of the dynamic learning:
Recurrent model: compute the recurrent state from the previous latent space, the action taken by the agent,

i.e., it computes the deterministic state (or ht).

Transition model: predict the prior from the recurrent output. Representation model: compute the posterior from the recurrent state and from

the embedded observations provided by the environment.

For more information see [https://arxiv.org/abs/1811.04551](https://arxiv.org/abs/1811.04551) and [https://arxiv.org/abs/2010.02193](https://arxiv.org/abs/2010.02193).

Parameters:

  • posterior (Tensor) – the stochastic state computed by the representation model (posterior). It is expected to be of dimension [stoch_size, self.discrete], which by default is [32, 32].

  • recurrent_state (Tensor) – a tuple representing the recurrent state of the recurrent model.

  • action (Tensor) – the action taken by the agent.

  • embedded_obs (Tensor) – the embedded observations provided by the environment.

  • is_first (Tensor) – if this is the first step in the episode.

Returns:

the recurrent state of the recurrent model. The posterior stochastic state (Tensor): computed by the representation model The prior stochastic state (Tensor): computed by the transition model The logits of the posterior state (Tensor): computed by the transition model from the recurrent state. The logits of the prior state (Tensor): computed by the transition model from the recurrent state. from the recurrent state and the embbedded observation.

Return type:

The recurrent state (Tensor)

get_initial_states(batch_shape: Sequence[int] | torch.Size) → Tuple[torch.Tensor, torch.Tensor][source]
imagination(prior: torch.Tensor, recurrent_state: torch.Tensor, actions: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]

One-step imagination of the next latent state. It can be used several times to imagine trajectories in the latent space (Transition Model).

Parameters:

  • prior (Tensor) – the prior state.

  • recurrent_state (Tensor) – the recurrent state of the recurrent model.

  • actions (Tensor) – the actions taken by the agent.

Returns:

the imagined prior state. The recurrent state (Tensor).

Return type:

The imagined prior state (Tuple[Tensor, Tensor])

class xuance.mindspore.representations.world_model.RecurrentModel(*args: Any, **kwargs: Any)[source]

Bases: Module

Recurrent model for the model-base Dreamer-V3 agent. This implementation uses the models.LayerNormGRUCell, which combines the standard GRUCell from PyTorch with the nn.LayerNorm, where the normalization is applied right after having computed the projection from the input to the weight space.

Parameters:

  • input_size (int) – the input size of the model.

  • dense_units (int) – the number of dense units.

  • recurrent_state_size (int) – the size of the recurrent state.

  • activation_fn (nn.Module) – the activation function. Default to SiLU.

  • layer_norm_cls (Callable[..., nn.Module]) – the layer norm to apply after the input projection. Defaults to LayerNorm.

  • layer_norm_kw (Dict[str, Any]) – the kwargs of the layer norm. Default to {“eps”: 1e-3}.

forward(input: torch.Tensor, recurrent_state: torch.Tensor) → torch.Tensor[source]

Compute the next recurrent state from the latent state (stochastic and recurrent states) and the actions.

Parameters:

  • input (Tensor) – the input tensor composed by the stochastic state and the actions concatenated together.

  • recurrent_state (Tensor) – the previous recurrent state.

Returns:

the computed recurrent output and recurrent state.

class xuance.mindspore.representations.world_model.WorldModel(*args: Any, **kwargs: Any)[source]

Bases: Module

Wrapper class for the World model.

Parameters:

  • encoder (Module) – the encoder.

  • rssm (RSSM) – the rssm.

  • observation_model (Module) – the observation model.

  • reward_model (Module) – the reward model.

  • continue_model (Module, optional) – the continue model.

world_model_v2

class xuance.mindspore.representations.world_model_v2.Actor(*args: Any, **kwargs: Any)[source]

Bases: Module

The wrapper class of the Dreamer_v2 Actor model.

Parameters:

  • latent_state_size (int) – the dimension of the latent state (stochastic size + recurrent_state_size).

  • actions_dim (Sequence[int]) – the dimension in output of the actor. The number of actions if continuous, the dimension of the action if discrete.

  • is_continuous (bool) – whether or not the actions are continuous.

  • distribution_config (Dict[str, Any]) – The configs of the distributions.

  • init_std (float) – the amount to sum to the input of the softplus function for the standard deviation. Default to 5.

  • min_std (float) – the minimum standard deviation for the actions. Default to 0.1.

  • dense_units (int) – the dimension of the hidden dense layers. Default to 400.

  • activation (int) – the activation function to apply after the dense layers. Default to nn.ELU.

  • mlp_layers (int) – the number of linear layers. Default to 4.

  • layer_norm (bool) – whether or not to use the layer norm. Default to False.

  • expl_amount (float) – the exploration amount to use during training. Default to 0.0.

  • expl_decay (float) – the exploration decay to use during training. Default to 0.0.

  • expl_min (float) – the exploration amount minimum to use during training. Default to 0.0.

forward(state: torch.Tensor, greedy: bool = False, mask: Dict[str, torch.Tensor] | None = None) → Tuple[Sequence[torch.Tensor], Sequence[torch.distributions.Distribution]][source]

Call the forward method of the actor model and reorganizes the result with shape (batch_size, *, num_actions), where * means any number of dimensions including None.

Parameters:

  • state (Tensor) – the current state of shape (batch_size, *, stochastic_size + recurrent_state_size).

  • greedy (bool) – whether or not to sample the actions. Default to False.

  • mask (Dict[str, Tensor], optional) – the action mask (which actions can be selected). Default to None.

Returns:

The tensor of the actions taken by the agent with shape (batch_size, *, num_actions). The distribution of the actions

class xuance.mindspore.representations.world_model_v2.CNNDecoder(*args: Any, **kwargs: Any)[source]

Bases: Module

The almost-exact inverse of the CNNEncoder class, where in 4 stages it reconstructs the observation image to 64x64. If multiple images are to be reconstructed, then it will create a dictionary with an entry for every reconstructed image. No bias is used if a nn.LayerNorm is used after the nn.Conv2dTranspose layer.

Parameters:

  • keys (Sequence[str]) – the keys of the image observation to be reconstructed.

  • output_channels (Sequence[int]) – the output channels, one for every image observation.

  • channels_multiplier (int) – the channels multiplier, same for the encoder network.

  • latent_state_size (int) – the size of the latent state. Before applying the decoder, a nn.Linear layer is used to project the latent state to a feature vector.

  • cnn_encoder_output_dim (int) – the output of the image encoder.

  • image_size (Tuple[int, int]) – the final image size.

  • activation (nn.Module, optional) – the activation function. Defaults to nn.ELU.

  • layer_norm (bool, optional) – whether to apply the layer normalization. Defaults to True.

forward(latent_states: torch.Tensor) → List[torch.Tensor][source]
class xuance.mindspore.representations.world_model_v2.CNNEncoder(*args: Any, **kwargs: Any)[source]

Bases: Module

The Dreamer-V2 image encoder. This is composed of 4 nn.Conv2d with kernel_size=3, stride=2 and padding=1. No bias is used if a nn.LayerNorm is used after the convolution. This 4-stages model assumes that the image is a 64x64. If more than one image is to be encoded, then those will be concatenated on the channel dimension and fed to the encoder.

Parameters:

  • keys (Sequence[str]) – the keys representing the image observations to encode.

  • input_channels (Sequence[int]) – the input channels, one for each image observation to encode.

  • image_size (Tuple[int, int]) – the image size as (Height,Width).

  • channels_multiplier (int) – the multiplier for the output channels. Given the 4 stages, the 4 output channels will be [1, 2, 4, 8] * channels_multiplier.

  • layer_norm (bool, optional) – whether to apply the layer normalization. Defaults to True.

  • activation (ModuleType, optional) – the activation function. Defaults to nn.ELU.

forward(obs: torch.Tensor) → torch.Tensor[source]
class xuance.mindspore.representations.world_model_v2.DreamerV2WorldModel(*args: Any, **kwargs: Any)[source]

Bases: Module

obs_space

world_model, actor, critic, target_critic for agent: player (link to policy.world_model.~ & policy.actor)

Type:

for policy

class xuance.mindspore.representations.world_model_v2.MLPDecoder(*args: Any, **kwargs: Any)[source]

Bases: Module

The exact inverse of the MLPEncoder. This is composed of N nn.Linear layers, where N is specified by mlp_layers. No bias is used if a nn.LayerNorm is used after the linear layer. If more than one vector is to be decoded, then it will create a dictionary with an entry for every reconstructed vector.

Parameters:

  • keys (Sequence[str]) – the keys representing the vector observations to decode.

  • output_dims (Sequence[int]) – the dimensions of every vector to decode.

  • latent_state_size (int) – the dimension of the latent state.

  • mlp_layers (int, optional) – how many mlp layers. Defaults to 4.

  • dense_units (int, optional) – the dimension of every mlp. Defaults to 512.

  • layer_norm (bool, optional) – whether to apply the layer normalization. Defaults to True.

  • activation (ModuleType, optional) – the activation function after every layer. Defaults to nn.ELU.

forward(latent_states: torch.Tensor) → Dict[str, torch.Tensor][source]
class xuance.mindspore.representations.world_model_v2.MLPEncoder(*args: Any, **kwargs: Any)[source]

Bases: Module

The Dreamer-V3 vector encoder. This is composed of N nn.Linear layers, where N is specified by mlp_layers. No bias is used if a nn.LayerNorm is used after the linear layer. If more than one vector is to be encoded, then those will concatenated on the last dimension before being fed to the encoder.

Parameters:

  • keys (Sequence[str]) – the keys representing the vector observations to encode.

  • input_dims (Sequence[int]) – the dimensions of every vector to encode.

  • mlp_layers (int, optional) – how many mlp layers. Defaults to 4.

  • dense_units (int, optional) – the dimension of every mlp. Defaults to 512.

  • layer_norm (bool, optional) – whether to apply the layer normalization. Defaults to True.

  • activation (ModuleType, optional) – the activation function after every layer. Defaults to nn.ELU.

forward(obs: torch.Tensor) → torch.Tensor[source]
class xuance.mindspore.representations.world_model_v2.PlayerDV2(*args: Any, **kwargs: Any)[source]

Bases: Module

The model of the Dreamer_v2 player.

Parameters:

  • encoder (nn.Module) – the encoder.

  • recurrent_model (nn.Module) – the recurrent model.

  • representation_model (nn.Module) – the representation model.

  • actor (nn.Module) – the actor.

  • actions_dim (Sequence[int]) – the dimension of the actions.

  • num_envs (int) – the number of environments.

  • stochastic_size (int) – the size of the stochastic state.

  • recurrent_state_size (int) – the size of the recurrent state.

  • device (str | torch.device) – the device where the model is stored.

  • discrete_size (int) – the dimension of a single Categorical variable in the stochastic state (prior or posterior). Defaults to 32.

  • actor_type (str, optional) – which actor the player is using (‘task’ or ‘exploration’). Default to None.

get_actions(obs: Dict[str, torch.Tensor], greedy: bool = False, mask: Dict[str, torch.Tensor] | None = None) → Sequence[torch.Tensor][source]

Return the greedy actions.

Parameters:

  • obs (Dict[str, Tensor]) – the current observations.

  • greedy (bool) – whether or not to sample the actions. Default to False.

  • mask (Dict[str, Tensor], optional) – the action mask (which actions can be selected). Default to None.

Returns:

The actions the agent has to perform.

init_states(reset_envs: Sequence[int] | None = None, num_envs: int | None = None) → None[source]

Initialize the states and the actions for the ended environments.

Parameters:

  • reset_envs (Optional[Sequence[int]], optional) – which environments’ states to reset. If None, then all environments’ states are reset. Defaults to None.

  • num_envs (Optional[int]) – the number of environments. If None, then it will be self.num_envs # prop added to deal with xuance test

class xuance.mindspore.representations.world_model_v2.RSSM(*args: Any, **kwargs: Any)[source]

Bases: Module

RSSM model for the model-base Dreamer agent.

Parameters:
dynamic(posterior: torch.Tensor, recurrent_state: torch.Tensor, action: torch.Tensor, embedded_obs: torch.Tensor, is_first: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor][source]
Perform one step of the dynamic learning:
Recurrent model: compute the recurrent state from the previous latent space, the action taken by the agent,

i.e., it computes the deterministic state (or ht).

Transition model: predict the prior from the recurrent output. Representation model: compute the posterior from the recurrent state and from

the embedded observations provided by the environment.

For more information see [https://arxiv.org/abs/1811.04551](https://arxiv.org/abs/1811.04551) and [https://arxiv.org/abs/2010.02193](https://arxiv.org/abs/2010.02193).

Parameters:

  • posterior (Tensor) – the stochastic state computed by the representation model (posterior). It is expected to be of dimension [stoch_size, self.discrete], which by default is [32, 32].

  • recurrent_state (Tensor) – a tuple representing the recurrent state of the recurrent model.

  • action (Tensor) – the action taken by the agent.

  • embedded_obs (Tensor) – the embedded observations provided by the environment.

  • is_first (Tensor) – if this is the first step in the episode.

Returns:

the recurrent state of the recurrent model. The posterior stochastic state (Tensor): computed by the representation model The prior stochastic state (Tensor): computed by the transition model The logits of the posterior state (Tensor): computed by the transition model from the recurrent state. The logits of the prior state (Tensor): computed by the transition model from the recurrent state. from the recurrent state and the embbedded observation.

Return type:

The recurrent state (Tensor)

imagination(prior: torch.Tensor, recurrent_state: torch.Tensor, actions: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]

One-step imagination of the next latent state. It can be used several times to imagine trajectories in the latent space (Transition Model).

Parameters:

  • prior (Tensor) – the prior state.

  • recurrent_state (Tensor) – the recurrent state of the recurrent model.

  • actions (Tensor) – the actions taken by the agent.

Returns:

the imagined prior state. The recurrent state (Tensor).

Return type:

The imagined prior state (Tuple[Tensor, Tensor])

class xuance.mindspore.representations.world_model_v2.RecurrentModel(*args: Any, **kwargs: Any)[source]

Bases: Module

Recurrent model for the model-base Dreamer-V3 agent. This implementation uses the models.LayerNormGRUCell, which combines the standard GRUCell from PyTorch with the nn.LayerNorm, where the normalization is applied right after having computed the projection from the input to the weight space.

Parameters:

  • input_size (int) – the input size of the model.

  • recurrent_state_size (int) – the size of the recurrent state.

  • dense_units (int) – the number of dense units.

  • activation (nn.Module) – the activation function. Default to ELU.

  • layer_norm (bool) – whether to use the LayerNorm inside the GRU. Defaults to True.

forward(input: torch.Tensor, recurrent_state: torch.Tensor) → torch.Tensor[source]

Compute the next recurrent state from the latent state (stochastic and recurrent states) and the actions.

Parameters:

  • input (Tensor) – the input tensor composed by the stochastic state and the actions concatenated together.

  • recurrent_state (Tensor) – the previous recurrent state.

Returns:

the computed recurrent output and recurrent state.

class xuance.mindspore.representations.world_model_v2.WorldModel(*args: Any, **kwargs: Any)[source]

Bases: Module

Wrapper class for the World model.

Parameters:

  • encoder (nn.Module) – the encoder.

  • rssm (RSSM) – the rssm.

  • observation_model (nn.Module) – the observation model.

  • reward_model (nn.Module) – the reward model.

  • continue_model (nn.Module, optional) – the continue model.

encoder: MultiEncoder