protomotions.agents.supervised.agent module

Supervised rollout imitation agent.

This agent collects rollouts with the student policy, labels those states with an expert policy, and optimizes a configured supervision loss. Algorithms such as MaskedMimic are experiment/model configurations of this generic loop.

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

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()

create_optimizers(model)

register_algorithm_experience_buffer_keys()

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,

)

Register algorithm keys whose shapes need a sample observation.

collect_rollout_step(obs_td, step)

Collect student action and expert label for the current state.

perform_optimization_step(batch_dict, batch_idx)

supervised_step(batch_dict)

Compute supervised imitation loss from a rollout batch.

calculate_extra_loss(batch_dict, actions)

get_state_dict(state_dict)

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.