protomotions.agents.base_agent.agent module#

Base agent implementation for reinforcement learning.

This module provides the core agent class that all RL algorithms extend. It handles the complete training lifecycle including rollout collection, experience buffering, optimization, checkpointing, evaluation, and distributed training coordination.

Key Classes:

BaseAgent: Abstract base class for all RL agents

Key Features:

Distributed training with Lightning Fabric
Experience buffer management
Automatic checkpoint saving/loading
Periodic evaluation during training
Reward normalization
Episode statistics tracking

class protomotions.agents.base_agent.agent.BaseAgent(fabric, env, config, root_dir=None)[source]#

Bases: object

Base class for reinforcement learning agents.

Provides the core training infrastructure that all algorithm implementations extend. Handles experience collection, optimization loop, checkpointing, and evaluation. Subclasses must implement model creation and algorithm-specific training logic.

Parameters:

fabric (MockFabric) – Lightning Fabric instance for distributed training.
env (BaseEnv) – Environment instance for interaction.
config (BaseAgentConfig) – Agent configuration with hyperparameters.
root_dir (Path | None) – Directory for saving checkpoints and logs (optional, uses logger dir if available).

model#: Neural network model (created by subclass).

optimizer#: Optimizer for model parameters.

experience_buffer#: Buffer for storing rollout data.

evaluator#: Evaluator for computing performance metrics.

require_reward_norm_on_load: bool = True#

__init__(fabric, env, config, root_dir=None)[source]#

Initialize the base agent.

Sets up distributed training infrastructure, initializes tracking metrics, and creates the evaluator. Subclasses should call super().__init__() first.

Parameters:

fabric (MockFabric) – Lightning Fabric for distributed training and device management.
env (BaseEnv) – Environment instance for agent-environment interaction.
config (BaseAgentConfig) – Configuration containing hyperparameters and training settings.
root_dir (Path | None) – Optional directory for saving outputs (uses logger dir if None).

property should_stop#

property has_critic: bool#

Whether this agent has a value critic.

Base agents do not assume critic ownership; PPO-style subclasses define that contract explicitly from their model config. Algorithm code should branch on this property instead of probing for critic attributes.

setup()[source]#

abstractmethod create_model()[source]#

abstractmethod create_optimizers(model)[source]#

load(checkpoint, load_env=True, load_training_state=True)[source]#

load_parameters(state_dict, load_training_state=True)[source]#

Load agent parameters from state dictionary.

Always restores model weights. When load_training_state is true, also restores optimizer-era state such as counters, reward normalization, and evaluator state.

Parameters:: state_dict – Dictionary containing saved agent state from checkpoint. Expected keys: epoch, step_count, run_start_time, best_evaluated_score, running_reward_norm (if normalization enabled), model.

get_state_dict(state_dict)[source]#

Get complete state dictionary for checkpointing.

Collects all agent state including model weights, training progress, and normalization statistics into a single dictionary for saving.

Parameters:: state_dict – Existing state dict to update (typically empty dict).
Returns:: Updated state dictionary containing all agent state.

get_inference_state_dict(state_dict, model_state_dict=None)[source]#

Get a checkpoint containing only state needed to run inference.

static inference_checkpoint_name(checkpoint_name)[source]#

save_inference_checkpoint( checkpoint_name, inference_state_dict, )[source]#

save(checkpoint_name='last.ckpt', new_high_score=False)[source]#

Save model checkpoint and environment state.

Rank 0 saves the main checkpoint (shared model weights). Each unique-task rank saves its env checkpoint. Global barrier ensures all ranks wait until saving is complete.

Parameters:

checkpoint_name (str) – Name of checkpoint file (e.g., “last.ckpt” or “epoch_100.ckpt”)
new_high_score (bool) – Whether this is a new high score (will also save as score_based.ckpt)

register_algorithm_experience_buffer_keys()[source]#

Subclasses override this to add custom keys to the experience buffer (e.g., AMP adds discriminator observations, ASE adds latent codes).

register_algorithm_experience_buffer_keys_from_obs(obs_td)[source]#

collect_rollout_step(obs_td, step)[source]#

Collect experience data during rollout at current timestep.

