Preprint · University of Wisconsin–Madison, Simulation-Based Engineering Lab

Learning the Right Abstraction

Neural Reduced Dynamics for Complex Robot Control

Harry Zhang  ·  Dan Negrut

Department of Mechanical Engineering, University of Wisconsin–Madison

High-fidelity simulators like Project Chrono give you the physics to evaluate a controller, but not the throughput to find one. We distill Chrono into a task-specific neural reduced dynamics model (NN-ROM), train the policy entirely inside the frozen model, and send it back to Chrono for closed-loop validation. The question the paper answers is not how to compress state, but what a learned model should keep, receive, reconstruct, and omit so that it is both fast and useful for control.

15-D / 3-D / 8-Dreduced states across three control tasks
3 / 3 terrainsone generalist beats both specialists
100/100 & 97/100Chrono goals reached, base & arm
~104Γ—faster simulated time than Chrono
Four-stage pipeline: Chrono data collection, causal-transformer NN-ROM, vectorized PPO in the frozen model, closed-loop Chrono validation.
One reusable pipeline. Chrono generates excited state–action trajectories (1); a causal transformer learns next-step reduced dynamics from a history window (2); PPO trains a policy against thousands of copies of the frozen model (3); the policy is returned unchanged to the full Chrono scene (4). Models are selected by open-loop rollout fidelity, not one-step loss.

Study Case I β€” one policy, three terrains

Nine closed-loop Chrono rollouts. Three tracking policies (columns), identical in every way except the frozen dynamics model they were trained inside, on three terrains (rows). Green is the reference, red the path actually driven β€” where you see only green, the vehicle is still on the line.

Rigid flat

A held-out reference from the rigid validation set.

Generalistmixture β€” rigid + CRMmedian XY RMSE 0.157 m
Rigid-onlyspecialistmedian XY RMSE 0.174 m
CRM-onlyspecialistmedian XY RMSE 0.232 m

CRM soil

Deformable terrain β€” the wheels sink and slip, and the rigid-only policy overshoots every turn.

Generalistmixture β€” rigid + CRMmedian XY RMSE 0.180 m
Rigid-onlyspecialistmedian XY RMSE 0.854 m
CRM-onlyspecialistmedian XY RMSE 0.231 m

Rigid bumpy out of distribution

No bumpy data enters training, checkpoint selection, normalization or reward tuning, and bumpy carries the same rigid context code as flat.

Generalistmixture β€” rigid + CRMmedian XY RMSE 0.149 m
Rigid-onlyspecialistmedian XY RMSE 0.187 m
CRM-onlyspecialistmedian XY RMSE 0.213 m

Each clip is one representative held-out reference for that terrain, rendered in Blender from the Chrono rollout; the quoted medians are over all 20 held-out references, so an individual clip can run against its policy's median. Clips loop β€” replay row restarts three in sync.

The generalist wins on every terrain β€” including the one it never saw

Closed-loop Chrono XY RMSE over 20 held-out references per terrain, all 20/20 rollouts completing. The mixture generalist takes the lowest median and mean in all three regimes: it beats the rigid-only specialist on rigid ground and the CRM-only specialist on CRM soil, on their own home terrain. Bumpy is the out-of-distribution case β€” the policy never previews the height profile, so closed-loop tracking there is the honest measure of whether the conditioned abstraction generalizes or memorizes.

Bar chart of median XY RMSE per terrain for the three policies, with interquartile-range whiskers.
Median XY RMSE, whiskers span the IQR. The rigid-only policy collapses on CRM soil (0.85 m median).
Closed-loop Chrono XY RMSE (m), 20 references, iteration-1000 checkpoints. † out of distribution: rigid context code, no bumpy data anywhere in training or selection.
TerrainPolicyMedianMean
Rigid flatGeneralist0.1570.184
Rigid-only0.1740.219
CRM-only0.2320.259
CRM soilGeneralist0.1800.249
Rigid-only0.8541.000
CRM-only0.2310.361
Rigid bumpy†Generalist0.1490.229
Rigid-only0.1870.238
CRM-only0.2130.418
Grid of closed-loop XY paths for three maneuvers on each of the three terrains, comparing reference against the three policies.
Closed-loop XY paths for three held-out maneuvers (columns) on each terrain (rows). The rigid-only specialist overshoots every turn on CRM soil; the CRM-only specialist is the outlier on bumpy ground despite sharing the rigid contact code. Panel annotations are per-episode RMSE for the generalist (G), rigid-only (R), and CRM-only (C). Row (c) is the out-of-distribution case.

Why context, not just capacity

The same reduced state and action evolve differently on rigid ground than on deformable soil, so the reduced dynamics are ambiguous without a regime label. A two-class one-hot code appended to each input token resolves it. Removing that code barely moves one-step loss β€” but more than doubles open-loop rollout error on rigid terrain, 3.73% β†’ 8.26%. History alone does not preserve the rigid/CRM distinction over a long rollout.

Study Case II β€” same pipeline, two very different abstractions

One M113 tracked platform with a front-mounted 4-DOF arm, driven by two independent tasks. The base keeps a 3-D planar state [vx, vy, r] and integrates its pose outside the network. The arm keeps an 8-D joint state [q, q̇], with the end-effector recovered by forward kinematics rather than learned as a channel. Dimension is a consequence of the dominant physics, not the design target.

