Predictive Style Matching (PSM)

Code for Predictive Style Matching: Natural and Robust Humanoid Locomotion on Unitree G1 (first-draft implementation, work in progress).

Dependencies

RL training and simulation build on mjlab. This repo registers the Psm-G1 task via mjlab’s task entry point. See the mjlab docs for install notes, multi-GPU training, and play/export.

Hardware deployment is described for unitree_rl_mjlab (also mjlab-based).

Install

git clone https://github.com/simeon-ned/predictive-style-matching.git
cd predictive-style-matching
uv sync

Quick start

Predictor (offline)

# Prepare motion data (once per dataset; see data/README.md)
uv run psm-data-to-npz --input-path data/seed --output-dir data/motions   # CSV/PKL → extended NPZ
uv run psm-data-augment --input-path data/motions     # or upgrade existing NPZ
uv run psm-data-stack --dataset-path data/motions # → data/motions/motions.npz

uv run psm-predictor-train
uv run psm-predictor-play --npz data/motions/your_clip.npz
uv run psm-data-vis --npz data/motions/your_clip.npz   # motion playback only (Viser)

Logs go to logs/predictor/<timestamp>/ (predictor.pth, metadata.pkl, config.yaml).

RL (mjlab)

uv run psm-env-train Psm-G1
uv run psm-env-play Psm-G1
uv run psm-list-envs   # optional: list registered tasks

By default, RL uses the latest bundle under logs/predictor/, then falls back to src/psm/predictor/weights/ if none exist. Override:

uv run psm-env-train Psm-G1 --predictor-path /path/to/bundle
uv run psm-env-train Psm-G1 --predictor-bundled
# or: --env.commands.twist.predictor-path /path/to/bundle

Policy logs: logs/rsl_rl/g1_psm/ (each run snapshots the active predictor under params/predictor/). ONNX: <run>/params/latest.onnx on play.

Deploy: unitree_rl_mjlab on the robot; copy the exported policy from your play run.

Repository layout

├── README.md           # this file — code & usage
├── pyproject.toml
├── data/
├── src/psm/
│   ├── predictor/      # training; logs → logs/predictor/; weights/ = packaged fallback
│   ├── motion/         # psm-data-to-npz, psm-data-augment, psm-data-stack
│   ├── scripts/        # psm-env-train, psm-env-play, utilities
│   └── env/            # RL (Psm-G1): cfg/, mdp/, runner.py, utils/ (deploy, symmetry, predictor path/log)
└── docs/               # GitHub Pages project site (not Python)
    ├── index.html
    └── static/

Paper TeX may live in a separate private repo during review; figures for the site are under docs/static/images/.

Citation

If you use this work, please cite:

@misc{nedelchev2026psm,
  title         = {Predictive Style Matching: Natural and Robust Humanoid Locomotion},
  author        = {Simeon Nedelchev and Ekaterina Chaikovskaia and Egor Davydenko and Eduard Zaliaev and Roman Gorbachev},
  year          = {2026},
  eprint        = {2606.07083},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url           = {https://arxiv.org/abs/2606.07083},
}

Acknowledgments

• mjlab — RL environment and training stack
• MuJoCo — physics simulation
• unitree_rl_mjlab — G1 sim2real deploy reference