srb agent train — Train Agent

Train a reinforcement-learning agent in a registered SRB environment with an integrated framework (dreamer, tdmpc2, sb3_*, sbx_*, rsl_rl_ppo, skrl_*).

Usage

srb agent train --algo ALGO --env ENV_ID [options]

Options | Shared Agent Options

--continue_training and --model are mutually exclusive. Use --continue_training for the latest run or --model PATH for a specific checkpoint. Resume behavior depends on the framework; TD-MPC2 supports weights import rather than full training-state resume.

Framework Resume Details

For sb3_* and sbx_*, agent.n_timesteps is the total step target, including steps already recorded by the checkpoint. Resume trains only the remaining steps and keeps the original counter. If the checkpoint already meets the target, --continue_training does nothing; an explicit --model at that budget is an error. Raise agent.n_timesteps above the checkpoint’s step count to extend training.

Off-policy SB3/SBX algorithms restore a saved replay buffer when one is available. Without one they resume with an empty buffer and repeat the learning_starts warm-up; enable agent.save_replay_buffer to preserve replay state. When VecNormalize is enabled, each resumable checkpoint uses its own matching normalization statistics. A checkpoint missing its required statistics cannot resume with normalization enabled.

SKRL resumes at the absolute trainer timestep stored in checkpoints and treats agent.trainer.timesteps as the total target. Older checkpoints without a stored cursor need +agent.resume_timestep=N, where N is the saved trainer timestep; checkpoints saved with agent.experiment.store_separately=true do not include SRB’s cursor. Standard SKRL checkpoints do not preserve general off-policy replay memory. The supported demonstration-prefill path saves its native replay state, but resuming it requires the same demonstration options and source, as described below.

Choices include sbx_ppo and sb3_ppo; availability depends on the installed framework profile.

SKRL Demonstration Prefill

skrl_sac, skrl_ddpg, and skrl_td3 accept canonical SRB JSONL demonstrations through --demo-dataset and --demo-max-transitions. The demonstration observation and action layout must match the live environment; transition limits and replay capacity are checked before training.

srb agent train --algo skrl_sac --env peg_in_hole \
  --demo-dataset /data/peg_in_hole_canonical \
  --demo-max-transitions 1000 \
  env.num_envs=1 agent.memory.memory_size=8192 --hide_ui

To resume a demonstration-prefilled run, supply the same demonstration options and source. Evaluation of the resulting policy does not require the source.

Other algorithms do not support this prefill path. Use srb dataset compatibility DATASET --purpose train to check a dataset’s training compatibility. Robomimic offline learning remains available through srb agent learn --dataset.

Dreamer Checkpoints and Continuation

Dreamer can resume training only when the run saved replay state with both agent.replay.save=true and agent.replay.save_wait=true. A weights-only checkpoint cannot restore missing replay; use it to start a new run instead.

Dreamer Demonstration Prefill

Dreamer also accepts the paired demonstration options and canonical SRB JSONL format. The observation projection, action layout, and runtime scales must match the live learner; actions use continuous [-1, 1] coordinates. Configure replay persistence with replay.save: true and replay.save_wait: true for resumable demonstration-prefilled runs. TD-MPC2 does not support demonstration ingestion.

Examples

# Fresh run on the landing environment with 512 parallel envs
srb agent train --algo sbx_ppo --env landing env.num_envs=512 --hide_ui

# Continue the latest run on the same env/algo combination
srb agent train --algo sbx_ppo --env landing --continue_training

# Resume from an arbitrary checkpoint
srb agent train --algo sbx_ppo --env landing \
  --model logs/landing/sbx_ppo/<run>/ckpt/<ckpt>

The path is resolved relative to the current working directory. The default log root is <repo>/logs; pass --logdir to select another root.

Model Artifacts

Pass --model-artifact-output OUT to package the run as a portable SRB model artifact. Add --model-artifact-dataset DATASET to record training-data provenance, or use srb model from-run after training for the same packaging workflow.

For demonstration-prefilled runs, packaging retains training-time manifest and member digests even if the source moves or is removed. An explicit --model-artifact-dataset sets the declared dataset link; the consumed demonstration binding remains in artifact metadata.

See the Reinforcement Learning Workflow for an end-to-end walkthrough including hyper-parameter overrides via Hydra and Weights & Biases logging.