Tracked-vehicle goal reaching and arm end-effector reaching, each trained by PPO inside its own frozen NN-ROM.
(a) Tracked-base goal reaching and (b) arm end-effector reaching, each trained inside its own frozen NN-ROM and validated in Chrono.
100 / 100tracked-base goals reached within 0.75 m in Chrono
97 / 100arm goals reached within 0.05 m
0arm contacts or joint-limit violations across all rollouts
10 min / 57 minsingle-GPU PPO training for base / arm policy

High fidelity supplies the physics to judge a controller, not the throughput to find one

Real-time factor (RTF) is simulated seconds advanced per wall-clock second; RTF = 1 is real time. Cost is normalized to each case's own NN-ROM, so it reads directly as how many times more wall-clock time the high-fidelity scene spends per simulated second than the surrogate that replaces it.

ConfigurationΞ”tSteps/sRTFCost
HMMWV β€” rigid flat2.0 ms3,6577.31Γ—239Γ—
HMMWV β€” bumpy rigid2.0 ms3,3246.65Γ—263Γ—
HMMWV β€” CRM soil0.5 ms3040.152Γ—11,500Γ—
HMMWV β€” NN-ROM (L8)‑10 ms175,0001,750Γ—1Γ—

Study Case I. The Chrono rows share one HMMWV_Full configuration and differ only in the tire–terrain contact model.

ConfigurationΞ”tSteps/sRTFCost
M113 tracked + arm β€” rigid0.5 ms5580.279Γ—58,400Γ— / 8,300Γ—
Tracked base β€” NN-ROM‑20 ms815,00016,300Γ—1Γ—
Arm β€” NN-ROM‑20 ms116,0002,320Γ—1Γ—

Study Case II. One Chrono scene serves both control modes, so its Cost cell gives the ratio over the drive- / reach-mode NN-ROM. The comparison is not symmetric: Chrono simulates the full coupled vehicle-plus-arm system, each surrogate only the subsystem its mode controls.

‑ NN-ROM rows are aggregate batched-GPU figures β€” dynamics forward passes per second summed across the 2,048–4,096 parallel rollouts used during PPO β€” not single-stream processes like the Chrono rows, which are one process on one Intel 14900KF.

What the margin actually buys

CRM is the dominant Chrono cost, and the reason is the terrain rather than the vehicle: ~14.1M SPH particles, ~10.8M boundary markers, and a 4Γ— smaller step. A 15 s episode costs ~100 s of wall clock. PPO consumes thousands of parallel rollouts over hundreds of iterations, so at that rate the search is simply not affordable.

A single converged run understates the point. The expensive part of applying RL is not that run but the design loop around it β€” shaping the reward and tuning the optimizer, each adjustment validated only by training again. Both Study Case II policies needed exactly that: spin and drive-against-brake penalties for the base; terminal bonus, action regularization, clearance observation and command shield for the arm. The effective cost of a policy is one run times the number of design iterations, so a run finishing in 10 minutes (base) or 57 minutes (arm) instead of the orders of magnitude implied above is what keeps that loop interactive β€” more reward and hyperparameter iterations per day, which is what ultimately yields a working policy.

What makes an abstraction the right one

Follow the dominant physics

The HMMWV needs tire normal loads and wheel speeds because tire–terrain interaction decides whether it holds a line. The base needs three planar velocities. The arm needs joint positions and rates. Different in kind, not just in count.

Don't learn what you can derive

Global pose is integrated from predicted velocities; the end-effector is FK(q); contact clearance is geometric. Reduction discards redundancy, not fidelity.

Add context only where ambiguous

A two-class terrain code where rigid and deformable dynamics diverge β€” nothing more. One backbone then shares what the regimes have in common.

Judge by transfer, not by loss

One-step error is necessary and insufficient; a policy can farm reward from a surrogate's mistakes. Closed-loop Chrono is the exam.

Every dataset behind these models is public

Every Chrono collection behind this work is released on Hugging Face — 70 GB across five datasets, nothing to re-collect. hmmwv_flat (32,768 episodes, 160M transitions at 100 Hz) and hmmwv_crm (2,000 episodes on deformable SPH soil) train the terrain-conditioned Study Case I model; hmmwv_bumpy (1,360 episodes over 100 random heightmaps) is the out-of-distribution test regime and enters no training, selection or normalization; tracked (2,160) and arm (15,000) train the two Study Case II models. Each ships in two tiers: raw, every recorded channel of every episode as float32 Parquet with per-episode metadata and the bumpy heightmaps — enough to define a different reduced state than ours; and processed, the exact training caches the deployed checkpoints read, so the repository's configs run verbatim.

from datasets import load_dataset

ds = load_dataset("harryzhang1018/NeDM",
                  "hmmwv_crm", split="val",
                  streaming=True)

BibTeX

@article{zhang2026abstraction,
  title   = {Learning the Right Abstraction: Neural Reduced Dynamics for
             Complex Robot Control},
  author  = {Zhang, Harry and Negrut, Dan},
  journal = {Preprint},
  year    = {2026}
}