Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

docs/source/training/rsl_rl.rst

@@ -0,0 +1,277 @@
+.. _rsl_rl:
+
+Training with RSL-RL
+====================
+
+mjlab uses `RSL-RL <https://github.com/leggedrobotics/rsl_rl>`_ for on-policy
+reinforcement learning. The integration has three parts: a **task registry**
+that bundles environment and training configs under a single name, a
+**VecEnv wrapper** that adapts mjlab environments to the interface RSL-RL
+expects, and a set of **configuration dataclasses** that control the training
+run.
+
+
+Task registry
+-------------
+
+Every task in mjlab is a pair: an environment configuration
+(``ManagerBasedRlEnvCfg``) and a training configuration
+(``RslRlOnPolicyRunnerCfg``). The task registry maps a string name to this
+pair so that training can be launched by name from the CLI.
+
+Tasks are registered by calling ``register_mjlab_task`` in the task's
+``__init__.py``:
+
+.. code-block:: python
+
+    from mjlab.tasks.registry import register_mjlab_task
+    from mjlab.tasks.velocity.rl import VelocityOnPolicyRunner
+
+    from .env_cfgs import unitree_g1_rough_env_cfg, unitree_g1_flat_env_cfg
+    from .rl_cfg import unitree_g1_ppo_runner_cfg
+
+    register_mjlab_task(
+      task_id="Mjlab-Velocity-Rough-Unitree-G1",
+      env_cfg=unitree_g1_rough_env_cfg(),
+      play_env_cfg=unitree_g1_rough_env_cfg(play=True),
+      rl_cfg=unitree_g1_ppo_runner_cfg(),
+      runner_cls=VelocityOnPolicyRunner,
+    )
+
+Each registration takes:
+
+- ``task_id``: a unique name following the convention
+  ``Mjlab-{Category}-{Terrain}-{Robot}``
+- ``env_cfg``: the ``ManagerBasedRlEnvCfg`` used for training
+- ``play_env_cfg``: a variant with randomization disabled and episode length
+  set to infinity, used for evaluation
+- ``rl_cfg``: the ``RslRlOnPolicyRunnerCfg`` with PPO hyperparameters and
+  network architecture
+- ``runner_cls``: an optional custom runner class (defaults to
+  ``MjlabOnPolicyRunner``)
+
+All task packages under ``src/mjlab/tasks/`` are auto-discovered at import
+time, so adding a new task only requires creating the config package and
+calling ``register_mjlab_task``.
+
+
+Training and playback
+---------------------
+
+**Launching a training run:**
+
+.. code-block:: bash
+
+    uv run train Mjlab-Velocity-Flat-Unitree-G1 --num-envs 4096
+
+The task name is the first positional argument. The entire configuration
+hierarchy (environment, scene, rewards, PPO hyperparameters, etc.) is
+exposed as CLI flags through `tyro <https://brentyi.github.io/tyro/>`_.
+Every field in ``ManagerBasedRlEnvCfg`` and ``RslRlOnPolicyRunnerCfg`` can
+be overridden from the command line using dot-separated paths:
+
+.. code-block:: bash
+
+    uv run train Mjlab-Velocity-Flat-Unitree-G1 \
+        --num-envs 4096 \
+        --agent.max-iterations 10000 \
+        --agent.algorithm.learning-rate 3e-4 \
+        --env.decimation 2
+
+.. important::
+
+   - **Hyphens, not underscores**: Python field names use underscores
+     (``num_envs``), but CLI flags use POSIX-style hyphens (``--num-envs``).
+   - **Explicit booleans**: boolean flags require an explicit ``True`` or
+     ``False`` value (e.g., ``--agent.resume True``, not ``--agent.resume``).
+     This is intentional for compatibility with W&B sweep configs.
+
+To discover available flags, use ``--help`` and pipe through ``grep``:
+
+.. code-block:: bash
+
+    # See all flags.
+    uv run train Mjlab-Velocity-Flat-Unitree-G1 --help
+
+    # Search for a specific field.
+    uv run train Mjlab-Velocity-Flat-Unitree-G1 --help | grep learning-rate
+
+Some commonly used top-level flags:
+
+``--num-envs``
+    Number of parallel simulation environments.
+
+``--gpu-ids``
+    GPU indices to use. Pass multiple indices for multi-GPU training (see
+    :ref:`distributed-training`), or ``None`` for CPU mode.
+
+``--video``
+    Record training rollout videos to ``{log_dir}/videos/train/``.
+
+``--enable-nan-guard``
+    Enable NaN detection and state capture (see :ref:`nan-guard`).
+
+
+**Playing back a trained policy:**
+
+.. code-block:: bash
+
+    # From W&B.
+    uv run play Mjlab-Velocity-Flat-Unitree-G1 \
+        --wandb-run-path your-entity/mjlab/run-id
+
+    # From a local checkpoint.
+    uv run play Mjlab-Velocity-Flat-Unitree-G1 \
+        --checkpoint-file logs/rsl_rl/g1_velocity/2025-01-27_14-30-00/model_1000.pt
+
+Key ``play`` arguments:
+
+``--agent``
+    Policy mode: ``"trained"`` (default), ``"zero"`` (zero actions), or
+    ``"random"`` (uniform random).
+
+``--viewer``
+    Viewer backend: ``"native"`` (MuJoCo viewer) or ``"viser"``
+    (browser-based).
+
+``--no-terminations``
+    Disable termination conditions so the policy runs indefinitely.
+
+
+VecEnv wrapper
+--------------
+
+``RslRlVecEnvWrapper`` adapts a ``ManagerBasedRlEnv`` to RSL-RL's ``VecEnv``
+interface. It handles three things:
+
+1. **Observation format**: translates observation dictionaries into the
+   ``TensorDict`` format RSL-RL expects.
+2. **Done signal**: merges ``terminated`` and ``truncated`` into a single
+   ``dones`` tensor and passes ``time_outs`` through ``extras`` so RSL-RL can
+   bootstrap correctly on truncated episodes.
+3. **Action clipping**: applies optional action clipping when ``clip_actions``
+   is set in the runner config.
+
+The wrapper also calls ``env.reset()`` during construction because RSL-RL does
+not call reset before beginning rollout collection.
+
+In normal usage you do not interact with the wrapper directly. The training
+script handles wrapping automatically.
+
+
+Configuration
+-------------
+
+``RslRlOnPolicyRunnerCfg`` is the top-level training configuration. It groups
+runner settings, network architecture (``RslRlModelCfg``), and PPO
+hyperparameters (``RslRlPpoAlgorithmCfg``). The following example from the
+Unitree G1 velocity task shows a typical configuration:
+
+.. code-block:: python
+
+    from mjlab.rl import (
+        RslRlModelCfg,
+        RslRlOnPolicyRunnerCfg,
+        RslRlPpoAlgorithmCfg,
+    )
+
+    def unitree_g1_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
+        return RslRlOnPolicyRunnerCfg(
+            actor=RslRlModelCfg(
+                hidden_dims=(512, 256, 128),
+                activation="elu",
+                obs_normalization=True,
+            ),
+            critic=RslRlModelCfg(
+                hidden_dims=(512, 256, 128),
+                activation="elu",
+                obs_normalization=True,
+            ),
+            algorithm=RslRlPpoAlgorithmCfg(
+                value_loss_coef=1.0,
+                use_clipped_value_loss=True,
+                clip_param=0.2,
+                entropy_coef=0.01,
+                num_learning_epochs=5,
+                num_mini_batches=4,
+                learning_rate=1.0e-3,
+                schedule="adaptive",
+                gamma=0.99,
+                lam=0.95,
+                desired_kl=0.01,
+                max_grad_norm=1.0,
+            ),
+            experiment_name="g1_velocity",
+            save_interval=50,
+            num_steps_per_env=24,
+            max_iterations=30_000,
+        )
+
+All fields have sensible defaults and can be overridden from the command line
+(e.g., ``--agent.algorithm.learning-rate 3e-4``). Use ``--help`` to see the
+full list of available fields and their defaults.
+
+
+Checkpoints and logging
+-----------------------
+
+Training artifacts are written to:
+
+.. code-block:: text
+
+    logs/rsl_rl/{experiment_name}/{timestamp}/
+        model_{iteration}.pt      # policy checkpoints
+        params/
+            env.yaml              # full environment config
+            agent.yaml            # full runner config
+
+Checkpoints are saved every ``save_interval`` iterations and uploaded to W&B
+as model artifacts by default. Set ``upload_model=False`` in the runner
+config to disable uploads while keeping metric logging.
+
+.. rubric:: Resuming from a checkpoint
+
+.. code-block:: bash
+
+    uv run train Mjlab-Velocity-Flat-Unitree-G1 \
+        --num-envs 4096 \
+        --agent.resume True
+
+The runner searches for the most recent run directory under
+``logs/rsl_rl/{experiment_name}/`` and loads the highest-numbered checkpoint.
+Narrow the search with ``--agent.load-run`` (regex on directory names) and
+``--agent.load-checkpoint`` (regex on checkpoint filenames).
+
+``--agent.max-iterations`` controls how many *additional* iterations to run
+from the checkpoint. If you are resuming from iteration 11500 with
+``--agent.max-iterations 300`` (the default), training will run iterations
+11500 through 11800. Set this to the number of new iterations you want.
+
+To resume from a W&B run:
+
+.. code-block:: bash
+
+    uv run train Mjlab-Velocity-Flat-Unitree-G1 \
+        --num-envs 4096 \
+        --agent.resume True \
+        --wandb-run-path your-entity/mjlab/run-id
+
+
+Citation
+--------
+
+If you use RSL-RL in your research, consider citing:
+
+.. code-block:: bibtex
+
+    @article{schwarke2025rslrl,
+        title={RSL-RL: A Learning Library for Robotics Research},
+        author={Schwarke, Clemens and Mittal, Mayank and Rudin, Nikita and Hoeller, David and Hutter, Marco},
+        journal={arXiv preprint arXiv:2509.10771},
+        year={2025}
+    }
+
+.. toctree::
+   :maxdepth: 1
+
+   motion_imitation