utils

device

xuance.mindspore.utils.device.collect_device_info(rank: int = 0, agent=None) → dict[source]

Collect runtime device / system info for reproducibility (MindSpore).

Returns a JSON-serializable dict.

xuance.mindspore.utils.device.set_device(expected_device: str)[source]

Set the computing device for a given deep learning framework.

Parameters:

  • dl_toolbox (str) – The deep learning framework to use. Options: “torch”, “tensorflow”, “mindspore”.

  • expected_device (str) – The desired computing device. Options: “cuda”, “GPU”, “gpu”, “Ascend”, “cpu”, “CPU.

Returns:

The assigned computing device, which may differ from expected_device if the requested device is unavailable.

Return type:

str

distributions

class xuance.mindspore.utils.distributions.ActivatedDiagGaussianDistribution(action_dim: int, activation_action)[source]

Bases: DiagGaussianDistribution

activated_rsample()[source]
activated_rsample_and_logprob()[source]
class xuance.mindspore.utils.distributions.CategoricalDistribution(action_dim: int)[source]

Bases: Distribution

deterministic_sample()[source]
entropy()[source]
get_param()[source]
kl_divergence(other: Distribution)[source]
log_prob(x)[source]
set_param(probs=None, logits=None)[source]
stochastic_sample()[source]
class xuance.mindspore.utils.distributions.DiagGaussianDistribution(action_dim: int)[source]

Bases: Distribution

deterministic_sample()[source]
entropy()[source]
get_param()[source]
kl_divergence(other: Distribution)[source]
log_prob(x: mindspore.Tensor)[source]
set_param(mu, std)[source]
stochastic_sample()[source]
class xuance.mindspore.utils.distributions.Distribution[source]

Bases: ABC

abstractmethod deterministic_sample()[source]
abstractmethod entropy()[source]
abstractmethod get_param()[source]
abstractmethod log_prob(x: mindspore.Tensor)[source]
abstractmethod set_param(*args)[source]
abstractmethod stochastic_sample()[source]
xuance.mindspore.utils.distributions.merge_distributions(distribution_list)[source]
xuance.mindspore.utils.distributions.split_distributions(distribution)[source]

harmonizer

class xuance.mindspore.utils.harmonizer.Harmonizer(*args: Any, **kwargs: Any)[source]

Bases: Module

Learnable parameter for loss_scale balancing Ref: https://github.com/thuml/HarmonyDream/blob/main/dreamerv3-jax/dreamerv3/nets.py

forward(loss: torch.Tensor, regularize=True)[source]
get_harmony()[source]

layers

xuance.mindspore.utils.layers.cnn_block(input_shape: Sequence[int], filter: int, kernel_size: int, stride: int, normalize: mindspore.nn.BatchNorm2d | mindspore.nn.LayerNorm | mindspore.nn.GroupNorm | mindspore.nn.InstanceNorm2d | None = None, activation: Type[mindspore.nn.Cell] | None = None, initialize: Callable[[mindspore.Tensor], mindspore.Tensor] | None = None) → Tuple[Sequence[Type[mindspore.nn.Cell]], Tuple][source]
xuance.mindspore.utils.layers.gru_block(input_dim: int, output_dim: int, num_layers: int = 1, dropout: float = 0, initialize: Callable[[mindspore.Tensor], mindspore.Tensor] | None = None) → Tuple[mindspore.nn.Cell, int][source]
xuance.mindspore.utils.layers.lstm_block(input_dim: int, output_dim: int, num_layers: int = 1, dropout: float = 0, initialize: Callable[[mindspore.Tensor], mindspore.Tensor] | None = None) → Tuple[mindspore.nn.Cell, int][source]
xuance.mindspore.utils.layers.mlp_block(input_dim: int, output_dim: int, normalize: mindspore.nn.BatchNorm1d | mindspore.nn.LayerNorm | None = None, activation: Type[mindspore.nn.Cell] | None = None, initialize: Callable[[mindspore.Tensor], mindspore.Tensor] | None = None) → Tuple[Sequence[Type[mindspore.nn.Cell]], Tuple[int]][source]
xuance.mindspore.utils.layers.pooling_block(input_shape: Sequence[int], scale: int, pooling: mindspore.nn.AdaptiveMaxPool2d | mindspore.nn.AdaptiveAvgPool2d) → Sequence[Type[mindspore.nn.Cell]][source]

