protomotions.agents.base_agent.model module

Base model interface for agent neural networks.

This module defines the abstract base class that all agent models must implement. It provides a TensorDictModule interface for clean, compilable models.

Key Classes:

• BaseModel: Abstract base class for all agent models (TensorDictModule)

class protomotions.agents.base_agent.model.RolloutStateSpec(shape=(), init='zeros', dtype=<Mock object>)

Bases: object

Declared model-owned per-env rollout context.

A model declares the tensor shape, dtype, and initial-value distribution for per-env context that must be carried during rollout and replayed during optimization. The framework owns when context is initialized and which rows are reset after episode boundaries.

shape: tuple[int, ...] = ()

init: str | Callable = 'zeros'

make_initial_value(num_envs, device)

__init__(shape=(), init='zeros', dtype=<Mock object>)

class protomotions.agents.base_agent.model.ProtoMotionsTensorDictModule(*args, **kwargs)

Bases: TensorDictModuleBase

Common contract for ProtoMotions-owned TensorDict modules.

Use this for modules that read and write named TensorDict keys and may own rollout state that must be saved in the experience buffer. Examples include policy/critic modules, normalizing observation processors, container modules, and models with rollout-local latent state.

forward(tensordict, log_internals=False)

Forward pass for ProtoMotions TensorDict modules.

log_internals is part of the common contract. Modules that produce diagnostics can honor it; stateless modules can ignore it.

rollout_state_specs()

Local model-owned per-env rollout state declarations.

reset_rollout_context(

env_ids=None,

num_envs=None,

device=None,

)

Initialize or reseed model-owned context carried across rollout steps.

num_envs and device are supplied once during setup, after the simulator is built. Later calls pass only env_ids to reseed rows for environments that just ended. This method intentionally does not infer shape/device changes mid-run; rebuilding the simulator should call setup again and reinitialize explicitly.

init_rollout_state(num_envs, device)

Allocate declared rollout-context buffers for this module tree.

reset_rollout_state(env_ids=None)

Reseed selected rows for this module tree’s rollout context.

read_rollout_state(tensordict)

Inject carried rollout context into tensordict when absent.

Rollout forwards read from model-owned buffers; optimization forwards read replayed values already present in the batch. This helper makes the distinction explicit and avoids each model hand-writing buffer fallbacks.

rollout_context_keys()

Module-owned state that must be stored from rollout and replayed.

experience_buffer_keys()

Keys produced during rollout that the experience buffer must store.

compute_model_loss(

tensordict,

current_epoch,

zero_loss,

log_prefix='model',

)

Optional module-owned auxiliary loss for agent optimization loops.

Most modules do not own an auxiliary loss. Models that do, such as a VAE-backed policy head, override this and return (loss, log_dict).

class protomotions.agents.base_agent.model.BaseModel(*args, **kwargs)

Bases: ProtoMotionsTensorDictModule

Base class for all agent models.

All models are TensorDictModules with a single forward method that processes observations and returns all model outputs in a TensorDict.

Parameters:

config (BaseModelConfig) – Model configuration with architecture parameters.

config

Stored configuration for the model.

in_keys

Input keys for TensorDict (set by subclasses).

out_keys

Output keys for TensorDict (default: [“action”]).

__init__(config)

optimization_module()

Module whose parameters should be optimized by the owning agent.

materialize_from_state_dict(state_dict)

Create lazily-owned modules needed for strict state-dict loading.

materialize(tensordict)

Run the setup-time pass used to create lazy module parameters.

abstractmethod forward(tensordict, log_internals=False)

Forward pass through the model.

Parameters:

tensordict (MockTensorDict) – TensorDict containing observations.

Returns:

TensorDict with model outputs added.

Return type:

MockTensorDict