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

arXiv:2610.21666 · cs.RO · Submitted 2026-10-15 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: Robotics Radio. Generated commentary on the latest robotics and control papers.

Rosa: Today's paper: "Diffeomorphic Hidden-State Steering".

Dev: The gist The Diffeomorphic Hidden-State Steering method uses expert hidden-state geometry to improve policy robustness at test time by applying variance-dependent rescaling in nonlinear coordinates to produce location-dependent steering directions.

Rosa: First, who's behind it and why it matters.

Title and authors: Rosa: So, we're talking about this paper today called "Diffeomorphic Hidden-State Steering: Improving Imitation Learning Policy Robustness through Expert Hidden-State Geometry". Rosa and Dev have been looking at it, and they want to walk us through what this actually means for how we make these AI policies more reliable.

Dev: Exactly. The core idea here is that small errors in observation or action can pile up during closed-loop execution, which messes up a trained imitation learning policy. This paper investigates whether we can use the expert demonstrations themselves to improve robustness at test time, without having to retrain the policy or collect new data.

Taro: So, it’s about taking those expert hidden states and using them as a guide to steer the current hidden state in a way that makes the policy more stable when things get noisy. That sounds like something we need when we're testing robots outside the lab or in tricky real-world scenarios.

Rosa: Right. The paper finds that expert hidden state variation isn't spread out randomly; instead, it’s concentrated in relatively few dimensions, and the directions capturing that variation change depending on which region of the task you're looking at.

Dev: That concentration is key because it motivates what they call Diffeomorphic Hidden State Steering, or DHS. This method learns an invertible nonlinear coordinate map from those expert hidden states by replaying demonstrations through a frozen policy.

Taro: So, it’s not just taking the expert states and moving them around; they are learning a specific mathematical way to map those states so that low-variance coordinates get attenuated much more strongly than high-variance ones.

Rosa: That attenuation is done in nonlinear coordinates, and then they map everything back to the original hidden state space. This clever setup means the steering directions change depending on where you are in the hidden state space right now.

Dev: And it says this allows them to retain every coordinate while still getting these location-dependent updates, which is pretty powerful because it’s not just a simple linear adjustment.

Title and authors: Taro: I'm thinking about the math they use, because if you're doing real-time control, you need to know how this map works operationally. They fit a diffeomorphic chart and an anisotropic Gaussian model of expert variation into their coordinates.

Rosa: And then they define the steering operation by finding an optimal coordinate modification that minimizes a combined penalty function. The unique minimizer gives them the formula z plus equals W beta z t, where W beta is this diagonal matrix with ai over ai plus beta for each coordinate i.

Dev: And when you set steer strength beta greater than zero, they find that coordinates with ai much larger than beta undergo small relative changes, but those with ai much smaller are strongly attenuated. This means every multiplier is strictly positive for any finite beta, so all coordinates are retained and the map T t beta is invertible in exact arithmetic.

Taro: It’s interesting how they manage to make sure that even when you're attenuating things heavily, you don't lose any part of the state representation by keeping every coordinate. That’s a tight constraint for any steering mechanism.

Rosa: The empirical results show this method works across four simulated manipulation tasks. DHS improves the mean success under observation noise over the unmodified Transformer policies they tested on these tasks, which is pretty significant for deployment.

Dev: Controlled comparisons showed that using these learned nonlinear coordinates gives higher mean success than using linear coordinates, even when they use both hard projection and full-rank rescaling methods. That suggests the choice of coordinate system really matters for task success.

Taro: I saw something in the results that was interesting about selective attenuation control. They found that DHS strongly attenuates the low-variance block while only making limited changes to the high-variance block during noisy Square roll-outs.

Rosa: And they tested five different strengths of beta, from zero point one up to one point zero, and their mean success ranged from forty-eight point zero percent to fifty-three point three percent, which is better than the no steering reference which was at only thirty-nine point seven percent.

Title and authors: Dev: But here’s a caveat they mentioned: increasing beta strengthens the attenuation in every chart coordinate, including those high-variance ones, and it can actually impair execution even as it reduces the low-variance chart norm.

Taro: So, there is a clear trade-off there between how much you try to dampen noise and how much you affect the coordinates that are already important for the task. That’s a real constraint for any practical system.

Rosa: The conclusion of this paper is that DHS uses expert hidden state geometry to improve the robustness of frozen imitation learning policies at test time, and it serves as a useful reference for selective test-time steering without needing extra demonstrations or policy retraining.