Called once per timestep during the data collection (rollout) phase. Subclasses implement this to: 1. Query the policy to select actions from observations 2. Store intermediate training data in experience buffer (e.g., values, log probs) 3. Return actor outputs for the environment step

Parameters:

obs_td (MockTensorDict) – TensorDict of observations from environment
step – Current timestep index in the rollout [0, num_steps)

Returns:

TensorDict with actor outputs (action, etc.) for env.step()

fit()[source]#

Main training loop for the agent.

Executes the complete training process including: 1. Experience buffer setup (auto-registers keys from model outputs) 2. Environment rollouts (data collection) 3. Model optimization 4. Periodic evaluation 5. Checkpoint saving 6. Logging and metrics

The loop runs for max_epochs epochs, where each epoch collects num_steps of experience from num_envs parallel environments, then performs multiple optimization steps on the collected data.

Note

This is the main entry point for training. Call setup() before fit().

pre_collect_step(step)[source]#

Advance agent-side state once per rollout step.

Called exactly once per step in the rollout loop, before add_agent_info_to_obs. Use this for stateful operations that must happen once per step (e.g., advancing hold counters, latent reset timers). Keep add_agent_info_to_obs as a stateless read.

add_agent_info_to_obs(obs)[source]#

Add agent-specific observations to the environment observations.

Must be stateless — reads agent state set by pre_collect_step() but does not modify it. Called for the actor’s obs during rollout. Also called during evaluation and setup (where pre_collect_step is not called), so it must not depend on pre_collect_step having run.

add_agent_info_to_next_obs(obs)[source]#

Add agent info to the next-step observations used for critic value.

Called once per rollout step on next_obs (after env.step) to prepare observations for the critic’s next_value computation. By default delegates to add_agent_info_to_obs. Override to provide different observations to the critic (e.g., fresh targets instead of held).

obs_dict_to_tensordict(obs_dict)[source]#

Convert observation dict to TensorDict.

Parameters:: obs_dict (Dict) – Dictionary of observation tensors from environment.
Returns:: TensorDict with same keys and values.
Return type:: MockTensorDict

post_env_step_modifications(dones, terminated, extras)[source]#

Allow subclasses to modify dones/terminated after env.step().

This hook allows algorithm-specific modifications (e.g., AMP discriminator termination) and then clears model-owned rollout context for any environment that is done.

Parameters:

dones – Reset flags from environment
terminated – Termination flags from environment
extras – Info dictionary from environment

Returns:

Modified (dones, terminated, extras) tuple

check_for_nans(*tensordicts)[source]#

record_rollout_step( next_obs_td, actions, rewards, dones, terminated, done_indices, extras, step, )[source]#

Record metrics and store data after environment step during rollout.

Called once per timestep during data collection phase, after the environment has been stepped. Handles: 1. Episode statistics tracking (rewards, lengths) 2. Environment extras logging 3. Experience buffer updates (actions, rewards, dones) 4. Reward normalization if enabled

Parameters:

next_obs – Observations after environment step
actions (MockTensor) – Actions that were applied
rewards (MockTensor) – Rewards from environment step
dones (MockTensor) – Reset flags indicating episode completion
terminated (MockTensor) – Termination flags indicating early termination
done_indices (MockTensor) – Indices of environments that finished episodes
extras (Dict) – Additional info dictionary from environment
step (int) – Current timestep index in the rollout [0, num_steps)

optimize_model()[source]#

abstractmethod perform_optimization_step(batch_dict, batch_idx)[source]#

post_epoch_logging(training_log_dict)[source]#

eval()[source]#

Set the model to evaluation mode.

Disables training-specific behaviors like dropout and batch normalization updates. Typically called before collecting experience or during evaluation.

train()[source]#

max_num_batches()[source]#

Calculate maximum number of minibatches per epoch.

Returns:: Integer number of minibatches needed to process all collected experience.

get_step_count_increment()[source]#

Calculate step count increment for distributed training.

Accounts for multiple GPUs/nodes and heterogeneous num_envs across ranks.

Returns:: Number of environment steps per training iteration across all processes.

terminate_early()[source]#

Request early termination of training.

Sets a flag that will cause the training loop to exit gracefully after the current epoch completes.