Physics-Informed Reinforcement Learning for Gait Optimization

Gait optimization for hyper-redundant snake / continuum robots, where a recurrent RL agent modulates the parameters of a serpenoid travelling-wave prior instead of commanding raw joint torques.

📖 New here? EXPLAINER.md explains the whole project in plain language — no maths, no jargon — including what worked, what didn't, and the two silent bugs that nearly invalidated the results.

Controlling a many-jointed robot by emitting one torque per joint suffers from the curse of dimensionality: random exploration rarely finds the narrow manifold of coordinated, wave-like motions, so the agent burns its budget on uncoordinated flailing. This project constrains the action space to four wave parameters — regardless of body length — so that even an untrained policy produces a coherent, physically meaningful gait.

            ┌──────────────┐  [Δα, Δω, Δφ, γ]  ┌────────────┐  θ_ref  ┌──────────────┐  τ   ┌─────────┐
 obs ──────▶│  RL agent    │──────────────────▶│ serpenoid  │───────▶│ PD low-level │────▶│ MuJoCo  │
 (history)  │ (LSTM policy)│  4-D constrained  │ generator  │        │  controller  │     │  snake  │
            └──────────────┘   action space    └────────────┘        └──────────────┘     └─────────┘
                  ▲                                                                             │
                  └─────────────────── joint state, COM velocity, contact sensors ──────────────┘

Results

Four controllers, 3 random seeds each, 150k steps, all under the same algorithm (recurrent PPO), evaluated across a hidden friction sweep (μ ∈ {0.4, 0.7, 1.0, 1.3, 1.6}, never shown to the agent). Reported as mean ± s.d. over seeds.

Flat terrain

Controller	Speed (m/s)	Energy (Στ²)	Cost of transport
Residual PIRL	1.218 ± 0.066	1.715 ± 0.205	1.411 ± 0.195
PIRL (serpenoid)	1.270 ± 0.041	1.458 ± 0.127	1.147 ± 0.067
Learned CPG	0.795 ± 0.442	0.593 ± 0.493	0.674 ± 0.443
Raw-torque PPO	1.762 ± 0.048	9.190 ± 0.273	5.214 ± 0.044

Hard terrain (position-varying friction field + random ground tilt)

Controller	Speed (m/s)	Energy (Στ²)	Cost of transport
Residual PIRL	0.654 ± 0.049	2.642 ± 0.076	4.058 ± 0.374
PIRL (serpenoid)	0.466 ± 0.214	1.838 ± 1.204	3.603 ± 1.147
Learned CPG	0.592 ± 0.025	2.149 ± 0.541	3.619 ± 0.834
Raw-torque PPO	0.784 ± 0.077	7.274 ± 0.817	9.267 ± 0.277

What holds up

