utils

device

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

Collect runtime device / system info for reproducibility.

Returns a JSON-serializable dict.

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

Bases: DiagGaussianDistribution

activated_rsample()

activated_rsample_and_logprob()

class xuance.torch.utils.distributions.BernoulliSafeMode(*args: Any, **kwargs: Any)

Bases: Bernoulli

property mode

class xuance.torch.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.torch.utils.distributions.DiagGaussianDistribution(action_dim: int)

Bases: Distribution

deterministic_sample()

entropy()

get_param()

kl_divergence(other: Distribution)

log_prob(x)

rsample()

set_param(mu, std)

stochastic_sample()

class xuance.torch.utils.distributions.Distribution

Bases: ABC

abstractmethod deterministic_sample()

abstractmethod entropy()

abstractmethod get_param()

abstractmethod log_prob(x: torch.Tensor)

abstractmethod set_param(*args)

abstractmethod stochastic_sample()

class xuance.torch.utils.distributions.MSEDistribution(mode: torch.Tensor, dims: int, agg: str = 'sum')

Bases: object

log_prob(value: torch.Tensor) → torch.Tensor

property mean: torch.Tensor

property mode: torch.Tensor

class xuance.torch.utils.distributions.SymLogDistribution(mode: torch.Tensor, dims: int, dist: str = 'mse', agg: str = 'sum', tol: float = 1e-08)

Bases: object

log_prob(value: torch.Tensor) → torch.Tensor

Computes the log probability of a value under this distribution.

Parameters:

value – The observed value (in original space) to evaluate.

Returns:

Log probability tensor, aggregated over event dimensions as specified.

Raises:

• AssertionError – If value shape does not match the distribution’s shape.

• NotImplementedError – If invalid distance or aggregation methods are provided.

property mean: torch.Tensor

property mode: torch.Tensor

class xuance.torch.utils.distributions.TruncatedNormal(*args: Any, **kwargs: Any)

Bases: TruncatedStandardNormal

Truncated Normal distribution https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf

cdf(value)

has_rsample = True

icdf(value)

log_prob(value)

class xuance.torch.utils.distributions.TruncatedStandardNormal(*args: Any, **kwargs: Any)

Bases: Distribution

Truncated Standard Normal distribution https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf

arg_constraints = {'a': torch.distributions.constraints.real, 'b': torch.distributions.constraints.real}

property auc

cdf(value)

entropy()

has_rsample = True

icdf(value)

log_prob(value)

property mean

rsample(sample_shape=torch.Size)

support()

property variance

class xuance.torch.utils.distributions.TwoHotEncodingDistribution(logits: torch.Tensor, dims: int = 0, low: int = -20, high: int = 20, transfwd: ~typing.Callable[[torch.Tensor], torch.Tensor] = <function sym_log>, transbwd: ~typing.Callable[[torch.Tensor], torch.Tensor] = <function sym_exp>)

Bases: object

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

property mean: torch.Tensor

property mode: torch.Tensor

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

Merges a list of individual distributions back into a batch distribution.

This function reconstructs a batched distribution from a list (or array) of individual distributions, supporting both categorical and diagonal Gaussian distributions.

Parameters:

distribution_list (list or np.ndarray) – A list or array of individual distribution instances.

Returns:

A merged batch distribution.

Return type:

CategoricalDistribution or DiagGaussianDistribution

Raises:

NotImplementedError – If the distribution type is not supported.

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

Splits a batch of distributions into individual instances.

This function separates a batch of distributions (either CategoricalDistribution or DiagGaussianDistribution) into individual distribution objects.

Parameters:

distribution (CategoricalDistribution or DiagGaussianDistribution) – The input distribution batch to be split.

Returns:

A reshaped array of individual distribution instances.

Return type:

np.ndarray

Raises:

NotImplementedError – If the distribution type is not supported.

harmonizer

class xuance.torch.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

class xuance.torch.utils.layers.Moments(*args: Any, **kwargs: Any)

Bases: Module

forward(x: torch.Tensor) → Any