Dev: It also points out that the expert hidden state density doesn't uniquely determine the update, and they suggest combining this expert geometry with predictions of an intervention’s future consequences.

Taro: The limitation they flag is that the fitted geometry depends on preprocessing and chart regularization, and their coordinate statistics are used as a steering reference rather than a direct model of task error.

Rosa: So, to wrap up the idea behind "Diffeomorphic Hidden-State Steering: Improving Imitation Learning Policy Robustness through Expert Hidden-State Geometry," it’s about using the structure found in expert hidden states to guide test-time corrections for imitation learning models.

Dev: It’s a method that allows you to tailor how much you correct different parts of the policy's internal state based on the geometry derived from those expert demonstrations.

Taro: For someone who just listens to this show, the implication is that we can potentially make policies more resilient to noise in real-world robotics by leveraging what we already have in our expert data.

Rosa: Right. So, this gives us a better tool for making frozen policies stronger when they encounter unexpected stuff during execution. Next up, we’ll look at how other recent papers are handling similar problems and where this fits into the bigger picture of autonomous systems.

The paper's summary: Rosa: So, we're looking at how this Diffeomorphic Hidden-State Steering method works as a whole now. It’s basically taking the structure of those expert hidden states we saw earlier and using it to fix the policy when it gets noisy during real operation.

Dev: Exactly. Think of it like having an expert blueprint for how a policy *should* be thinking at different moments, and you use that blueprint to subtly nudge the current thinking process away from bad noise, without having to retrain anything.

Taro: So, the main thing is this learned map from those expert states—the DHS—it lets you selectively attenuate certain parts of the hidden state based on how much variation they have in the expert data.

Rosa: That’s right. It doesn't just blindly cut everything that looks different; it specifically targets coordinates that have low variance in the expert data, making sure those stable parts are handled gently, while strongly damping the ones that are wildly fluctuating.

Dev: And this happens all within a nonlinear coordinate system—a diffeomorphism—which is what lets you get these location-dependent steering directions instead of just applying a fixed correction everywhere.

Taro: So, when the robot is in one state configuration, say near a corner, the steering direction changes because the underlying geometry of those expert states is different there. It adapts to where you are in the hidden state space.

Rosa: And that adaptation means you get these steering directions that vary with position while still retaining every single coordinate from your original representation. It’s like getting a highly customized correction tailored precisely to the current situation.

Dev: The empirical results across those manipulation tasks show this actually helps the policy succeed under observation noise more than it does without any steering at all, and even better than using simple linear corrections or PCA.

Taro: That's a big deal because it shows that picking your coordinates is important—the nonlinear ones learned by DHS actually beat the linear ones in terms of task success, even when you’re just keeping the same number of dimensions.

Rosa: So, this means for someone out there building a robot, you can use what your expert demonstrations teach you about how those hidden states behave to build a much more robust policy that handles real-world noise better at test time.

Dev: But we have to watch the trade-off they found—stronger attenuation always affects every coordinate, even the ones with high variance, so there's a limit on how much you can aggressively dampen the noise before you start messing up the intended action.

Taro: That makes sense. It’s a balancing act between cleaning up the state and making sure you don't destroy the actual maneuver being attempted.

Rosa: So, while this is cool for improving robustness, there are limitations they point out—the geometry they fit depends on how you prepare your data and regularize your chart, so it’s not just magic; it’s tied to that setup.

Dev: And the authors themselves say that the expert hidden state density isn't enough on its own to tell you exactly what update to apply, which leads them toward combining this geometry with predictions about what might happen later in the mission.

Taro: That future work sounds interesting because it suggests tying this steering mechanism not just to current noise, but also to predicting the consequences of an intervention down the line.

Rosa: Yeah, that’s where things could go next—moving from a test-time fix based on expert structure to a more predictive system that anticipates failure and steers accordingly.

Dev: So, what we're seeing here is using existing knowledge from demonstrations to create a dynamic steering mechanism that adapts to the robot's current internal state geometry.

Taro: It’s about leveraging the spatial organization of expert hidden states to guide test-time corrections for imitation learning policies in real manipulation tasks.

The paper's improvements: Tom: So, we're talking about what makes this Diffeomorphic Hidden-State Steering method actually better than just slapping some simple noise filter on top of an imitation learning policy. It's about using geometry to make the correction surgical and smart, not just blunt.

