$ show Axellwppr/motion_tracking
This repository implements a complete reinforcement learning pipeline for whole-body motion tracking on humanoid robots, specifically Unitree G1 and XBot L7. The system uses a teacher-student architecture with Proximal Policy Optimization (PPO) to train policies that track reference motions from retargeted motion capture datasets. It features three training phases: a base training phase, an adaptation phase with domain randomization, and a finetuning phase. The environment is built on MJLab (a MuJoCo wrapper) and includes extensive domain randomization (e.g., motor parameters, body mass, center of mass, joint offsets, external forces, and a window load capacity curriculum for payload-aware training). The codebase is highly modular, with registered MDP components (observations, rewards, terminations, randomizations) and a flexible configuration system using Hydra. It also includes advanced features such as symmetry augmentation, differentiable forward kinematics, and a custom motion dataset pipeline with memory-mapped storage. While primarily a reinforcement learning training framework, it has strong data engineering aspects for processing motion datasets and generating training data for the student encoder. The repository is well-structured for research and sim-to-real transfer, but it is tightly coupled to specific robot assets and dependencies (torch, torchrl, mjlab) with pinned versions that may require adaptation for other environments.
The LLM identified 4 files worth highlighting for reuse:
A differentiable forward kinematics module implemented as a PyTorch layer. It computes body positions, velocities, and quaternions from joint positions for an arbitrary kinematic chain defined in a robot description (URDF/MJCF). This is a highly reusable component for any robotics project that requires kinematic computation within a neural network, such as RL or imitation learning. It is self-contained, depending only on torch and tensordict, and can be extracted with minimal changes.
A symmetry transformation module that automatically generates joint and spatial symmetry mappings from a robot's XML description (MJCF or URDF). It provides a SymmetryTransform class that can mirror observations and actions, and a function to auto-detect symmetric joint pairs by analyzing the robot's kinematic tree. This is a non-obvious, reusable pattern for reducing policy learning complexity in symmetric robots, applicable to any humanoid or symmetric manipulator.
A SymmetricVecNorm class that extends VecNorm (from torchrl) by doubling the batch with mirrored observations during running statistics updates. This reduces bias in normalization when the environment has symmetries. It is a clever, reusable component for any RL training pipeline that exploits symmetry, and can be easily extracted into a standalone module.
A collection of math utilities for robotics and reinforcement learning, including quaternion multiplication, application, axis-angle conversion, finite-difference velocity computation, signal smoothing, and timestamp interpolation. These functions are self-contained, well-documented, and can be directly reused in any robotics project for state estimation or control.