Diffeomorphic Hidden-State Steering: Improving Imitation Learning Policy Robustness through Expert Hidden-State Geometry

Can the original expert demonstrations alone make a trained imitation-learning policy more robust in closed loop, with no extra demonstrations, no reward signal and no change to the policy's weights?

Small errors in observations or actions push an imitation-learning policy into states its demonstrations never covered, and the errors compound until the task fails. The usual fixes, DAgger-style expert relabelling, recovery demonstrations or reinforcement-learning fine-tuning, all need information beyond the original dataset, which is expensive on a production line.

Earlier hidden-state corrections either trained a correction module together with the policy or fitted steering filters to successful and failed rollouts. A simple alternative fits a linear filter (PCA) to the expert hidden states and suppresses low-variance directions, but one fixed set of directions cannot follow a structure whose orientation changes from one region of the state space to another. At the action level, smoothing such as an exponentially weighted moving average damps jitter but does not correct the internal state.

DHS replays the demonstrations through the frozen policy and records the hidden states at one layer. The author found that this expert variation sits in few dimensions (29 principal components explain 90% of it on the Square task) but that its orientation differs between clusters. DHS therefore learns an invertible nonlinear coordinate map (a chart, built on score-based pullback Riemannian geometry) from those hidden states. At run time it maps the current hidden state into the chart, shrinks coordinates in which experts varied little more strongly than those in which they varied a lot, keeps every coordinate, and maps the result back before the remaining layers produce the action. One strength parameter controls how hard it steers.

On Robomimic Square, Transport and Can and MimicGen Coffee-D1, with the same frozen BC-Transformer per task and 300 episodes per cell (three seeds of 100), DHS raised noisy success from 39.7% to 51.7% on Square, 37.3% to 41.0% on Transport, 55.3% to 61.7% on Can and 67.0% to 73.0% on Coffee-D1. Among hidden-state methods (PCA projection, autoencoder reconstruction, nearest-neighbour replacement, a Gaussian score step) it had the best noisy score on Square, Can and Coffee-D1; on Transport the autoencoder scored higher, 43.3% against 41.0%. Combined with action smoothing it lifted Square to 61.7% from 57.7% and Transport to 62.0% from 56.3%. With the same 29 retained coordinates, the learned chart scored 48.7% under noise against 21.7% for PCA; with all 512 coordinates kept, DHS scored 50.3% against 44.0% for a PCA-based filter. During noisy Square runs the low-variance part of the state shrank from a median norm of 8.49 to 0.95 while the high-variance part barely moved (cosine similarity 0.9989). Success stayed between 48.0% and 53.3% across steering strengths from 0.1 to 1.0.

Everything is in simulation, with low-dimensional state inputs, one policy architecture and one intervention layer. DHS cut clean-observation success on Transport from 45.0% to 36.0%, and adding smoothing to DHS lowered noisy success on Can and Coffee-D1, so the steering strength needs per-task tuning; the strengths were chosen by hyperparameter search. The method models where expert states lie, not how they evolve over time, so a state can look expert-like and still be wrong for the current task phase. It is a single-author preprint with no hardware test.