Alopezcordero's picture
Update README.md
bab4c1d verified
|
Raw
History Blame Contribute Delete
6.28 kB
---
library_name: stable-baselines3
pipeline_tag: reinforcement-learning
tags:
- reinforcement-learning
- deep-reinforcement-learning
- stable-baselines3
- ppo
- mujoco
- humanoid
- adversarial-motion-priors
- amp
- character-animation
- motion-capture
model-index:
- name: ppo-amp-humanoid-direction
results:
- task:
type: reinforcement-learning
name: reinforcement-learning
dataset:
name: HumanoidDirection-v0
type: HumanoidDirection-v0
metrics:
- type: mean_reward
value: "971.74 ± 13.51"
name: mean_reward
---
# PPO + AMP — Direction-Following MuJoCo Humanoid
A **PPO** policy (Stable-Baselines3) trained with **Adversarial Motion Priors (AMP)** ([Peng et al., 2021](https://arxiv.org/abs/2104.02180)) on a custom `HumanoidDirection-v0` environment — Gymnasium's MuJoCo `Humanoid-v5` extended with a random target heading each episode.
## Model description
MuJoCo `Humanoid-v5` that walks in a random commanded direction with a natural, human-like gait, using **Adversarial Motion Priors (AMP)** ([Peng et al., 2021](https://arxiv.org/abs/2104.02180)) + **Stable-Baselines3 PPO**.
Each episode the humanoid gets a random 2D target heading. Reward is `0.5 · task + 0.5 · style`: the task reward (bounded in [0, 1]) pays full credit for walking at ~1.4 m/s along the target, and the style reward comes from an LSGAN discriminator trained to tell policy transitions from retargeted mocap transitions. See `replay.mp4` for a sample rollout.
Eval (10 episodes, deterministic): **971.74 ± 13.51** episode reward.
Three artifacts are trained jointly and shipped together:
| File | What it is |
|---|---|
| `ppo_humanoid_direction_amp_fixed.zip` | SB3 PPO policy (MLP, π and V nets [1024, 512], ReLU) |
| `vecnormalize_amp_fixed.pkl` | `VecNormalize` observation statistics — **required at inference** |
| `amp_discriminator_fixed.pt` | AMP discriminator state dict (2×512 MLP, 90-D input) — training artifact, not needed for inference |
- **Observation space**: `Humanoid-v5` observation + 2-D target direction (unit vector).
- **Action space**: `Humanoid-v5` continuous torques (17-D).
- **Task reward** ∈ [0, 1]: `exp(-2 · max(0, 1.4 − v_proj)²)` — full credit at walking speed along the target, no bonus for sprinting.
- **Style reward** ∈ [0, 1]: `clamp(1 − 0.25(d − 1)², 0, 1)` from the discriminator, computed on heading-invariant 45-D feature transitions so it judges gait quality, not global heading.
## Files
- `humanoid_direction_env.py` / `register_env.py``HumanoidDirection-v0` (Humanoid-v5 + target heading, bounded directional reward)
- `amp_obs.py` — shared heading-invariant 45-D AMP features (used by both policy and mocap sides)
- `amp_env.py` — AMP reward wrapper (task/style mixing, reference state init)
- `amp_discriminator.py` / `amp_callback.py` — LSGAN discriminator + training callback (syncs weights into `SubprocVecEnv` workers)
- `motion_lib.py` — loads mocap `.pkl` clips, filters bad contacts, samples expert transitions interpolated at the exact env control dt
- `train_sb3_real_amp.py` — training (exp 1); `train_sb3_real_amp_exp2.py` — more conservative discriminator (exp 2)
- `humanoid_direction_evaluate.py` — roll out a trained policy, print per-episode and mean ± std reward
- `models/`, `models_exp2/` — trained PPO policies, `VecNormalize` stats, discriminators
## Training procedure
Trained with `train_sb3_real_amp.py` on 8 `SubprocVecEnv` workers (CPU) + 1 GPU for PPO/discriminator updates, up to 50M environment steps. The discriminator is updated once per rollout (8 × 2048 steps) against expert transitions sampled from mocap at the exact environment control dt, and its weights are pushed to every worker after each update. Episodes start from a random mocap frame (Reference State Initialization) with probability 0.5.
### Hyperparameters
**PPO**: lr 1e-4, n_steps 2048, batch 512, 5 epochs, target_kl 0.02, γ 0.99, GAE λ 0.95, clip 0.2, ent_coef 0.0, obs normalization (clip 10).
**AMP discriminator**: 90-D input (two 45-D heading-invariant frames), 2×512 MLP; LSGAN loss to +1/−1 targets; Adam lr 1e-4, weight decay 1e-4; 8 updates × batch 512 per rollout; gradient penalty 5.0 on real samples; score regularization 1e-4; grad-norm clip 1.0; 100k fake-transition replay buffer.
**Reward mixing**: 0.5 · task + 0.5 · style.
A second experiment (`models_exp2/`) uses a more conservative discriminator: disc lr 3e-5, 4 updates/rollout, gradient penalty 10.0, PPO lr 5e-5, RSI 0.3.
### Training data
~60 motion-capture clips (stand, walk, run, turns 45°–135°, backwards locomotion, gait transitions) retargeted to the Humanoid-v5 skeleton, stored as `qpos`/`qvel`/`fps` pickles. Clips originate from AMASS-format mocap (ACCAD subject sets, among others). **The original mocap datasets carry their own licenses (typically non-commercial for AMASS subsets) — verify before redistribution or commercial use.**
## Usage
```bash
pip install "gymnasium[mujoco]" stable-baselines3 torch numpy mujoco
python train_sb3_real_amp.py # train (SLURM scripts included)
python humanoid_direction_evaluate.py # evaluate + render
```
The `VecNormalize` stats must be loaded alongside the PPO model at eval time — the policy was trained on normalized observations.
## Evaluation
Qualitative: see `replay.mp4` in the repo for a sample rollout. Run `humanoid_direction_evaluate.py` to render episodes and print per-episode rewards.
Tensorboard data is available in repository.
## Limitations
- Locomotion only (walk/run/turn/stand); no other skills are in the motion prior.
- Target speed fixed at 1.4 m/s; direction is randomized but speed is not commanded.
- Flat-ground MuJoCo simulation; no domain randomization — not directly transferable to a real robot.
- The policy is tied to this exact observation layout and `VecNormalize` statistics.
## Citation
```bibtex
@article{peng2021amp,
title={AMP: Adversarial Motion Priors for Stylized Physics-Based Character Control},
author={Peng, Xue Bin and Ma, Ze and Abbeel, Pieter and Levine, Sergey and Kanazawa, Angjoo},
journal={ACM Transactions on Graphics (TOG)},
volume={40}, number={4}, pages={1--20}, year={2021}
}
```