The efficiency advantage is large and statistically significant. Against raw-torque PPO, the physics-informed agents cut cost of transport by 3.7× on flat (p = 0.001, Hedges' g = −21.6) and 2.3× on hard terrain (p < 0.001, g = −12.7), with actuation energy 5.4× and 2.8× lower respectively. These gaps survive the seed-to-seed spread comfortably.

Raw-torque is still the fastest — but its edge narrows with difficulty. On flat ground it is decisively quicker (p = 0.001); on hard terrain the gap over residual PIRL is no longer significant at the 5% level (p = 0.079).

What did not

The residual agent did not break the speed ceiling. It was the headline hypothesis — add a learned correction, keep the efficiency, escape the prior's limit. At the original penalty weight it came out marginally slower than the plain serpenoid agent on flat ground (1.218 vs 1.270 m/s, p = 0.32) while using more energy.

That regression turned out to be a hyperparameter artifact. The penalty w_residual had been set once by hand and never swept; a 3 × 3 grid found it was roughly 4× too strict. Re-running the full multi-seed study with w_residual = 0.005 (the only change) recovers the entire deficit:

Flat terrain	Speed (m/s)	Energy	CoT
PIRL (prior only)	1.270 ± 0.041	1.458	1.147
Residual, w = 0.02	1.218 ± 0.066	1.715	1.411
Residual, w = 0.005	1.269 ± 0.060	1.900 ± 0.616	1.515

Hard terrain	Speed (m/s)	Energy	CoT
PIRL (prior only)	0.466 ± 0.214	1.838	3.603
Residual, w = 0.02	0.654 ± 0.049	2.642	4.058
Residual, w = 0.005	0.705 ± 0.049	2.774	3.953

But fixing it bought nothing net. The tuned agent reaches parity with the prior on flat ground (p = 0.978, g = 0.02) rather than beating it, and pays for that in energy — it has the worst actuation energy and worst cost of transport of the three. That is precisely the trade this project argues against. Plain PIRL retains the best cost of transport in both conditions.

The durable effect is reliability on hard terrain, and it survived tuning. The tuned residual is 51% faster than the plain prior (0.705 vs 0.466) with a 4.4× tighter spread (± 0.049 vs ± 0.214). Per seed that is [0.739, 0.726, 0.648] against PIRL's [0.237, 0.662, 0.499] — two of three PIRL seeds fall below every residual seed.

This experiment cannot establish that claim, by construction. The gap is p = 0.189, and with 3 seeds per arm no test can reach p < 0.05 for an effect this size: a rank-based test at n = 3 vs 3 has a floor of p = 0.1 even under perfect separation. Settling it needs more seeds, not better luck — a design limitation that should have been priced in before the runs.

The learned CPG's apparent efficiency win on flat ground is an artifact. Its per-seed speeds were [1.158, 0.925, 0.303] — one seed learned to barely move, and a near-stationary robot scores well on energy-per-distance. Treat its flat-terrain cost of transport as unreliable.

Two measurement notes

Cost of transport is aggregated as total energy ÷ total distance, not as the mean of per-episode ratios. On tilted ground a single stalled episode (distance → 0) sends its own ratio into the hundreds and hijacks the average; the earlier per-episode averaging reported CoT values of 36–104 with ±63 spread for exactly this reason.

All agents now share one algorithm. The earlier comparison paired the physics-informed agent with a recurrent policy and the raw-torque baseline with a feed-forward one, which confounded the action space with the policy class.

Speed and energy alone cannot tell a clean travelling wave from thrashing that happens to drift forward, so gaits are also inspected directly — overhead path plus the body waveform in a head-to-tail-aligned frame:

Method

Serpenoid prior

Each joint i tracks a reference angle

θ_i(t) = α · sin(ω·t + (i−1)·φ) + γ

symbol	meaning	who sets it
α	wave amplitude (body curvature)	agent (Δα)
ω	temporal frequency (wave speed)	agent (Δω)
φ	inter-joint phase lag (wavelength)	agent (Δφ)
γ	turning bias (steering)	agent (γ)

The action is A_t = [Δα, Δω, Δφ, γ] — 4-D instead of 12-D.

Why it moves

Locomotion requires anisotropic ground friction: a body wave only produces net thrust if sliding sideways is harder than rolling forward. Following Hirose's classic wheeled snake robots, every segment carries a passive wheel that rolls freely along the body axis but resists lateral slip.

Friction is declared through explicit wheel–floor contact pairs, not the default geom-combination rule — MuJoCo combines two geoms' friction by the element-wise maximum, which would silently clamp any floor coefficient below the wheel's own value and make most of the sweep inert.

Hidden terrain, recurrent policy

The friction coefficient is randomized per episode and never observed, so an LSTM policy must infer the terrain from the history of contact and slip signals.

Multi-objective reward

R_t = w1·v_∥ − w2·Σ τ_i² − w3·p_slip − w4·(1 − cos ψ) + R_alive

Progress v_∥ is measured along a fixed goal axis and ψ is the heading error. This matters: crediting progress along the instantaneous heading pays out identically whether the robot moves straight or banks through a wide circle — and agents exploit exactly that. The fixed axis plus heading penalty is what makes the learned gait genuinely straight. Weights live in configs/default.yaml.

Install

pip install -r requirements.txt

Python ≥ 3.9. Uses the official DeepMind mujoco bindings (no license needed).

Quick start

pytest -q

Run the whole study end to end:

python scripts/run_experiment.py --timesteps 300000

Or step by step:

python -m pirl.train --agent pirl --algo recurrent_ppo --timesteps 300000

python -m pirl.train --agent torque --algo ppo --timesteps 300000

python -m pirl.evaluate --agent pirl --algo recurrent_ppo --model runs/pirl_recurrent_ppo/final_model --out runs/eval_pirl.json

python -m pirl.viz.compare runs/eval_pirl.json runs/eval_torque.json runs/eval_cpg.json --out runs/comparison.png

Record a gait video (the camera tracks the robot):

python -m pirl.viz.render_gait --agent pirl --algo recurrent_ppo --model runs/pirl_recurrent_ppo/final_model --out runs/gait.mp4

Visualizing results

All figures regenerate from an existing run without retraining:

python -m pirl.viz.plot_summary --out runs/summary.png

python -m pirl.viz.plot_curves runs/pirl_recurrent_ppo runs/torque_ppo --labels "PIRL (serpenoid)" "Raw-torque PPO"

python -m pirl.viz.plot_gait_trace --out runs/gait_traces.png

python -m pirl.viz.stack_videos runs/gait_pirl.mp4 runs/gait_torque.mp4 runs/gait_cpg.mp4 --labels PIRL Raw-torque CPG --out runs/gait_compare.mp4

Experimental comparison

Controller	Action space	Prior	Implementation
Physics-informed RL	4 wave params	serpenoid curve	pirl/envs/serpenoid_env.py
Residual PIRL	4 wave + 12 residual	serpenoid + learned correction	pirl/envs/residual_env.py
Learned CPG	4 oscillator params	coupled Hopf oscillators	pirl/envs/cpg_env.py
Raw-torque RL	12 joint torques	none	pirl/envs/torque_env.py
Fixed CPG	none (open loop)	fixed wave, no learning	pirl/control/cpg.py

The learned CPG matters as a control: it gives an RL agent the same 4-D authority over a different prior, so any remaining gap is attributable to the structure of the serpenoid prior rather than to the size of the action space or to the presence of learning.

Reproducing the study

python scripts/run_study.py --seeds 3 --timesteps 150000 --condition flat

python scripts/run_study.py --seeds 3 --timesteps 150000 --condition hard

python -m pirl.viz.aggregate runs/study/flat --reference residual

python -m pirl.viz.plot_study runs/study/flat runs/study/hard --labels flat hard --out runs/study/conditions.png

Add seeds to a finished study without retraining the ones already done:

python scripts/run_study.py --condition flat --seeds 3 --seed-offset 3 --skip-existing

Terrain presets for --condition: flat, hard (friction field + slope), obstacles, hard_obstacles.

Zero-shot morphology transfer — train one policy, run it unchanged on bodies with different link counts (a raw-torque policy structurally cannot, since its output layer is welded to the joint count it trained on):

python scripts/run_transfer.py --timesteps 150000 --joints 8,10,12,16,20

Sweep the residual agent's two free parameters (this is what identified w_residual = 0.02 as ~4× too strict):

python scripts/sweep_residual.py --scales 0.15,0.35,0.7 --weights 0.0,0.005,0.02 --timesteps 80000

Re-run the study at the tuned weight, changing only that one parameter:

python scripts/run_study.py --agents residual --seeds 3 --timesteps 150000 --condition flat --outroot runs/study_tuned --env-override '{"reward": {"w_residual": 0.005}}'

Any environment setting can be overridden from the command line for a one-off run, without editing the config:

python -m pirl.train --agent residual --env-override '{"residual_scale": 0.7, "reward": {"w_residual": 0.005}}'

Project layout

pirl/
├── models/build_snake.py    # procedural MJCF snake (wheeled, variable friction)
├── control/
│   ├── serpenoid.py         # serpenoid wave generator (the physics prior)
│   ├── pd_controller.py     # low-level PD torque tracker
│   └── cpg.py               # Baseline 2: fixed central pattern generator
├── envs/
│   ├── snake_base.py        # dynamics, observation, reward, friction, MDP
│   ├── serpenoid_env.py     # physics-informed env (4-D action)
│   └── torque_env.py        # Baseline 1 env (12-D raw-torque action)
├── train.py                 # RecurrentPPO / PPO / SAC training
├── evaluate.py              # metrics + friction-sweep rollouts
├── utils.py                 # config loading + environment factory
└── viz/                     # visualization & reporting layer
    ├── compare.py           # comparison table & per-metric plots
    ├── plot_summary.py      # headline multi-panel results figure
    ├── plot_curves.py       # learning curves from the TensorBoard logs
    ├── plot_gait_trace.py   # overhead path + body-frame waveform
    ├── stack_videos.py      # side-by-side gait comparison video
    └── render_gait.py       # rollout -> mp4/gif
configs/default.yaml         # all hyper-parameters
scripts/run_experiment.py    # end-to-end train -> evaluate -> compare
tests/                       # 13 unit + integration tests
assets/                      # result figures used by this README

Scope and limitations

The physics is a planar, wheeled abstraction of anisotropic snake–ground friction. It reproduces the locomotion behavior and supports the full research loop, but it is not a substitute for a 3-D contact model or hardware — no sim-to-real claims are made. Terrain friction is a single scalar per episode rather than spatially varying, and results are reported for a single seed; multi-seed statistics are the next step.

Roadmap

Implemented (see Results): residual PIRL, a learned coupled-oscillator CPG competitor, harder terrain (friction field, slopes, obstacles), morphology-agnostic observations for zero-shot transfer, and a multi-seed harness with Welch tests and effect sizes.

The tooling for every remaining item is in place; what is left is compute.

• More seeds — now the binding constraint. Three seeds cannot establish the residual agent's hard-terrain advantage no matter how large it is (see the power note above). Five to ten would settle it; --seed-offset with --skip-existing appends them without redoing finished runs.
• Longer budget. 150k steps was chosen to fit 24 runs; the residual agent's 16-D action space plausibly needs more than the plain prior's 4-D to pay off.
• Run the transfer experiment. scripts/run_transfer.py trains the fixed-width policy and sweeps it over body sizes; it has not been run yet.
• Obstacle conditions. Implemented and selectable (--condition obstacles), but not yet part of a study.
• The "stand still" local optimum. One learned-CPG seed converged to barely moving, which scores acceptably because a stationary robot spends almost no energy. stall_steps / stall_distance end such an episode early; disabled by default so existing results stay comparable.
• Hardware. Sim-to-real remains untouched.

License

Released under the MIT License.