utils

device

xuance.tensorflow.utils.device.collect_device_info(rank: int = 0, agent=None) → dict

Collect runtime device / system info for reproducibility (TensorFlow 2.x).

Returns a JSON-serializable dict.

xuance.tensorflow.utils.device.set_device(expected_device: str)

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.tensorflow.utils.distributions.ActivatedDiagGaussianDistribution(action_dim: int, activation_action)

Bases: DiagGaussianDistribution

activated_rsample()

activated_rsample_and_logprob()

class xuance.tensorflow.utils.distributions.CategoricalDistribution(action_dim: int)

Bases: Distribution

deterministic_sample()

entropy()

get_param()

kl_divergence(other: Distribution)

log_prob(x)

set_param(probs=None, logits=None)

stochastic_sample()

class xuance.tensorflow.utils.distributions.DiagGaussianDistribution(action_dim: int)

Bases: Distribution

deterministic_sample()

entropy()

get_param()

kl_divergence(other: Distribution)

log_prob(x)

set_param(mu, std)

stochastic_sample()

class xuance.tensorflow.utils.distributions.Distribution

Bases: object

abstractmethod deterministic_sample()

abstractmethod entropy()

abstractmethod get_param()

abstractmethod log_prob(x: tensorflow.Tensor)

abstractmethod set_param(*args)

abstractmethod stochastic_sample()

xuance.tensorflow.utils.distributions.merge_distributions(distribution_list)

xuance.tensorflow.utils.distributions.split_distributions(distribution)

harmonizer

class xuance.tensorflow.utils.harmonizer.Harmonizer(*args: Any, **kwargs: Any)

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)

get_harmony()

layers

xuance.tensorflow.utils.layers.cnn_block(input_shape: Sequence[int], filters: int, kernel_size: int, stride: int, normalize: tensorflow.keras.layers.Layer | None = None, activation: tensorflow.keras.layers.Layer | None = None, initializer: tensorflow.keras.initializers.Initializer | None = None)

xuance.tensorflow.utils.layers.gru_block(input_dim: Sequence[int], output_dim: int, num_layers: int = 1, dropout: float = 0, initialize: Callable[[tensorflow.Tensor], tensorflow.Tensor] | None = None)

xuance.tensorflow.utils.layers.lstm_block(input_dim: Sequence[int], output_dim: int, num_layers: int = 1, dropout: float = 0, initialize: Callable[[tensorflow.Tensor], tensorflow.Tensor] | None = None)

xuance.tensorflow.utils.layers.mlp_block(input_dim: int, output_dim: int, normalize: tensorflow.keras.layers.Layer | None = None, activation: tensorflow.keras.layers.Layer | None = None, initializer: tensorflow.keras.initializers.Initializer | None = None)

xuance.tensorflow.utils.layers.pooling_block(input_shape: Sequence[int], scale: int, pooling: tensorflow.keras.layers.Layer | None = None) → Sequence[Type[tensorflow.keras.Model]]

layers4dreamer

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

class xuance.tensorflow.utils.layers4dreamer.CNN(*args: Any, **kwargs: Any)

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

property model: torch.nn.Module

property output_dim: int

class xuance.tensorflow.utils.layers4dreamer.DeCNN(*args: Any, **kwargs: Any)

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

property model: torch.nn.Module

property output_dim: int

class xuance.tensorflow.utils.layers4dreamer.LayerNorm(*args: Any, **kwargs: Any)

Bases: LayerNorm

forward(x: torch.Tensor) → torch.Tensor

class xuance.tensorflow.utils.layers4dreamer.LayerNormChannelLast(*args: Any, **kwargs: Any)

Bases: LayerNorm

forward(x: torch.Tensor) → torch.Tensor

class xuance.tensorflow.utils.layers4dreamer.LayerNormGRUCell(*args: Any, **kwargs: Any)

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

class xuance.tensorflow.utils.layers4dreamer.MLP(*args: Any, **kwargs: Any)

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

property model: torch.nn.Module

property output_dim: int

class xuance.tensorflow.utils.layers4dreamer.MultiDecoder(*args: Any, **kwargs: Any)

Bases: Module

property cnn_keys: Sequence[str]

forward(x: torch.Tensor) → torch.Tensor

property mlp_keys: Sequence[str]

class xuance.tensorflow.utils.layers4dreamer.MultiEncoder(*args: Any, **kwargs: Any)

Bases: Module

property cnn_keys: Sequence[str]

forward(obs: Dict[str, torch.Tensor], *args, **kwargs) → torch.Tensor

property mlp_keys: Sequence[str]

class xuance.tensorflow.utils.layers4dreamer.NatureCNN(*args: Any, **kwargs: Any)

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

property output_dim: int

xuance.tensorflow.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

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.tensorflow.utils.layers4dreamer.create_layer_with_args(layer_type: Type[torch.nn.Module] | None, layer_args: Tuple[Any, ...] | Dict[Any, Any] | None) → torch.nn.Module

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.tensorflow.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]]

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.tensorflow.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]

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.tensorflow.utils.layers4dreamer.per_layer_ortho_init_weights(module: torch.nn.Module, gain: float = 1.0, bias: float = 0.0)

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

class xuance.tensorflow.utils.operations.MyLinearLR(*args: Any, **kwargs: Any)

Bases: LearningRateSchedule

xuance.tensorflow.utils.operations.assign_from_flat_grads(flat_grads: tensorflow.Tensor, model: tensorflow.keras.Model) → tensorflow.keras.Model

xuance.tensorflow.utils.operations.assign_from_flat_params(flat_params: tensorflow.Tensor, model: tensorflow.keras.Model) → tensorflow.keras.Model

xuance.tensorflow.utils.operations.get_flat_params(model: tensorflow.keras.Model) → tensorflow.Tensor

xuance.tensorflow.utils.operations.set_seed(seed)

xuance.tensorflow.utils.operations.update_linear_decay(optimizer, step, total_steps, initial_lr, end_factor)

value_norm

class xuance.tensorflow.utils.value_norm.ValueNorm(input_shape, norm_axes=1, beta=0.99999, per_element_update=False, epsilon=1e-05)

Bases: object

Normalize a vector of observations - across the first norm_axes dimensions

denormalize(input_vector)

Transform normalized data back into original distribution

normalize(input_vector)

reset_parameters()

running_mean_var()

update(input_vector)