$ show unitreerobotics/unitree_rl_mjlab
This repository provides an end-to-end reinforcement learning framework for Unitree legged robots (Go2, G1, H1, H2, A2, AS2, R1). It uses MuJoCo physics (via the custom `mjlab` package) to simulate locomotion and motion imitation tasks. The codebase includes training scripts (`scripts/train.py`) that run PPO (from `rsl_rl`) with configurable network architectures and multi-GPU support, exporting trained policies to ONNX format. The `deploy/` directory contains C++ applications that execute these policies on real hardware using ONNX Runtime and Unitree SDK2, orchestrated by a finite state machine. The system is split into two main tasks: velocity tracking (gait generation) and motion tracking (whole-body imitation from reference trajectories). Key architectural patterns include manager-based environment configuration, composable MDP terms (rewards, observations, terminations), curriculum learning, and a YAML-driven FSM for deployment. The repository is actively maintained, with recent commits and bundled third-party dependencies (MuJoCo binaries, ONNX Runtime) for cross-platform deployment.
The LLM identified 4 files worth highlighting for reuse:
Implements adaptive motion replay sampling for whole-body imitation learning. Stores failure histories and biases motion segment selection based on past performance, improving sample efficiency and robustness. This algorithm is dataset-agnostic and could be reused in any motion tracking RL pipeline that samples from a library of reference trajectories.
A recursive descent parser that compiles a human-readable domain-specific language for defining joystick button combinations (e.g., 'LT(2s)+A.on_pressed') into predicate functions. This pattern is non-obvious and could be extracted as a lightweight C++ DSL for event-based control in any embedded or robotics application.
A comprehensive set of reward functions for quadruped locomotion, covering base height, foot air time, slip, soft landing, gait symmetry, and energy efficiency. Each reward is defined as a separate function with clear hyperparameters and is registered via the mjlab manager pattern. These reward functions are reusable in other legged locomotion RL environments with similar state/action spaces.
A generic finite state machine for robot control that dynamically registers states and transitions. States are derived from `FSMState` and registered via registers the state machine. Transition conditions are evaluated from a joystick DSL at runtime. This pattern is reusable for any robot control application requiring mode switching (e.g., passive, standby, locomotion, recovery) with hardware abstraction.