Rosa: Exactly. The paper suggests a few key improvements over standard methods, mainly focusing on how they select which part of the hidden state gets modified based on its statistical properties in the expert data.

Dev: Specifically, they move beyond simple linear projections and show that using these learned nonlinear coordinates gives higher task success than PCA under both hard projection and full-rank rescaling.

Taro: So, if you look at it from a failure mode perspective, this means the policy isn't just being nudged randomly; it’s being steered along directions that are specifically informed by where the expert states are concentrated in that high-dimensional space.

Rosa: That’s right. It implies that we can design better steering mechanisms because we aren't treating all dimensions equally; we know which ones are important and which ones are just noisy noise in the context of the expert behavior.

Dev: And they show this selection process is very precise; they achieve location-dependent steering directions, meaning the correction changes depending on where you are in the hidden state space right now.

Taro: That level of adaptation is crucial for real-world recovery; it suggests that when something goes wrong, the system doesn't just apply a generic fix but applies a fix tailored to its specific current situation.

Rosa: And they also found this method works better than just action smoothing when you combine it with some simple action smoothing techniques in certain tasks, which is a good practical improvement for deployment.

Dev: But the authors are also very careful about the control aspect; they show how varying the steering strength, that beta value, lets you precisely control how much of that attenuation affects your actual actions.

Taro: That control over the trade-off between noise reduction and action stability is important because it gives engineers a way to tune the system for different operational requirements.

Rosa: So, the implication here is that we can build more resilient robots by using expert geometry to create test-time steering that adapts dynamically to where they are in their internal state.

Dev: It’s about making the correction smarter by understanding the underlying structure of what a good policy looks like based on expert behavior.

Taro: This moves us toward systems that don't just recover from errors, but recover in a way that respects the underlying dynamics of the task.

Conclusion: Rosa: So we're wrapping up on "Diffeomorphic Hidden-State Steering: Improving Imitation Learning Policy Robustness through Expert Hidden-State Geometry." Basically, we saw how using the structure of expert hidden states allows us to build a steering map that fixes policy errors dynamically at test time.

Dev: Yeah, the main takeaway is that this method gives us a powerful tool to improve closed-loop robustness without needing to retrain the imitation learning policy itself.

Taro: It means we can make policies more resilient when they encounter unexpected things outside of what they saw during training, which is a big deal for real-world deployment.

Rosa: I think it really shows how much we can learn from our expert demonstrations and use that knowledge to guide the AI's behavior when it’s operating in the field.

Dev: And while the results are strong across those four simulation tasks, there are limitations they mentioned about how dependent the fitted geometry is on their data preprocessing and chart regularization methods.

Taro: That dependence is important; it means this steering reference isn't a perfect model of task error itself, but rather a guide based on how they chose to represent that error.

Rosa: Exactly. So, if you’re building something, remember that the success of this method hinges on your initial setup and how you structure those expert hidden states.

Dev: It's a solid way to tackle test-time failures in real-time loops while keeping the latency manageable, provided you can handle that geometric fitting upfront.

Taro: For me, it shows that we don't need perfect models of the world; we can use structural information from demonstrations to guide recovery, which is more practical for autonomy researchers.

Rosa: Well said. So this paper provides a solid reference point for test-time steering using expert geometry to boost policy robustness against observation noise.

Dev: It’s a useful piece of work showing how we can leverage existing expert knowledge to make frozen policies more reliable under real-world conditions.

Taro: Moving on, next time we'll look at papers dealing with agentic systems and mission-state governance for industrial multi-robot setups.

KwangBin Lee

KAIST

cs.RO

Submitted: 2026-10-15

Updated: 2026-10-15

Comments: 9 pages, 2 figures

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 88/100

The gist: The gist The Diffeomorphic Hidden-State Steering method uses expert hidden-state geometry to improve policy robustness at test time by applying variance-dependent rescaling in nonlinear coordinates

Key concepts

Expert Hidden State Geometry: Analysis shows that expert hidden states are concentrated near a low-dimensional manifold. Variation along this manifold differs across regions, suggesting that the expert distribution is structured in a specific, non-uniform way. This geometry is used as a reference to guide how the policy should be steered.
Diffeomorphic Hidden State Steering (DHS): DHS learns an invertible nonlinear coordinate map from expert hidden states by replaying demonstrations through the frozen policy. It then modifies these coordinates based on their variance: low-variance coordinates are attenuated more strongly, while high-variance ones are changed less. Finally, it maps these modified coordinates back to the original hidden state space.
Steering by Variance-Dependent Rescaling: This operation defines steering strength using a specific mathematical formula where coordinates with low expert variance are strongly reduced (attenuated) while those with high variance undergo smaller relative changes. This ensures that every coordinate is retained and the transformation remains invertible, allowing the steering direction to vary intelligently based on the current hidden state.

