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

Five controllers, 8 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. Two of the five constrain the action space to a 4-D physics prior (PIRL and the learned CPG); one relaxes that constraint by adding back per-joint corrections (residual); one is unconstrained 12-D torque control.

Flat terrain

Controller	Action space	Speed (m/s)	Energy (Στ²)	Cost of transport
Learned CPG	4-D prior	0.843 ± 0.255	0.732 ± 0.337	0.840 ± 0.332
PIRL (serpenoid)	4-D prior	1.249 ± 0.179	1.535 ± 0.213	1.263 ± 0.313
Residual PIRL	4-D + 12-D	1.224 ± 0.071	1.702 ± 0.324	1.404 ± 0.319
Residual (tuned)	4-D + 12-D	1.272 ± 0.045	1.766 ± 0.406	1.399 ± 0.370
Raw-torque PPO	12-D	1.591 ± 0.186	9.315 ± 0.273	5.934 ± 0.787

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

Controller	Action space	Speed (m/s)	Energy (Στ²)	Cost of transport
Learned CPG	4-D prior	0.585 ± 0.061	1.878 ± 0.488	3.206 ± 0.775
PIRL (serpenoid)	4-D prior	0.591 ± 0.172	2.419 ± 0.842	3.970 ± 0.782
Residual PIRL	4-D + 12-D	0.675 ± 0.075	2.951 ± 0.354	4.422 ± 0.768
Residual (tuned)	4-D + 12-D	0.651 ± 0.064	2.826 ± 0.317	4.393 ± 0.795
Raw-torque PPO	12-D	0.861 ± 0.106	7.513 ± 0.692	8.787 ± 0.865

The central claim holds, decisively

