Language Language English 简体中文

unilab.algos.torch.common

class unilab.algos.torch.common.EmpiricalNormalization

Bases: Module

Normalize mean and variance of observations using running statistics.

count: torch.Tensor

__init__(shape, device, eps=0.01)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

property mean

property std

forward(x, center=True, update=True)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• x (Tensor)

• center (bool)

• update (bool)

Return type:

Tensor

update(x)

inverse(y)

class unilab.algos.torch.common.DistributionalQNetwork

Bases: Module

Single distributional Q-network (C51 variant).

Architecture: Linear→ReLU → Linear→ReLU → Linear→ReLU → Linear Outputs num_atoms logits over the value distribution.

Parameters:

• obs_dim (int)

• n_act (int)

• num_atoms (int)

• v_min (float)

• v_max (float)

• hidden_dim (int)

• device (Optional[device])

__init__(obs_dim, n_act, num_atoms, v_min, v_max, hidden_dim, device=None)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• obs_dim (int)

• n_act (int)

• num_atoms (int)

• v_min (float)

• v_max (float)

• hidden_dim (int)

• device (Optional[device])

forward(obs, actions)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• obs (Tensor)

• actions (Tensor)

Return type:

Tensor

projection(obs, actions, rewards, bootstrap, discount, q_support, device)

Categorical projection (Bellman update on the distribution support).

Parameters:

• obs (Tensor)

• actions (Tensor)

• rewards (Tensor)

• bootstrap (Tensor)

• discount (Tensor)

• q_support (Tensor)

• device (device)

Return type:

Tensor

class unilab.algos.torch.common.Critic

Bases: Module

Twin distributional Q-networks for off-policy RL (SAC/TD3).

Parameters:

• obs_dim (int)

• n_act (int)

• num_atoms (int)

• v_min (float)

• v_max (float)

• hidden_dim (int)

• device (Optional[device])

q_support: torch.Tensor

__init__(obs_dim, n_act, num_atoms, v_min, v_max, hidden_dim, device=None)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• obs_dim (int)

• n_act (int)

• num_atoms (int)

• v_min (float)

• v_max (float)

• hidden_dim (int)

• device (Optional[device])

forward(obs, actions)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• obs (Tensor)

• actions (Tensor)

projection(obs, actions, rewards, bootstrap, discount)

Projection operation using both Q-networks.

Parameters:

• obs (Tensor)

• actions (Tensor)

• rewards (Tensor)

• bootstrap (Tensor)

• discount (Tensor)

get_value(probs)

Calculate value from probability distribution using support.

Parameters:

probs (Tensor)

Return type:

Tensor

unilab.algos.torch.common.get_env_dims(env_name, sim_backend='mujoco', env_cfg_override=None)

Get (actor_obs_dim, action_dim, critic_obs_dim) from environment.

Parameters:

• env_name (str)

• sim_backend (str)

• env_cfg_override (dict | None)

Return type:

tuple[int, int, int]

unilab.algos.torch.common.check_nan_loss(loss, default_metrics)

Check if loss contains NaN or Inf values.

Parameters:

• loss (Tensor) – Loss tensor to check

• default_metrics (dict) – Default metric values to return if NaN detected

Return type:

tuple[Tensor | None, dict | None]

Returns:

(loss, None) if valid, (None, nan_metrics) if invalid

unilab.algos.torch.common.clip_gradients(parameters, max_norm=10.0)

Clip gradients by global norm.

Parameters:

• parameters – Model parameters

• max_norm (float) – Maximum gradient norm

unilab.algos.torch.common.safe_tensor(tensor, nan_value=0.0, clamp_range=(-10.0, 10.0))

Make tensor numerically safe by clamping and replacing NaN values.

Parameters:

• tensor (Tensor) – Input tensor

• nan_value (float) – Value to replace NaN with

• clamp_range (tuple) – (min, max) range to clamp values

Return type:

Tensor

Returns:

Safe tensor

unilab.algos.torch.common.ensure_registries(packages=None, *, optional_packages=None, fail_on_error=True)

Import env registry bootstrap modules.

Parameters:

• packages (Sequence[str] | None)

• optional_packages (Sequence[str] | None)

• fail_on_error (bool)

Return type:

None

unilab.algos.torch.common.build_actor(algo_type, obs_dim, action_dim, actor_hidden_dim, use_layer_norm, device, num_envs=1, actor_num_blocks=2, actor_noise_zeta_mu=2.0, actor_noise_zeta_max=16, priv_info_dim=None, priv_info_embed_dim=9, priv_mlp_hidden_dims=(256, 128, 9), **kwargs)

Build the correct actor model based on algorithm type.

Parameters:

• actor_num_blocks (int)

• actor_noise_zeta_mu (float)

• actor_noise_zeta_max (int)

• priv_info_dim (int | None)

• priv_info_embed_dim (int)

• priv_mlp_hidden_dims (tuple[int, ...] | list[int])

Modules

actor_factory	Actor factory helpers for torch off-policy algorithms.
ane_inference	Simplified ANE backend using deterministic inference.
ane_wrapper	ANE (Apple Neural Engine) inference wrapper for Fast SAC.
base_collector	Base collector class for shared functionality.
device
networks	Neural network architectures for RL algorithms.
normalization	Observation normalization for RL training.
stability	Numerical stability utilities for RL training.