Terminology

Steering — 4× in this paper · explained in 2 episodes
Diffeomorphic Hidden-State Steering — 2× in this paper
Robustness — 2× in this paper · explained in 6 episodes
Transformer — 1× in this paper · explained in 2 episodes

Summary

The gist The Diffeomorphic Hidden-State Steering method uses expert hidden-state geometry to improve policy robustness at test time by applying variance-dependent rescaling in nonlinear coordinates to produce location-dependent steering directions.

Characterization of Expert Hidden State Geometry

The analysis of a trained policy finds that expert hidden state variation is concentrated in relatively few dimensions, while the directions capturing this variation differ across regions Our empirical analysis finds that expert hidden states vary predominantly within relatively few dimensions, while linear subspaces fitted to different clusters of expert hidden-states differ in orientation (Section IV-A). This observations motivate viewing the expert distribution as concentrated near a low-dimensional manifold, with broad variation along local tangent subspaces and narrow variation along local normal subspaces

The Diffeomorphic Hidden State Steering (DHS) Mechanism

DHS is proposed to learn an invertible nonlinear coordinate map from expert hidden states obtained by replaying the demonstrations through the frozen policy DHS attenuates coordinates with low expert variance more strongly than those with high variance, then maps the modified coordinates back to hidden-state space This allows steering directions to vary with the current hidden state while retaining every coordinate The process involves fitting a diffeomorphic chart from expert hidden states and an anisotropic Gaussian model of expert variation in its coordinates

Steering by Variance-Dependent Rescaling

The steering operation is defined by finding the optimal coordinate modification to minimize a combined penalty The unique minimizer gives z+t = Wβzt, where Wβ = diag(ai / (ai + β) for i=1 to d), and Tβ(h) = Φ−1(WβΦ(h)) For steering strength β > 0, coordinates with ai ≫ β undergo small relative changes, while those with ai ≪ β are strongly attenuated This means that every multiplier is strictly positive for finite β, so all coordinates are retained and Tβ is invertible in exact arithmetic

Empirical Results and Performance

Across four simulated manipulation tasks, DHS improves mean success under observation noise over the unmodified Transformer policies Controlled comparisons yield higher mean success with learned nonlinear coordinates than with linear coordinates under both hard projection and full-rank rescaling DHS achieves higher noisy means than the Gaussian score step on all four tasks and the highest noisy mean among the hidden-state methods on Square, Can, and Coffee-D1 The learned chart yields higher mean success than PCA in both clean and noisy conditions (Table II), showing that the choice of coordinates affects task success even when the same number of coordinates is retained

Selective Attenuation and Steering Strength Tradeoff

The analysis shows that DHS strongly attenuates the low-variance block while making limited changes to the high-variance block during noisy Square roll-outs Across the five tested strengths from β = 0.1 to 1.0, mean success ranges from 48.0% to 53.3%, above the 39.7% no steering reference in Table I Increasing β strengthens attenuation in every chart coordinate, including high-variance coordinates (Eq. (5)), showing that stronger attenuation can impair execution despite further reducing the low-variance chart norm

Conclusion and Limitations

DHS uses expert hidden-state geometry to improve the robustness of frozen IL policies at test time The expert hidden state density does not uniquely determine the update, and a natural extension involves combining expert geometry with predictions of an intervention’s future consequences [12], [13] Future work could also combine DHS with out-of-distribution (OOD) detection to adapt when and how strongly steering should be applied based on departures from the expert hidden-state distribution. The evaluation covers fixed BC-Transformer checkpoints with low-dimensional observations in simulation, and a single intervention site. The fitted geometry depends on preprocessing and chart regularization; its coordinate statistics provide a steering reference rather than a model of task error. The paper demonstrates that expert hidden states can serve as a useful reference for selective test-time steering without additional demonstrations or policy retraining. This work originated from research conducted at KAIST and the author thanks Prof. SooJean Han and the Autonomous Control for Stochastic Systems (ACSS) Lab at KAIST for their support and the research environment in which this work began. The OpenAI’s ChatGPT was used to assist with language polishing, and codes for implementation of the experiments. The author takes full responsibility for all statements and final formulations.

