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.

forward_deployable(tensordict)

The deployable forward: by default the same graph the policy runs.

A policy family whose exported pass differs from its training forward (an actor that samples exploration noise) overrides this; everything else inherits it, so the exporter, evaluator and teacher call one name without an isinstance check on the policy class.

property deployable_in_keys: list

The input keys the deployable graph reads; in_keys by default.

property deployable_action_key: str

The output key the host drives the robot from.

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.

save_rollout_context()

Copy every rollout-context buffer in this module tree.

A caller that runs its own episodes on the training envs, such as evaluation, saves this first and passes it to load_rollout_context afterwards, so training resumes with the context it had.

load_rollout_context(state)

Write back rollout context returned by save_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.

build_training_targets(rollout)

Derive model-owned targets from [time, env, …] rollout tensors.

BaseAgent calls this once per rollout under no_grad, after agent preprocessing and before flattening/shuffling destroys time adjacency. Inputs are borrowed experience-buffer views: do not mutate them. Return additional [time, env, …] targets/masks for compute_model_loss; the agent owns batching, not their meaning. The default gathers child targets; overrides should call super() to preserve those results.

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.

Variables:

• 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.

optional_checkpoint_modules()

Networks allowed to keep initialization when wholly absent on load.

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