protomotions.agents.common.common module

Common neural network components and utilities for agents.

This module provides shared building blocks used across different agent architectures, including observation normalization, weight initialization, and specialized layers.

Key Classes:

• NormObsBase: Base class for modules with observation normalization

• ObsProcessor: General observation processor for reshaping and normalizing

• Embedding: Embedding layer for discrete inputs

Key Functions:

• weight_init: Initialize network weights

• get_params: Extract parameters from optimizer groups

protomotions.agents.common.common.get_params(obj)

Extract parameters from optimizer parameter groups.

Handles both flat lists of parameters and grouped parameters (as used by optimizers).

Parameters:

obj – Either a list of nn.Parameter or a list of parameter groups (dicts).

Returns:

Flat list of all parameters.

Return type:

List[<Mock object at 0x7fbf798dabd0>[]]

protomotions.agents.common.common.weight_init(m, orthogonal=False)

Initialize weights for neural network modules.

Applies appropriate initialization to linear layers and other modules. Linear layers get orthogonal or default initialization with zero bias.

Parameters:

• m – Neural network module to initialize.

• orthogonal – If True, use orthogonal initialization for linear layers.

class protomotions.agents.common.common.NormObsBase(config)

Bases: <Mock object at 0x7fbf798db690>[]

Base class for modules with observation normalization.

Provides running mean/std normalization of observations using exponential moving averages. Normalization statistics are updated during training and frozen during evaluation.

Uses lazy initialization - input shape is inferred on first forward pass.

This is a simple tensor-to-tensor module. Subclasses handle TensorDict extraction/insertion.

Parameters:

config (NormObsBaseConfig) – Configuration specifying output dimensions and normalization parameters.

running_obs_norm

RunningMeanStd module for observation normalization (lazy).

num_out

Output dimension.

__init__(config)

build_norm()

forward(obs)

Forward pass that normalizes observations.

Parameters:

obs (MockTensor) – Observation tensor to normalize.

Returns:

Normalized observation tensor.

Return type:

MockTensor

protomotions.agents.common.common.apply_module_operations(

obs,

module_operations,

normalizer,

forward_model=None,

)

class protomotions.agents.common.common.ObsProcessor(*args, **kwargs)

Bases: TensorDictModuleBase

General observation processor - applies operations and normalization.

ForwardConfig applies normalization but skips the forward model (no MLP).

__init__(config)

config: ObsProcessorConfig

forward(tensordict, *args, **kwargs)

class protomotions.agents.common.common.ModuleContainer(*args, **kwargs)

Bases: TensorDictModuleBase

Generic container that runs a list of modules sequentially on a TensorDict.

With TensorDict, the distinction between “sequential”, “parallel input”, and “parallel output” is implicit in how keys flow between modules.

__init__(config)

config: ModuleContainerConfig

forward(tensordict, *args, **kwargs)

Forward through all models sequentially.