--- Page 1 ---

Diffeomorphic Hidden-State Steering: Improving Imitation Learning Policy Robustness through Expert Hidden-State Geometry Abstract— Small observation or action errors can compound during closed-loop execution and disrupt a trained imitation learning policy. We study whether the original expert demonstrations can improve robustness through test-time hidden-state steering without additional demonstrations or policy retraining. Our analysis of a trained policy finds that expert hidden-state variation is concentrated in relatively few dimensions, while the directions capturing this variation differ across regions Motivated by this geometry, we propose Diffeomorphic HiddenState Steering (DHS), which learns an invertible nonlinear coordinate map from expert hidden states obtained by replaying the demonstrations through the frozen policy DHS attenuates coordinates with low expert variance more strongly than those with high variance, then maps the modified coordinates back to hidden-state space Across four simulated manipulation tasks, DHS improves mean success under observation noise over the unmodified Transformer policies Controlled comparisons yield higher mean success with learned nonlinear coordinates than with linear coordinates under both hard projection and full-rank rescaling Combining DHS with action smoothing improves mean success rate under observation noise in certain tasks compared with action smoothing alone

--- Page 3 ---

III. PROBLEM FORMULATION Let DE denote the expert demonstrations used to train an IL policy πθ. At control step t, the policy receives an observation history Ot and generates an action using its original action selection rule. We select a single intermediate layer l and denote the hidden-state to be modified at control step t by ht ∈ R d,where d is its dimension. A steering map T: R d → R d replaces this hidden state with h+t = T(ht). (1) At the selected layer, all other token representations remain unchanged. The policy then resumes its original computation and generates an action using its original action selection rule. We denote the policy with this intervention by πθ,T. Before deployment, we replay observation histories from DE through the frozen policy and collect hidden states at the same layer and token position. These expert hidden states form the dataset HE, which provides the fitting reference for T. The objective is for πθ,T to achieve higher task success than πθ under observation noise.

--- Page 4 ---

IV. DIFFEOMORPHIC HIDDEN-STATE STEERING A. Empirical Geometry of Expert Hidden States Studies of image classifiers and Transformers have reported evidence of low-dimensional manifold structure in hidden representations [8], [9]. Motivated by these findings, we examine whether the expert hidden-state distribution from an IL robot policy is concentrated in relatively few dimensions and whether regional linear subspaces differ in orientation We select the Square manipulation task from Robomimic [15] as a representative task and use a Transformer policy trained by behavior cloning (BC-Transformer).

Improvements for AI systems

Expert Hidden-State Geometry Characterization: The system can characterize expert hidden-state geometry in a trained IL policy, observing concentrated variation and differences in the orientation of linear subspaces fitted to different clusters of expert hidden-states. This allows for a deep understanding of how expert behavior is organized within the policy's internal representations.
Diffeomorphic Hidden-State Steering (DHS) Implementation: The system can implement DHS, which learn[s] an invertible nonlinear coordinate map from expert hidden states obtained by replaying the demonstrations through the frozen policy, enabling it to attenuate coordinates with low expert variance more strongly than those with high variance. This results in a steering mechanism that allows steering directions to vary with the current hidden state while retaining every coordinate.
Improved Closed-Loop Robustness: The improved AI system can achieve higher task success under observation noise, as DHS improves mean success under observation noise over the unmodified Transformer policies across simulated manipulation tasks. Specifically, it can yield gains of 3.7 to 12.0 percentage points in noisy success compared to no steering on certain tasks like Square and Transport when combined with Action EWMA.
Coordinate Selection Superiority: The system can select learned nonlinear coordinates over linear ones, as the paper demonstrates that learned nonlinear coordinates yield higher mean success than linear coordinates from principal component analysis (PCA) under both hard projection and full-rank rescaling. This proves that the choice of coordinates affects task success even when the same number of coordinates is retained.
Selective Attenuation Control: The system can precisely control the intervention effect by varying the steering strength, as shown by Table III where varying β on Square at noise level σ = 0.05 shows that stronger attenuation can impair execution despite further reducing the low-variance chart norm. This highlights a crucial trade-off between perturbation attenuation and changes to coordinates used for action generation.
Location-Dependent Steering Directions: The system can utilize location-dependent steering directions, as the nonlinear chart allows the same coordinate rescaling to act along directions that vary with position, effectively capturing local geometry changes in hidden state space.