layers4dreamer

Adapted from: https://github.com/thu-ml/tianshou/blob/master/tianshou/utils/net/common.py

class xuance.mindspore.utils.layers4dreamer.CNN(*args: Any, **kwargs: Any)[source]

Bases: Module

Simple CNN backbone.

Parameters:

  • input_channels (int) – dimensions of the input channels.

  • hidden_channels (Sequence[int], optional) – intermediate number of channels for the CNN, including the output channels.

  • dropout_layer (Union[ModuleType, Sequence[ModuleType]], optional) – which dropout layer to be used before activation (possibly before the normalization layer), e.g., nn.Dropout. You can also pass a list of dropout modules with the same length of hidden_sizes to use different dropout modules in different layers. If None, then no dropout layer is used. Defaults to None.

  • norm_layer (Union[ModuleType, Sequence[ModuleType]], optional) – which normalization layer to be used before activation, e.g., nn.LayerNorm and nn.BatchNorm1d. You can also pass a list of normalization modules with the same length of hidden_sizes to use different normalization modules in different layers. If None, then no normalization layer is used. Defaults to None.

  • activation (Union[ModuleType, Sequence[ModuleType]], optional) – which activation to use after each layer, can be both the same activation for all layers if a single nn.Module is passed, or different activations for different layers if a list is passed. Defaults to nn.ReLU.

forward(obs: torch.Tensor) → torch.Tensor[source]
property model: torch.nn.Module
property output_dim: int
class xuance.mindspore.utils.layers4dreamer.DeCNN(*args: Any, **kwargs: Any)[source]

Bases: Module

Simple DeCNN backbone.

Parameters:

  • input_channels (int) – dimensions of the input channels.

  • hidden_channels (Sequence[int], optional) – intermediate number of channels for the CNN, including the output channels.

  • dropout_layer (Union[ModuleType, Sequence[ModuleType]], optional) – which dropout layer to be used before activation (possibly before the normalization layer), e.g., nn.Dropout. You can also pass a list of dropout modules with the same length of hidden_sizes to use different dropout modules in different layers. If None, then no dropout layer is used. Defaults to None.

  • norm_layer (Union[ModuleType, Sequence[ModuleType]], optional) – which normalization layer to be used before activation, e.g., nn.LayerNorm and nn.BatchNorm1d. You can also pass a list of normalization modules with the same length of hidden_sizes to use different normalization modules in different layers. If None, then no normalization layer is used. Defaults to None.

  • activation (Union[ModuleType, Sequence[ModuleType]], optional) – which activation to use after each layer, can be both the same activation for all layers if a single nn.Module is passed, or different activations for different layers if a list is passed. Defaults to nn.ReLU.

forward(obs: torch.Tensor) → torch.Tensor[source]
property model: torch.nn.Module
property output_dim: int
class xuance.mindspore.utils.layers4dreamer.LayerNorm(*args: Any, **kwargs: Any)[source]

Bases: LayerNorm

forward(x: torch.Tensor) → torch.Tensor[source]
class xuance.mindspore.utils.layers4dreamer.LayerNormChannelLast(*args: Any, **kwargs: Any)[source]

Bases: LayerNorm

forward(x: torch.Tensor) → torch.Tensor[source]
class xuance.mindspore.utils.layers4dreamer.LayerNormGRUCell(*args: Any, **kwargs: Any)[source]

Bases: Module

A GRU cell with a LayerNorm, taken from https://github.com/danijar/dreamerv2/blob/main/dreamerv2/common/nets.py#L317.

This particular GRU cell accepts 3-D inputs, with a sequence of length 1, and applies a LayerNorm after the projection of the inputs.

