protomotions.agents.supervised.agent module#

Supervised rollout imitation agent.

This agent collects simulator rollouts and optimizes a configured supervision loss. Labels can come from observations, model-owned target encoders, or an optional frozen external expert. It is not an offline dataset training loop.

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

Bases: BaseAgent

Student/expert rollout agent for supervised distillation.

The agent collects TensorDict rollouts, writes configured model outputs into the experience buffer, and optimizes SupervisionLossConfig. Models and experiment files define which keys are predictions and labels.

model: BaseModel#
create_model()[source]#
create_optimizers(model)[source]#
register_algorithm_experience_buffer_keys()[source]#

Register algorithm-specific keys in the experience buffer.

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]#

Register algorithm keys whose shapes need a sample observation.

collect_rollout_step(obs_td, step)[source]#

Collect the executed action, model outputs and labels for this state.

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

perform_optimization_step(batch_dict, batch_idx)[source]#
supervised_step(batch_dict)[source]#

Compute supervised imitation loss from a rollout batch.

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.