Constraining the action space to a physics prior buys a large, significant efficiency advantage. Against unconstrained 12-D torque control, the serpenoid agent cuts actuation energy 6.1× on flat (p < 0.001, Hedges' g = −30.0) and 3.1× on hard terrain (p < 0.001, g = −6.25). Cost of transport falls 4.7× and 2.2× respectively, both at p < 0.001. These are not marginal effects that a few unlucky seeds could produce; at n = 8 they are among the largest effect sizes in the study.

Raw torque remains the fastest, and it should be. It has strictly more authority — 12 independent degrees of freedom against 4 wave parameters — so the interesting question was never whether it could go faster (p = 0.002 flat, p = 0.003 hard) but what that speed costs. It costs roughly six times the energy per metre. That trade is the thesis, and it is now established at n = 8 on two terrains.

Relaxing the constraint does not help. The residual agent adds 12 learned per-joint corrections on top of the prior — the natural "best of both" move. It matches the plain prior and never beats it: flat speed p = 0.719, CoT p = 0.385; hard speed p = 0.236, CoT p = 0.262. Both tunings, both terrains, every metric. It consistently costs more energy for no speed it can demonstrate. Handing the policy back the degrees of freedom the prior removed bought nothing, which is positive evidence for constraining them in the first place.

Which prior matters

The two physics-informed controllers share an action space of exactly four parameters; they differ only in the structure of the prior. The serpenoid agent prescribes joint angles kinematically — the policy commands a shape and the PD controller is asked to track it. The learned CPG modulates a coupled Hopf oscillator, so the shape emerges from a dynamical system relaxing onto its limit cycle rather than being demanded of the joints.

That structural difference is measurable:

	Flat CoT	Hard CoT	Hard speed
Learned CPG	0.840 ± 0.332	3.206 ± 0.775	0.585 ± 0.061
PIRL (serpenoid)	1.263 ± 0.313	3.970 ± 0.782	0.591 ± 0.172
p	0.020	0.070	0.928

On hard terrain the two are statistically indistinguishable on speed — p = 0.928, g = +0.04, about as close as two separately trained controllers get — while the CPG uses 22% less energy to do it. On flat ground it is significantly more efficient (p = 0.020, g = +1.24) at the cost of being 33% slower (p = 0.003).

This is not the "stand still and score well" artifact. An earlier version of this README dismissed the CPG's flat-terrain efficiency on exactly that ground: at n = 3 one seed had converged to 0.303 m/s, and a near-stationary robot spends almost no energy per metre. With 8 seeds the claim does not survive scrutiny — dropping that seed entirely leaves CoT at 0.912 ± 0.284 (n = 7) against the serpenoid agent's 1.263 ± 0.313, still significant at p = 0.040. Within the CPG arm, speed and CoT are uncorrelated (flat r = +0.33, p = 0.42; hard r = +0.03, p = 0.95), so the efficiency is not being bought by moving slowly. The n = 3 verdict was a small-sample artifact, not the result.

What the two priors actually converged to

scripts/analyse_cpg.py rolls both physics-informed agents out over the full grid the study evaluates on — 8 seeds × 5 friction values, 80 rollouts each — and records the wave parameters they settle at. It produced one problem with the study and one explanation of its main result.

Both agents pin the wave frequency to its ceiling. Across all 160 rollouts, on both terrains, ω = 6.000 rad/s with a standard deviation of exactly zero. On flat ground the phase lag is likewise pinned to its floor (30.000° ± 0.000), which is the longest wavelength the bounds permit. These are hand-set constants in WAVE_BOUNDS, chosen before any of these experiments ran.

This matters more than it first looks. The speed numbers in this README are not measurements of what the physics priors can do — they are measurements of what they do while frequency is capped at 6.0 rad/s. Both agents want to run the wave faster and are not permitted to. Any claim of the form "the physics prior limits top speed" is unsupported until that bound is widened, which is why the experiment now leads the roadmap.

On flat ground, the efficiency gap is an amplitude choice. Once both agents saturate frequency and phase lag, the only wave parameter left free is amplitude — and they choose differently:

Flat terrain, 40 rollouts each	PIRL (serpenoid)	Learned CPG
Frequency ω (rad/s)	6.000 ± 0.000	6.000 ± 0.000	identical (at bound)
Phase lag (deg)	30.000 ± 0.000	30.000 ± 0.001	identical (at bound)
Amplitude (deg)	27.49 ± 7.36	22.42 ± 8.14	0.82× (p = 0.005)
Energy per step	1.511 ± 0.608	0.732 ± 0.476	0.48× (p < 1e-7)
Distance (m)	25.29 ± 5.59	16.92 ± 5.59	0.67× (p < 1e-8)
Energy per metre	0.0598	0.0432	0.72×

The CPG runs an 18% shallower wave. That costs it a third of the distance — but it saves half the energy. Energy falls faster than distance does, and the ratio between them is exactly what cost of transport measures. Bending the body less is cheaper than the speed it gives up, and this analysis reproduces the study's efficiency ordering independently (0.0432 against 0.0598).

On hard terrain the same analysis does not reproduce the study's ordering, giving pooled energy-per-metre of 1.716 for the serpenoid agent against 1.989 for the CPG — the opposite sign. The likely reason is protocol: hard terrain randomises ground tilt and start pose per episode, and this analysis runs one fixed-seed episode per cell where the study averages three stochastic ones. One episode is too thin a sample of that variation. The study's number is the trustworthy one, so on hard terrain the efficiency result stands but its mechanism remains open.

What the hard-terrain rollouts do show is that the CPG converges to a far more stereotyped gait — phase lag 30.18° ± 0.41 against the serpenoid agent's 35.38° ± 20.90, a 51× tighter spread. That matches the outcome variance there, where the CPG's speed s.d. is 2.8× tighter (0.061 vs 0.172, Brown-Forsythe p = 0.071). It does not hold on flat, where the CPG is the more variable of the two, so this is a hard-terrain observation rather than a general property.

What the extra seeds took away

The residual's hard-terrain advantage evaporated. At n = 3 it looked decisive — 0.654 ± 0.049 against the prior's 0.466 ± 0.214 — and it was the headline hypothesis. The prior's three-seed sample happened to contain its bad seeds; five more brought its mean to 0.591 and the gap fell inside the noise (p = 0.236, Mann-Whitney p = 0.382). Nothing was wrong with the earlier arithmetic. There were simply not enough seeds to support the conclusion drawn from it.

The reliability effect is consistent but not significant. Speed s.d. for the prior against the residual is 2.53× on flat and 2.29× on hard — two independent conditions landing on nearly the same ratio. But Brown-Forsythe gives p = 0.107 and p = 0.141, and combining them (Fisher) gives p = 0.078. The CoT variance ratio is 0.98× and 1.02×, so whatever the effect is, it lives in speed alone.

The cleaner description is a failure rate rather than a spread: on hard terrain the prior has 2 of 8 seeds below 0.5 m/s (worst 0.237) where the residual has 0 of 8 (worst 0.534). That is also not significant at this n (Fisher exact p ≈ 0.47). Reported as an observation, not a finding.

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 was included as a control: it gives an RL agent the same 4-D authority over a different prior, so any gap between the two is attributable to the structure of the prior rather than to the size of the action space or to the presence of learning. That control is what makes its result interpretable -- it is the most efficient controller in both conditions (see Which prior matters), and because the action space is held fixed at four parameters, the difference can only come from the prior's structure.

Reproducing the study

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

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

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

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 5 --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 8 --timesteps 150000 --condition flat --outroot runs/study_tuned --env-override '{"reward": {"w_residual": 0.005}}'

Inspect what each physics prior converged to (wave parameters and the torque they cost, over the same seed x friction grid the study evaluates on):

python scripts/analyse_cpg.py --condition hard

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
├── run_study.py             # multi-seed study over one terrain condition
├── run_tier1.py             # drives the full 8-seed study in resumable stages
├── build_compare_tuned.py   # assembles the three-way tuned comparison
├── sweep_residual.py        # grid over residual_scale x w_residual
├── analyse_cpg.py           # what each physics prior converged to
└── run_transfer.py          # zero-shot morphology transfer (not yet run)
tests/                       # 34 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.

Results are 8 seeds per arm on two terrains. That is enough to establish the efficiency advantage of the physics prior with very large effect sizes, and enough to show the residual agent does not beat it. It is not enough to settle the variance difference between them (Fisher-combined p = 0.078), which would need roughly 15-20 seeds per arm.

Both physics-informed agents drive the wave frequency to the ceiling of its allowed range (omega = 6.0 rad/s) in every rollout measured, so the gaits reported here are shaped by that bound rather than freely chosen. Widening WAVE_BOUNDS is the obvious next experiment and would change the speed numbers -- though not, on the evidence here, the ordering by efficiency.

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 an 8-seed harness with Welch tests, effect sizes, and robust variance tests.

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

• Widen the wave-parameter bounds. Both priors saturate omega at 6.0 rad/s with zero variance across 80 rollouts. The bound, not the policy, is setting the gait frequency, and no result here can be read as the prior's true limit.
• Investigate the CPG prior further. It is the most efficient controller in both conditions and reaches the serpenoid agent's hard-terrain speed exactly. It entered the study as a control and finished as the most interesting arm.
• 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.
• More seeds for the variance question only. Everything else is settled at n = 8; --seed-offset with --skip-existing appends seeds without redoing finished runs.
• Hardware. Sim-to-real remains untouched.

License

Released under the MIT License.