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.BaseModel(*args, **kwargs)[source]¶

Bases: TensorDictModuleBase

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)[source]¶

abstractmethod forward(tensordict)[source]¶

Forward pass through the model.

Parameters:: tensordict (MockTensorDict) – TensorDict containing observations.
Returns:: TensorDict with model outputs added.
Return type:: MockTensorDict