Parameters:

  • input_size (int) – the input size.

  • hidden_size (int) – the hidden state size

  • bias (bool, optional) – whether to apply a bias to the input projection. Defaults to True.

  • batch_first (bool, optional) – whether the first dimension represent the batch dimension or not. Defaults to False.

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

  • layer_norm_kw (Dict[str, Any]) – the kwargs of the layer norm. Default to {}.

forward(input: torch.Tensor, hx: torch.Tensor | None = None) → torch.Tensor[source]
class xuance.mindspore.utils.layers4dreamer.MLP(*args: Any, **kwargs: Any)[source]

Bases: Module

Simple MLP backbone.

Parameters:

  • input_dims (Union[int, Sequence[int]]) – dimensions of the input vector.

  • output_dim (int, optional) – dimension of the output vector. If set to None, there is no final linear layer. Else, a final linear layer is added. Defaults to None.

  • hidden_sizes (Sequence[int], optional) – shape of MLP passed in as a list, not including input_dims and output_dim.

  • dropout_layer (Union[ModuleType, Sequence[ModuleType]], optional) – which dropout layer to be used before activation (possibly before the normalization layer), e.g., nn.Dropout. You can also pass a list of dropout modules with the same length of hidden_sizes to use different dropout modules in different layers. If None, then no dropout layer is used. Defaults to None.

  • norm_layer (Union[ModuleType, Sequence[ModuleType]], optional) – which normalization layer to be used before activation, e.g., nn.LayerNorm and nn.BatchNorm1d. You can also pass a list of normalization modules with the same length of hidden_sizes to use different normalization modules in different layers. If None, then no normalization layer is used. Defaults to None.

  • activation (Union[ModuleType, Sequence[ModuleType]], optional) – which activation to use after each layer, can be both the same activation for all layers if a single nn.Module is passed, or different activations for different layers if a list is passed. Defaults to nn.ReLU.

  • flatten_dim (int, optional) – whether to flatten input data. The flatten dimension starts from 1. Defaults to True.

property flatten_dim: int | None
forward(obs: torch.Tensor) → torch.Tensor[source]
property model: torch.nn.Module
property output_dim: int
class xuance.mindspore.utils.layers4dreamer.MultiDecoder(*args: Any, **kwargs: Any)[source]

Bases: Module

property cnn_keys: Sequence[str]
forward(x: torch.Tensor) → torch.Tensor[source]
property mlp_keys: Sequence[str]
class xuance.mindspore.utils.layers4dreamer.MultiEncoder(*args: Any, **kwargs: Any)[source]

Bases: Module

property cnn_keys: Sequence[str]
forward(obs: Dict[str, torch.Tensor], *args, **kwargs) → torch.Tensor[source]
property mlp_keys: Sequence[str]
class xuance.mindspore.utils.layers4dreamer.NatureCNN(*args: Any, **kwargs: Any)[source]

Bases: CNN

CNN from DQN Nature paper: Mnih, Volodymyr, et al. “Human-level control through deep reinforcement learning.” Nature 518.7540 (2015): 529-533.

Parameters:

  • in_channels (int) – the input channels to the first convolutional layer

  • features_dim (int) – the features dimension in output from the last convolutional layer

  • screen_size (int, optional) – the dimension of the input image as a single integer. Needed to extract the features and compute the output dimension after all the convolutional layers. Defaults to 64.

forward(x: torch.Tensor) → torch.Tensor[source]
property output_dim: int
xuance.mindspore.utils.layers4dreamer.cnn_forward(model: torch.nn.Module, input: torch.Tensor, input_dim: torch.Size | Tuple[int, ...], output_dim: torch.Size | Tuple[int, ...]) → torch.Tensor[source]

Compute the forward of a Convolutional neural network. It flattens all the dimensions before the model input_size, i.e., the dimensions before the (C_in, H, W) dimensions for the encoder and the dimensions before the (feature_size,) dimension for the decoder.