xuance.torch.utils.layers.cnn_block(input_shape: Sequence[int], filter: int, kernel_size: int, stride: int, normalize: torch.nn.BatchNorm2d | torch.nn.LayerNorm | torch.nn.GroupNorm | torch.nn.InstanceNorm2d | None = None, activation: Type[torch.nn.Module] | None = None, initialize: Callable[[torch.Tensor], torch.Tensor] | None = None, device: str | int | torch.device | None = None) → Tuple[Sequence[Type[torch.nn.Module]], Tuple]

xuance.torch.utils.layers.gru_block(input_dim: int, output_dim: int, num_layers: int = 1, dropout: float = 0, initialize: Callable[[torch.Tensor], torch.Tensor] | None = None, device: str | int | torch.device | None = None) → Tuple[torch.nn.Module, int]

xuance.torch.utils.layers.lstm_block(input_dim: int, output_dim: int, num_layers: int = 1, dropout: float = 0, initialize: Callable[[torch.Tensor], torch.Tensor] | None = None, device: str | int | torch.device | None = None) → Tuple[torch.nn.Module, int]

xuance.torch.utils.layers.mlp_block(input_dim: int, output_dim: int, normalize: torch.nn.BatchNorm1d | torch.nn.LayerNorm | None = None, activation: Type[torch.nn.Module] | None = None, initialize: Callable[[torch.Tensor], torch.Tensor] | None = None, device: str | int | torch.device | None = None) → Tuple[Sequence[Type[torch.nn.Module]], Tuple[int]]

xuance.torch.utils.layers.pooling_block(input_shape: Sequence[int], scale: int, pooling: torch.nn.AdaptiveMaxPool2d | torch.nn.AdaptiveAvgPool2d, device: str | int | torch.device | None = None) → Sequence[Type[torch.nn.Module]]

layers4dreamer

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

class xuance.torch.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.torch.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.torch.utils.layers4dreamer.LayerNorm(*args: Any, **kwargs: Any)

Bases: LayerNorm

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

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

Bases: LayerNorm

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

class xuance.torch.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.torch.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.torch.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.torch.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.torch.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.torch.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.torch.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.torch.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.torch.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.torch.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

xuance.torch.utils.operations.assign_from_flat_grads(flat_grads: torch.Tensor, model: torch.nn.Module) → torch.nn.Module

Assigns gradients from a flat tensor to the corresponding model parameters.

This function takes a flattened gradient tensor and reshapes it to match the model’s parameter shapes before assigning the values to the respective parameter gradients.

Parameters:

• flat_grads (torch.Tensor) – A 1D tensor containing all model gradients in a flattened form.

• model (torch.nn.Module) – The model whose parameters will be updated with the given gradients.

Returns:

The model with updated parameter gradients.

Return type:

torch.nn.Module

xuance.torch.utils.operations.assign_from_flat_params(flat_params: torch.Tensor, model: torch.nn.Module) → torch.nn.Module

Assigns parameter values from a flat tensor to the corresponding model parameters.

This function takes a flattened parameter tensor and reshapes it to match the model’s parameter shapes before assigning the values to the respective model parameters.

Parameters:

• flat_params (torch.Tensor) – A 1D tensor containing all model parameters in a flattened form.

• model (torch.nn.Module) – The model whose parameters will be updated with the given values.

Returns:

The model with updated parameters.

Return type:

torch.nn.Module

xuance.torch.utils.operations.compute_lambda_values(rewards: torch.Tensor, values: torch.Tensor, continues: torch.Tensor, lmbda: float = 0.95)

xuance.torch.utils.operations.compute_stochastic_state(logits: torch.Tensor, discrete: int = 32, sample=True) → torch.Tensor

Compute the stochastic state from the logits computed by the transition or representaiton model.

Parameters:

• logits (Tensor) – logits from either the representation model or the transition model.

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

• sample (bool) – whether or not to sample the stochastic state. Default to True.

Returns:

The sampled stochastic state.

class xuance.torch.utils.operations.dotdict(*args, **kwargs)

Bases: dict

A dictionary supporting dot notation.

as_dict() → Dict[str, Any]

xuance.torch.utils.operations.get_flat_grad(y: torch.Tensor, model: torch.nn.Module) → torch.Tensor

Computes and flattens the gradients of a given loss tensor with respect to model parameters.

Parameters:

• y (torch.Tensor) – The scalar loss tensor whose gradients will be computed.

• model (torch.nn.Module) – The neural network model whose parameters’ gradients need to be extracted.

Returns:

A 1D tensor containing all the gradients of the model parameters, concatenated into a single vector.

Return type:

torch.Tensor

xuance.torch.utils.operations.get_flat_params(model: torch.nn.Module) → torch.Tensor

Flattens and concatenates all parameters of a given model into a 1D tensor.

Parameters:

model (torch.nn.Module) – The neural network model whose parameters need to be flattened.

Returns:

A 1D tensor containing all model parameters concatenated together.

Return type:

torch.Tensor

xuance.torch.utils.operations.init_distributed_mode(master_port: str = None)

Initializes the distributed training environment.

This function sets up the necessary environment variables for distributed training, configures the CUDA device for each process, and initializes the process group for communication between GPUs using the NCCL backend.

Parameters:

master_port (str, optional) – The port number for the master process. If not provided, the default value “12355” is used.

xuance.torch.utils.operations.init_weights(m)

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

Sets the random seed for reproducibility across different libraries.

This function ensures that random number generation in PyTorch (CPU & GPU), NumPy, and Python’s built-in random module produces consistent results across runs, improving experiment reproducibility.

Parameters:

seed (int) – The seed value to set for all random number generators.

xuance.torch.utils.operations.sym_exp(x: torch.Tensor) → torch.Tensor

Computes the symmetric exponential of a tensor.

This function applies an exponential transformation while preserving the sign of the input. The operation is defined as: sign(x) * (exp(|x|) - 1).

Parameters:

x – Input tensor of any shape.

Returns:

A tensor of the same shape as input, with symmetric exponential applied element-wise.

xuance.torch.utils.operations.sym_log(x: torch.Tensor) → torch.Tensor

Computes the symmetric logarithm of a tensor.

This function applies a logarithm transformation while preserving the sign of the input. The operation is defined as: sign(x) * log(1 + |x|). Ref: https://github.com/danijar/dreamerv3/blob/8fa35f83eee1ce7e10f3dee0b766587d0a713a60/dreamerv3/jaxutils.py

Parameters:

x – Input tensor of any shape.

Returns:

A tensor of the same shape as input, with symmetric logarithm applied element-wise.

xuance.torch.utils.operations.two_hot_decoder(tensor: torch.Tensor, support_range: int) → torch.Tensor

Decode a tensor representing a two-hot vector as a tensor of floating point numbers.

Parameters:

• tensor (Tensor) – tensor to decode of shape (…, batch_size, support_size)

• support_range (int) – range of the support of the values, going from -support_range to support_range

Returns:

tensor of shape (…, batch_size, 1)

Return type:

Tensor

xuance.torch.utils.operations.two_hot_encoder(tensor: torch.Tensor, support_range: int = 300, num_buckets: int | None = None) → torch.Tensor

Encode a tensor representing a floating point number x as a tensor with all zeros except for two entries in the indexes of the two buckets closer to x in the support of the distribution. Check https://arxiv.org/pdf/2301.04104v1.pdf equation 9 for more details.

Parameters:

• tensor (Tensor) – tensor to encode of shape (…, batch_size, 1)

• support_range (int) – range of the support of the distribution, going from -support_range to support_range

• num_buckets (int) – number of buckets in the support of the distribution

Returns:

tensor of shape (…, batch_size, support_size)

Return type:

Tensor

xuance.torch.utils.operations.uniform_init_weights(given_scale)

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

Updates the learning rate using a linear decay schedule.

This function adjusts the learning rate linearly based on the current step. The learning rate starts from initial_lr and decreases linearly until it reaches a minimum value determined by end_factor * initial_lr.

Parameters:

• optimizer (torch.optim.Optimizer) – The optimizer whose learning rate will be updated.

• step (int) – The current training step.

• total_steps (int) – The total number of training steps.

• initial_lr (float) – The initial learning rate at step 0.

• end_factor (float) – The minimum learning rate factor. The final learning rate is at least end_factor * initial_lr.

value_norm

class xuance.torch.utils.value_norm.ValueNorm(*args: Any, **kwargs: Any)

Bases: Module

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)