protomotions.agents.supervised.config module#

Configuration for generic supervised rollout training.

class protomotions.agents.supervised.config.RolloutActor(value)[source]#

Bases: Enum

Which model forward produces the rollout.

STUDENT uses the model’s forward. EXPERT uses the model’s own collect_expert_rollout. An optional external expert labels every state and drives a fraction of episodes through expert_rollout_probability.

STUDENT = 'student'#

EXPERT = 'expert'#

class protomotions.agents.supervised.config.SupervisedAgentConfig( batch_size, training_max_steps, training_max_iterations=None, _target_='protomotions.agents.supervised.agent.SupervisedAgent', model=<factory>, num_steps=32, gradient_clip_val=0.0, fail_on_bad_grads=False, check_grad_mag=True, gamma=0.99, bounds_loss_coef=0.0, task_reward_w=1.0, num_mini_epochs=1, training_early_termination=None, save_epoch_checkpoint_every=1000, save_last_checkpoint_every=10, save_inference_checkpoint=False, max_episode_length_manager=None, evaluator=<factory>, normalize_rewards=True, normalized_reward_clamp_value=5.0, reward_norm_ema_decay=None, expert_model_path=None, rollout_actor=RolloutActor.STUDENT, loss=<factory>, expert_rollout_probability=0.0, )[source]#

Bases: BaseAgentConfig

Generic supervised imitation agent configuration.

Experiment files choose the rollout actor, optional external expert checkpoint, and supervised loss keys. The agent loop stays independent of the specific student model.

Attributes:: batch_size: Training batch size. training_max_steps: Maximum training steps. training_max_iterations: Maximum number of complete rollout and optimization iterations. When set, this takes precedence over training_max_steps. model: Model configuration. num_steps: Environment steps per update. gradient_clip_val: Max gradient norm. 0=disabled. fail_on_bad_grads: Fail on NaN/Inf gradients. check_grad_mag: Log gradient magnitude. gamma: Discount factor. bounds_loss_coef: Action bounds loss. 0 for tanh outputs. task_reward_w: Task reward weight. num_mini_epochs: Mini-epochs per update. training_early_termination: Stop early at this step. None=disabled. save_epoch_checkpoint_every: Save epoch_xxx.ckpt every N epochs. save_last_checkpoint_every: Save last.ckpt every K epochs. save_inference_checkpoint: Also save inference_<name>.ckpt without optimizer or other training-only state. max_episode_length_manager: Episode length curriculum. evaluator: Evaluation config. normalize_rewards: Normalize rewards. normalized_reward_clamp_value: Clamp normalized rewards to [-val, val]. reward_norm_ema_decay: EMA decay for reward normalization (None = Welford). Set to e.g. 0.99 to track non-stationary reward distributions. expert_model_path: External expert checkpoint that labels rollouts; requires rollout_actor=student. rollout_actor: Model forward used for collecting rollout actions. loss: Supervised loss over model outputs and labels. expert_rollout_probability: Per-episode probability that the external expert drives the episode.

model: BaseModelConfig#

expert_model_path: str | None = None#

rollout_actor: RolloutActor = 'student'#

loss: SupervisionLossConfig#

expert_rollout_probability: float = 0.0#

apply_inference_defaults(env)[source]#

Run the trained student alone: drop the expert and its observations.

The expert labels rollouts during training; at inference the student acts on its own, so the expert checkpoint, its rollout share and the expert-prefixed observations it added are removed.

__init__( batch_size, training_max_steps, training_max_iterations=None, _target_='protomotions.agents.supervised.agent.SupervisedAgent', model=<factory>, num_steps=32, gradient_clip_val=0.0, fail_on_bad_grads=False, check_grad_mag=True, gamma=0.99, bounds_loss_coef=0.0, task_reward_w=1.0, num_mini_epochs=1, training_early_termination=None, save_epoch_checkpoint_every=1000, save_last_checkpoint_every=10, save_inference_checkpoint=False, max_episode_length_manager=None, evaluator=<factory>, normalize_rewards=True, normalized_reward_clamp_value=5.0, reward_norm_ema_decay=None, expert_model_path=None, rollout_actor=RolloutActor.STUDENT, loss=<factory>, expert_rollout_probability=0.0, )#