Parameters:

  • model (nn.Module) – the model.

  • input (Tensor) – the input tensor of dimension (*, C_in, H, W) or (*, feature_size), where * means any number of dimensions including None.

  • input_dim (Union[torch.Size, Tuple[int, ...]]) – the input dimensions, i.e., either (C_in, H, W) or (feature_size,).

  • output_dim (Union[torch.Size, Tuple[int, ...]]) – the desired dimensions in output.

Returns:

The output of dimensions (*, *output_dim).

Examples

>>> encoder
CNN(
    (network): Sequential(
        (0): Conv2d(3, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): ReLU()
        (2): Conv2d(4, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (3): ReLU()
        (4): Flatten(start_dim=1, end_dim=-1)
        (5): Linear(in_features=128, out_features=25, bias=True)
    )
)
>>> input = torch.rand(10, 20, 3, 4, 4)
>>> cnn_forward(encoder, input, (3, 4, 4), -1).shape
torch.Size([10, 20, 25])

>>> decoder
Sequential(
    (0): Linear(in_features=230, out_features=1024, bias=True)
    (1): Unflatten(dim=-1, unflattened_size=(1024, 1, 1))
    (2): ConvTranspose2d(1024, 128, kernel_size=(5, 5), stride=(2, 2))
    (3): ReLU()
    (4): ConvTranspose2d(128, 64, kernel_size=(5, 5), stride=(2, 2))
    (5): ReLU()
    (6): ConvTranspose2d(64, 32, kernel_size=(6, 6), stride=(2, 2))
    (7): ReLU()
    (8): ConvTranspose2d(32, 3, kernel_size=(6, 6), stride=(2, 2))
)
>>> input = torch.rand(10, 20, 230)
>>> cnn_forward(decoder, input, (230,), (3, 64, 64)).shape
torch.Size([10, 20, 3, 64, 64])
xuance.mindspore.utils.layers4dreamer.create_layer_with_args(layer_type: Type[torch.nn.Module] | None, layer_args: Tuple[Any, ...] | Dict[Any, Any] | None) → torch.nn.Module[source]

Create a single layer with given layer type and arguments.

Parameters:

  • layer_type (ModuleType) – the type of the layer to be created.

  • layer_args (ArgType, optional) – the arguments to be passed to the layer.

xuance.mindspore.utils.layers4dreamer.create_layers(layer_type: Type[torch.nn.Module] | None | List[Type[torch.nn.Module] | None], layer_args: Tuple[Any, ...] | Dict[Any, Any] | None | List[Tuple[Any, ...] | Dict[Any, Any] | None], num_layers: int) → Tuple[List[Type[torch.nn.Module] | None], Tuple[Any, ...] | Dict[Any, Any] | None | List[Tuple[Any, ...] | Dict[Any, Any] | None]][source]

Create a list of layers with given layer type and arguments.

If a layer_type is not specified, then the lists will be filled with None. If the layer type or the layer arguments are specified only once, they will be cast to a sequence of length num_layers.

Parameters:

  • layer_type (Union[ModuleType, Sequence[ModuleType]]) – the type of the layer to be created.

  • layer_args (ArgsType, optional) – the arguments to be passed to the layer.

  • num_layers (int) – the number of layers to be created.

Returns:

a list of layers and a list of args.

Return type:

Tuple[Sequence[ModuleType], ArgsType]

Examples

>>> create_layers(nn.Linear, None, 3)
([nn.Linear, nn.Linear, nn.Linear], [None, None, None])

>>> create_layers(nn.Linear, {"arg1":3, "arg2": "foo"}, 3)
(
    [nn.Linear, nn.Linear, nn.Linear],
    [{'arg1': 3, 'arg2': 'foo'}, {'arg1': 3, 'arg2': 'foo'}, {'arg1': 3, 'arg2': 'foo'}]
)

>>> create_layers([nn.Linear, nn.Conv2d], [{"bias":False}, {"kernel_size": 5, "bias": True}], 2)
([nn.Linear, nn.Conv2d], [{'bias': False}, {'kernel_size':5, 'bias': True}])

>>> create_layers([nn.Linear, nn.Linear], (64, 10), 2)
([nn.Linear, nn.Linear], [(64, 10), (64, 10)])
xuance.mindspore.utils.layers4dreamer.miniblock(input_size: int, output_size: int, layer_type: Type[torch.nn.Module] = torch.nn.Linear, layer_args: Tuple[Any, ...] | Dict[Any, Any] | None = None, dropout_layer: Type[torch.nn.Module] | None = None, dropout_args: Tuple[Any, ...] | Dict[Any, Any] | None = None, norm_layer: Type[torch.nn.Module] | None = None, norm_args: Tuple[Any, ...] | Dict[Any, Any] | None = None, activation: Type[torch.nn.Module] | None = None, act_args: Tuple[Any, ...] | Dict[Any, Any] | None = None) → List[torch.nn.Module][source]

Construct a miniblock with given input/output-size, dropout layer, norm layer and activation function.

Based on Tianshou’s miniblock function (https://github.com/thu-ml/tianshou/blob/master/tianshou/utils/net/common.py).

Parameters:

  • input_size (int) – the input size of the miniblock (in_features for Linear and in_channels for Conv2d).

  • output_size (int) – the output size of the miniblock.

  • layer_type (Type[nn.Linear], optional) – the type of the layer to be created. Defaults to nn.Linear.

  • layer_args (ArgType, optional) – the arguments to be passed to the layer. Defaults to None.

  • dropout_layer (ModuleType, optional) – the type of the dropout layer to be created. Defaults to None.

  • dropout_args (ArgType, optional) – the arguments to be passed to the dropout layer. Defaults to None.

  • norm_layer (ModuleType, optional) – the type of the norm layer to be created. Defaults to None.

  • norm_args (ArgType, optional) – the arguments to be passed to the norm layer. Defaults to None.

  • activation (ModuleType, optional) – the type of the activation function to be created. Defaults to None.

  • act_args (Tuple[Any, ...] | Dict[Any, Any] | None, optional) – the arguments to be passed to the activation function. Defaults to None.

Returns:

the miniblock as a list of layers.

Return type:

List[nn.Module]

xuance.mindspore.utils.layers4dreamer.per_layer_ortho_init_weights(module: torch.nn.Module, gain: float = 1.0, bias: float = 0.0)[source]

Initialize the weights of a module with orthogonal weights.

Parameters:

  • module (nn.Module) – module to initialize

  • gain (float, optional) – gain of the orthogonal initialization. Defaults to 1.0.

  • bias (float, optional) – bias of the orthogonal initialization. Defaults to 0.0.

operations

xuance.mindspore.utils.operations.assign_from_flat_grads(flat_grads: mindspore.Tensor, model: mindspore.nn.Cell) → mindspore.nn.Cell[source]
xuance.mindspore.utils.operations.assign_from_flat_params(flat_params: mindspore.Tensor, model: mindspore.nn.Cell) → mindspore.nn.Cell[source]
xuance.mindspore.utils.operations.clip_grads(grads, low, high)[source]
xuance.mindspore.utils.operations.get_flat_grad(y: mindspore.Tensor, model: mindspore.nn.Cell) → mindspore.Tensor[source]
xuance.mindspore.utils.operations.get_flat_params(model: mindspore.nn.Cell) → mindspore.Tensor[source]
xuance.mindspore.utils.operations.set_seed(seed)[source]
xuance.mindspore.utils.operations.update_linear_decay(optimizer, step, total_steps, initial_lr, end_factor)[source]

value_norm

class xuance.mindspore.utils.value_norm.ValueNorm(*args: Any, **kwargs: Any)[source]

Bases: Cell

Normalize a vector of observations - across the first norm_axes dimensions

denormalize(input_vector)[source]

Transform normalized data back into original distribution

normalize(input_vector)[source]
reset_parameters()[source]
running_mean_var()[source]
update(input_vector)[source]