Skip to main content

Mimic Policy Training

Train, preview, and export a K1 motion-imitation policy with Cyclo MJLab. Run the commands inside the Cyclo MJLab container from /workspace/cyclo_mjlab.

Train

python scripts/reinforcement_learning/train.py \
Cyclo-Mimic-K1-Rev1-Dance1 \
--env.scene.num-envs 4096

Training output is written to:

logs/rsl_rl/k1_mimic/<run>/

<run> is the timestamped directory created when training starts. The main files in each run directory are:

The training configuration saves a checkpoint every 500 iterations. Because Dance1 and Dance2 use the same k1_mimic output root, select the run directory created by the task you trained.

Play and Export

Select a checkpoint from the matching training run and pass its path to play.py.

python scripts/reinforcement_learning/play.py \
Cyclo-Mimic-K1-Rev1-Dance1 \
--checkpoint-file logs/rsl_rl/k1_mimic/<run>/model_<iteration>.pt \
--num-envs 1 \
--viewer native

Use the native viewer when a desktop display is available.

Replace <run> with the timestamped run directory and <iteration> with the checkpoint iteration to use.

If --viewer is omitted, the default auto mode uses the native MuJoCo viewer when a desktop display is available and otherwise uses the Viser viewer. Before opening the selected viewer, play.py exports the policy and its reference motion data to the same training run directory.

Display the Reference Motion Next to the Policy

Add --reference-ghost-offset to display the reference motion as a semi-transparent ghost next to the robot controlled by the trained policy:

python scripts/reinforcement_learning/play.py \
Cyclo-Mimic-K1-Rev1-Dance1 \
--checkpoint-file logs/rsl_rl/k1_mimic/<run>/model_<iteration>.pt \
--num-envs 1 \
--reference-ghost-offset 1.4,0,0

The three values are the reference ghost offset in the world-frame x, y, and z directions, in meters. In this example, 1.4,0,0 moves the reference robot 1.4 m in the positive X direction so that the policy and reference motions can be compared side by side.

Use the same option with the Dance2 task ID to compare a Dance2 policy with its reference motion.

After play.py exports the policy, the run directory contains:

How to Add Your Own Mimic Task

Use this flow when you want to add a new Mimic task for your own motion.

note

This guide assumes you already have a Soma-retargeter output CSV from the Kimodo > Soma-retargeter pipeline.

The examples below use my_motion as the new motion name. Use your own motion name in the same places.

  1. Place and convert the motion file.

    Place the Soma-retargeter output CSV under the K1 Rev.1 motion data directory:

    my_motion_soma.csv is the headered CSV exported by Soma-retargeter. Cyclo MJLab does not train directly from this file, so convert it to the MJLab motion format:

    python scripts/tools/motion/soma_retargeter_csv_converter.py \
    -f source/assets/motions/K1_rev1/my_motion/my_motion_soma.csv \
    --output_name source/assets/motions/K1_rev1/my_motion/my_motion.npz

    Validate the generated NPZ before training:

    python scripts/tools/motion/replay_npz.py \
    -f source/assets/motions/K1_rev1/my_motion/my_motion.npz \
    --validate-only

    Preview the motion and confirm that the robot pose, playback speed, and loop transition are correct:

    python scripts/tools/motion/replay_npz.py \
    -f source/assets/motions/K1_rev1/my_motion/my_motion.npz \
    --loop
  2. Create a new Mimic task config.

    Create source/tasks/mimic/config/k1_rev1/my_motion_env_cfg.py by duplicating dance1_env_cfg.py. In the new file, point TRAJECTORY_FILE to the generated NPZ:

    TRAJECTORY_FILE = (
    SRC_PATH / "assets" / "motions" / "K1_rev1" / "my_motion" / "my_motion.npz"
    )

    Rename k1_rev1_dance1_env_cfg to k1_rev1_my_motion_env_cfg and keep the remaining configuration logic unchanged.

  3. Register the new task.

    The files you edit are here:

    In source/tasks/mimic/config/k1_rev1/__init__.py, import the new configuration:

    from .my_motion_env_cfg import k1_rev1_my_motion_env_cfg

    Then register a unique task ID:

    register_mjlab_task(
    task_id="Cyclo-Mimic-K1-Rev1-MyMotion",
    env_cfg=k1_rev1_my_motion_env_cfg(),
    play_env_cfg=k1_rev1_my_motion_env_cfg(play=True),
    rl_cfg=k1_rev1_mimic_ppo_runner_cfg(),
    runner_cls=MotionTrackingOnPolicyRunner,
    )

    Check that the task is registered:

    python scripts/list_envs.py --keyword MyMotion

    Check that the output includes Cyclo-Mimic-K1-Rev1-MyMotion.

  4. Train and export.

    Run train.py with the new task ID. Then run play.py with the same task ID to preview the policy and export exported/policy.onnx. Keep the generated my_motion.csv with the policy for Sim2Real deployment.