How to Make Robots Learn from Demonstration Without Losing Their Freedom
The Problem Every Robot Controller Faces
Imagine you're teaching a robot to walk. You could let it figure out every step from scratch—learning through trial and error, optimizing for efficiency while avoiding falls. Or you could show it how a human walks and let it use that as a template, deviating only when necessary to accomplish specific tasks. The second approach seems intuitively smarter. We routinely learn motor skills by imitating others. But getting a controller to do the same thing mathematically? That's surprisingly hard.
The difficulty lies in a conceptual mismatch. Classical optimal control asks: what's the best action right now given the current state? It doesn't natively support the question: how can I stay close to a reference behavior while still accomplishing my goal? Meanwhile, machine learning approaches can certainly learn from demonstrations, but they often treat the reference as a black box, losing the mathematical guarantees that make classical control theory so powerful.
A new paper from Mintae Kim and Koushil Sreenath at UC Berkeley cuts through this tension with a clean mathematical framework called Trajectory-Regularized Stochastic Optimal Control (TRSOC). The core insight is deceptively simple: by measuring how far a controlled trajectory distribution diverges from a reference trajectory distribution, you can interpolate continuously between pure performance and pure imitation. The math connects to a classic tool in probability theory—Girsanov's theorem—allowing the authors to transform a global trajectory comparison into a local penalty that's computationally tractable.
The implications reach across robotics, autonomous vehicles, and any system where you want to leverage prior knowledge while retaining the flexibility to optimize for new objectives. "Large deviations from reference behaviors can lead to brittle policies," the authors note. TRSOC offers a principled way to stay within guardrails.
The Mathematics of Staying Close to a Template
To understand what TRSOC actually does, we need to peel back the layers of its mathematical construction. At its heart is a concept every graduate student in probability has encountered: the Kullback-Leibler divergence. KL divergence measures how different two probability distributions are. If you have a distribution P (your controlled trajectories) and a distribution Q (your reference trajectories), the KL divergence KL(P || Q) tells you how much information is lost when you use Q to approximate P.
In the TRSOC framework, the authors want to minimize the KL divergence between where the system actually goes and where a reference behavior says it should go. This seems straightforward until you try to actually compute it. Trajectory distributions live in an infinite-dimensional space—each trajectory is a complete path through time, and there are uncountably many such paths. Computing distances between distributions in this space sounds intractable.
Here is where Girsanov's theorem becomes the key that unlocks the problem. This theorem, originally developed for studying quantum mechanics, describes what happens to the probability of a path when you change the drift of a stochastic process. A drift is the deterministic part of how a system evolves—think of it as the average direction a system moves. When you perturb the drift, Girsanov's theorem tells you exactly how the probability of any given trajectory changes.
The crucial consequence, stated in Lemma 1 of the paper, is that the trajectory-level KL divergence between controlled and reference dynamics reduces to a simple quadratic penalty on the drift mismatch. Formally, if the controlled system has drift f(x, u) and the reference has drift f₀(x), the divergence becomes proportional to the squared norm of σ⁻¹(x)(f(x, u) − f₀(x)), where σ(x) is the diffusion matrix capturing stochastic noise. This is a local cost—evaluated at each state and action—rather than a global quantity over entire trajectories.
The requirement is that the diffusion matrix σ(x) must be invertible everywhere and satisfy a technical condition called Novikov's condition. This is satisfied in most practical settings where there's always some noise entering the system. Under these assumptions, the trajectory distribution of the controlled process is absolutely continuous with respect to the reference distribution, meaning the KL divergence is well-defined.
This transformation from trajectory-level to local-level regularization is the mathematical heart of the paper. It means that instead of comparing entire paths (computationally ruinous), you simply add a penalty whenever your control tries to move the system away from where the reference would have moved it. The penalty is quadratic, which preserves convexity and makes the optimization problem tractable.
Dynamic Programming Without Compromise
One of the deepest results in optimal control theory is the Dynamic Programming (DP) principle, which states that an optimal policy can be computed by working backward from the terminal objective. If you've studied reinforcement learning, you've encountered this as the idea that the value of a state depends on the values of future states. The mathematical incarnation of this principle is the Hamilton-Jacobi-Bellman (HJB) equation—a partial differential equation that encodes the optimal decision at every point in state space.
Here's the catch: naive trajectory regularization can break DP. If you're comparing entire trajectories, the optimal decision at any moment might depend on choices you make much later, because those later choices affect the trajectory shape. This destroys the Markov structure that makes DP tractable. The optimization becomes non-causal: you'd need to know the future to decide the present.
The authors prove that TRSOC avoids this pitfall. Because the trajectory KL reduces to a time-additive running cost, the regularized problem remains time-consistent. The DP principle applies unchanged. Theorem 1 in the paper establishes the finite-horizon DP equation, showing that the value function satisfies the usual decomposition: the cost from now to some future time, plus the value at that future time. This means you can still compute the optimal policy using backward induction.
Theorem 2 provides a verification result: if you can find a smooth solution to the HJB equation with the modified running cost, that solution equals the value function, and the minimizing control is optimal. The authors also extend the framework to infinite-horizon discounted problems, deriving the stationary HJB equation that characterizes optimal policies for problems without a fixed terminal time.
The KL regularization enters the HJB equation only through the running cost L(x, u), which is the original cost ℓ(x, u) plus the drift-mismatch penalty. The HJB equation takes the standard form:
with terminal condition V(T, x) = g(x). The only change from vanilla SOC is that the running cost now includes the regularization term. Everything else—the DP structure, the Markov property, the existence of optimal feedback laws—carries over unchanged.
This mathematical elegance is a feature, not just a technical detail. Engineers who have spent years developing tools for solving HJB equations can apply their existing solvers to TRSOC problems. The framework doesn't require new algorithms or approximations; it simply modifies the cost function.
The Trade-Off Between Performance and Fidelity
The regularization parameter λ ≥ 0 governs how strongly the controller is pulled toward the reference behavior. When λ = 0, the KL penalty vanishes entirely, and TRSOC reduces to standard stochastic optimal control—pure performance optimization. When λ is large, the penalty dominates, forcing the controller to stay close to the reference trajectory distribution.
Proposition 3 establishes a monotonicity property: as λ increases, the value function (the optimal cost) never decreases. This makes intuitive sense—adding constraints can't make the best achievable cost lower. More precisely, λ₁ ≤ λ₂ implies V^{λ₁}(x) ≤ V^{λ₂}(x) for all states x.
Proposition 4 shows what happens in the limit as λ → ∞. Under a technical assumption (that there exists a control matching the reference drift exactly), the optimal trajectories converge to the reference trajectories. The expected squared drift mismatch over any finite horizon shrinks like O(1/λ). In other words, with enough regularization weight, the controller can be forced arbitrarily close to the reference behavior.
The finite-λ regime is where things get interesting. The authors quantify the trade-off explicitly. Let Jₓ⁰(u) be the unregularized performance cost of a control policy u, and let Dₓ(u) be the expected discounted drift-mismatch. Then the regularized objective takes the form:
The optimal policy for a given λ balances these two terms. The authors derive a bound on performance degradation: if u⁰ is optimal for the unregularized problem and u^λ is optimal for TRSOC with parameter λ, then:
This inequality reveals that the cost of deviating from the reference in the unregularized solution bounds how much performance you lose by regularizing. If your optimal unregularized controller already stays close to the reference, the penalty has little effect. If it diverges wildly, you pay for that divergence.
The implications for multi-objective optimization are profound. The set of optimal solutions as λ varies traces a Pareto frontier between task performance and reference fidelity. For free, you get a complete parameterization of the trade-off curve. In practice, you might compute the frontier once, then select λ based on how much deviation you can tolerate for your specific application.
The Optimal Policy in Closed Form
For a broad class of practical problems, the authors derive an explicit formula for the optimal feedback law. The setting is control-affine dynamics with quadratic costs. The drift takes the form f(x, u) = f₀(x) + B(x)u, where f₀(x) is the reference drift and B(x) maps control inputs to state changes. The running cost is q(x) + (1/2)u^⊤R(x)u, where q(x) captures state-dependent penalties and R(x) is a positive-definite control weight matrix.
Under these assumptions, the optimal control (Proposition 5) is:
where the effective control weight is:
This is strikingly similar to the standard linear-quadratic regulator (LQR), except that the control weight matrix is augmented by the regularization term. The KL penalty effectively increases the cost of control actions, but in a way that depends on how much each action changes the drift relative to the diffusion. Controls that induce large drift changes in low-noise directions are penalized more heavily.
The closed-loop dynamics become:
The system still follows the reference drift f₀(x) but with a correction term that minimizes the value function gradient. As λ grows, the correction shrinks, and the system behavior converges to the reference dynamics.
For the linear-quadratic-Gaussian (LQG) case with constant matrices, the solution is even more explicit. The value function is quadratic: V(x) = x^⊤Px + c, where P satisfies a modified Riccati equation with the augmented control weight. The optimal feedback is linear: u* = −(R + λB^⊤Σ⁻¹B)⁻¹B^⊤Px, where Σ = σσ^⊤ is the diffusion covariance matrix.
Stability Without Sacrifice
A controller is useless if it drives the system to instability. The authors analyze closed-loop stability using Foster-Lyapunov theory, establishing sufficient conditions under which the optimal TRSOC policy stabilizes the reference dynamics.
The generator of the closed-loop diffusion captures how the value function evolves along trajectories. If you can find a Lyapunov function—some function of the state that's positive everywhere and decreases along trajectories—then the system is stable in an appropriate probabilistic sense. The authors show that under standard growth conditions on the drift and diffusion, the TRSOC solution inherits the stability properties of the reference dynamics, with the regularization providing an additional stabilizing effect through the increased control cost.
This matters for practical deployment. Many learned controllers behave well in simulation but fail catastrophically when deployed on physical systems due to distribution shift. By constraining the controller to stay close to a reference that is known to be stable, TRSOC provides a form of robustness that purely performance-based optimization cannot guarantee.
From Theory to Robots
The paper's experimental section demonstrates these concepts on a 2D navigation task and a 3D quadrotor hover scenario. The reference dynamics can come from three sources: passive dynamics (how the system naturally evolves), a nominal controller (a baseline policy you want to improve upon), or a learned model (trained on offline data from an expert demonstrator).
The figures from the paper illustrate the trade-off beautifully. In the known reference control setting (
a), each point on the curve represents the optimal solution for a particular λ. Moving along the curve from upper-left to lower-right, you increase λ, trading performance for fidelity. The curve is convex, reflecting the fundamental structure of the optimization.
The same pattern emerges with learned references (
a). Even when the reference dynamics come from a neural network trained on trajectory data, the framework seamlessly handles the mismatch. The learned reference captures the nominal behavior used in the data-driven TRSOC, as shown in
.
The approximation quality of the learned reference is remarkable. Figure 4 compares the learned acceleration function with the nominal acceleration along reference trajectories. "The close agreement shows that the learned drift captures the nominal behavior used in the data-driven TRSOC," the authors note. This validates the practical pipeline: collect demonstrations, fit a dynamics model, use that model as the reference, and apply TRSOC to find an improved controller.
The cost and deviation metrics as functions of λ (
and
) show how regularization smoothly interpolates between the two objectives. As λ increases from 0, the deviation from reference drops monotonically while the performance cost rises. The slope of this trade-off depends on the problem structure—some tasks allow near-optimal performance with moderate regularization, while others require staying very close to the reference.
Why This Matters for the Real World
The robotics literature is littered with failed attempts to combine learning and control. Pure reinforcement learning can discover impressive behaviors, but it requires enormous amounts of data and provides no guarantees about safety or stability. Classical control offers mathematical precision but assumes perfect knowledge of the system dynamics—knowledge that's rarely available in practice.
TRSOC offers a middle path. By treating a reference policy (or a learned model of expert behavior) as a soft constraint rather than a hard requirement, it preserves the flexibility of optimization while incorporating the structure encoded in prior knowledge. The controller doesn't need to reproduce the reference exactly; it can deviate wherever that deviation improves performance. But it cannot deviate arbitrarily far, because each departure from the reference trajectory is penalized.
This has immediate applications in several domains:
Robot manipulation. Imagine a robot that has learned to grasp objects from human demonstrations. The demonstration provides a reference trajectory distribution. Rather than asking the robot to reproduce the demonstration exactly—which might fail when object positions differ slightly from training—a TRSOC controller can use the demonstration as a template while optimizing for the specific grasp geometry it encounters.
Autonomous vehicles. A vehicle trained on human driving data could use the learned policy as a reference. TRSOC would allow it to deviate from human-like behavior when necessary (to avoid obstacles, for example) while maintaining a human-like driving style that passengers find comfortable.
Safety-critical systems. The regularization provides a form of interpretable constraint. If the reference dynamics are known to be safe, then the λ parameter directly controls how far the controller can drift from safe behavior. This is far more interpretable than ad-hoc penalty functions in reinforcement learning.
Sim-to-real transfer. A common failure mode in robotics is that controllers learned in simulation fail on real hardware due to model mismatch. By using simulation trajectories as a reference and regularizing toward them, TRSOC can bridge the sim-to-real gap. The controller stays close to the simulated behavior (which is known to work) while adapting to the real dynamics.
The Relationship to Existing Approaches
TRSOC occupies a unique position in the landscape of control and learning methods. The authors carefully distinguish it from related frameworks to clarify what makes it different.
KL control, for instance, optimizes directly over trajectory measures rather than over controls. The optimal solution is found by exponentially reweighting trajectories according to their cost, leading to a risk-sensitive objective. TRSOC, by contrast, operates in the standard SOC framework. It admits DP and HJB characterizations, and the optimal policy is a state feedback law. The key structural difference is that TRSOC remains risk-neutral (linear expectation) while geometrically penalizing drift deviation. KL control induces risk-sensitivity through its formulation.
Path-integral control, another related method, can linearize the HJB equation under specific structural assumptions—specifically, when control and diffusion enter through the same channel. This enables importance sampling-based solutions. TRSOC makes no such assumption; it works for arbitrary nonlinear dynamics, paying for this generality with a standard (nonlinear) HJB equation rather than a linear one.
The Schrödinger bridge problem, which seeks an entropy-minimizing interpolation between endpoint distributions, also involves KL divergence on trajectory space. But that problem is fundamentally different: it asks for the most likely trajectory connecting two states, given a reference process. TRSOC, by contrast, asks for the optimal control that trades off performance and reference proximity, with the control entering through the drift.
The common thread through all these related approaches is KL divergence as a regularizer. But the specific way TRSOC deploys it—through Girsanov's theorem, reducing to a local drift penalty, preserving DP—differentiates it from the alternatives.
Caveats and Open Questions
No mathematical framework is without limitations, and TRSOC has its own.
The most significant assumption is that the diffusion matrix σ(x) must be invertible everywhere. In continuous-time stochastic systems, this means there must always be some noise in every direction of the state space. While common in theory (it ensures absolute continuity of trajectory measures), many real systems have degenerate noise—some state components evolve deterministically. Extending TRSOC to such cases is an open problem.
The reference dynamics must be specified explicitly. In many applications, you don't have a closed-form reference policy; you have trajectories from an expert or a simulator. The paper shows that a learned dynamics model can serve as the reference, but this introduces approximation error. How this error propagates through the framework is not fully characterized.
The optimal policy requires solving the HJB equation, which is generally intractable for high-dimensional systems. While the paper provides closed-form solutions for the LQ case, practical deployment in robotic systems with many degrees of freedom will require approximation methods. Model predictive control, neural network approximations, or other numerical schemes could bridge this gap, but the theory doesn't yet guarantee their performance.
Finally, the theoretical results assume infinite data and exact knowledge of the dynamics. In reality, the reference distribution is estimated from finite data, and the dynamics model is imperfect. How robust is TRSOC to these uncertainties? The experiments suggest reasonable performance, but a rigorous robustness analysis remains for future work.
The Path Forward
What makes TRSOC particularly promising is its generality. The framework doesn't require the reference dynamics to be optimal for anything. The reference could encode safety constraints, human preferences, energy efficiency, or any other prior knowledge. The controller will use it as a template while optimizing for the specified task cost.
The natural next step is to apply TRSOC to larger-scale problems. The 2D and 3D examples in the paper are proofs of concept; real-world deployment would involve complex manipulators, autonomous vehicles, or even chemical processes. Whether existing numerical methods for HJB equations scale to these domains is an open empirical question.
Another direction is adaptive reference learning. Currently, the reference is fixed. But in many applications, you might want to update the reference as you gather more data about the system. This would require a coupled optimization over the reference model and the controller, which is computationally challenging but potentially powerful.
The relationship to offline reinforcement learning is worth exploring further. The TRSOC framework suggests a new objective function for offline RL: minimize the task cost plus a penalty on deviation from a learned reference distribution. This is different from existing offline RL methods, which typically penalize out-of-distribution actions. The TRSOC perspective emphasizes trajectory-level deviation rather than action-level deviation, which might provide better stability and sample efficiency.
Finally, there are theoretical connections to explore. The TRSOC objective is closely related to certain variational formulations in statistical physics and information theory. The authors hint at these connections but don't pursue them. Understanding TRSOC through the lens of these related fields might yield new algorithmic insights.
A Unifying Principle
At its core, TRSOC is about leveraging structure without sacrificing optimality. Every engineer knows that prior knowledge is valuable—it reduces the space of policies you need to search, focuses learning on the relevant regions of state space, and provides safety margins. But incorporating prior knowledge into optimal control has traditionally required either ignoring it (pure optimization) or ceding control to it (pure imitation).
TRSOC offers a third way. By measuring deviation from a reference at the level of trajectory distributions, it captures what matters: not just whether you're taking similar actions, but whether your entire trajectory distribution matches what you'd expect from following the reference. This is a more meaningful measure of similarity than action-space penalties, and it's computationally tractable.
The mathematical machinery—Girsanov's theorem, KL divergence, HJB equations—is classical. What's new is the synthesis: bringing these pieces together to create a framework that is both theoretically clean and practically useful.
As robotic systems become more capable and as we ask them to operate in increasingly complex environments, frameworks like TRSOC will become more important. We'll want them to learn from human demonstrations, from simulation, from experience—but we'll also want to constrain how far they can drift from behaviors we trust. TRSOC provides a principled way to express this trade-off, and it does so in a language that control theorists already speak.
The framework won't solve every problem. For highly nonlinear, high-dimensional systems, solving the HJB equation remains challenging. For systems with deterministic dynamics (no noise), the current formulation doesn't apply. And the choice of λ—how much to regularize—requires domain knowledge or experimentation.
But for the problems it does address, TRSOC provides something valuable: a bridge between learning and control, between performance and fidelity, between what we know and what we want to optimize. In a field often divided between theory-averse learning and实践-averse control theory, that's a welcome synthesis.