FOCUS: First-Order Control Using Optimal Transport Supervision for Imitation from Observations

AI summary

Overview

Research area: Imitation learning from observations (ILfO) combined with first-order model-based reinforcement learning and optimal transport (OT) — specifically, learning control policies from state-only demonstrations inside differentiable physics simulators.

Technical level: Advanced. The paper assumes familiarity with Markov decision processes, entropy-regularized optimal transport / Sinkhorn algorithms, backpropagation through time, and actor–critic methods (SHAC, AHAC, SAPO).

Scope: A method paper introducing FOCUS, a training framework that converts temporally unaligned state-only demonstrations into short-horizon differentiable rewards via local optimal-transport matching with soft best-of-K expert-segment selection, evaluated on four contact-rich DFlex locomotion tasks (arXiv:2610.41540v1, workshop "Physical Understanding for Decision-Making: Bridging Foundation Models and Reliable Agents", CC BY 4.0).

What This Paper Is About

Imitation from observations gives a learner only sequences of expert states — no expert actions, and no indication of which moment in the demonstration corresponds to the learner's current moment. The learner must simultaneously figure out what to imitate and where it is in the demonstration.

Existing OT-based imitation methods typically turn transport costs into surrogate rewards and then optimize with zeroth-order reinforcement learning, which is sample-inefficient and high-variance; adversarial ILfO methods depend on discriminator stability and tuning. Differentiable simulators offer a third route (pathwise gradients), but naively differentiating an imitation loss through long, contact-rich trajectories is brittle. FOCUS's goal is to make that third route work: combine local OT alignment with short differentiable rollouts plus critic bootstrapping, so that unstable long-horizon gradients are avoided while credit still propagates beyond the matching window.

Key Contributions

1. A local OT + best-of-K supervision scheme. Each short agent rollout window is matched against K sampled expert segments, and a softmin (log-sum-exp) aggregation with temperature β (fixed at 0.5 in experiments) weights locally compatible segments more heavily, producing differentiable imitation rewards that de-emphasize phase-shifted or unreachable expert targets.

2. Coupling local transport costs to short-horizon differentiable actor–critic learning. Transport-based per-step costs are turned into rewards and combined with a learned critic that bootstraps value beyond the H-step rollout, following the SHAC/SAPO template but with an imitation signal derived from state-only demonstrations rather than an environment reward.

3. A reward-shaping design space over OT costs. Three monotone decreasing transformations of the aggregated cost are proposed — raw negative cost (−ℒ), exponential (exp(−ℒ)), and logarithmic (−log(1 − e^(−ℒ) + ε), with ε = 10⁻⁷) — to control reward scale and curvature for actor and critic learning.

4. An ILfO setting with no alignment aids and no expert actions. FOCUS requires neither temporally aligned trajectories, nor expert-aligned initialization, nor teacher forcing, nor behavior-cloning pretraining, and the authors report it remains effective when only a single expert trajectory is available.

Main Findings

• Best final return in every environment. Across Hopper, Ant, Humanoid and SNU Humanoid, at least one FOCUS OT variant attains the highest mean return at the final interaction budget among the compared methods (aggregate interquartile mean of expert-normalized return in Fig. 1; per-task curves in Fig. 4 over six seeds with bootstrap 95% confidence intervals).

• Strong result on the hardest task. On the SNU Humanoid — a contact-rich musculoskeletal model controlled by 152 muscle activations — the OT variants approach expert performance and attain substantially higher mean returns than baselines within both interaction and wall-clock budgets; all zeroth-order baselines remain far below expert performance.

• OT beats direct local matching. At the interaction-budget endpoint, both OT variants (FOCUS-OT-L2 and FOCUS-OT-Cos) attain higher mean returns than FOCUS-L2 (direct local ℓ2 matching with the same best-of-K and actor–critic machinery) in all four environments.

• Wall-clock advantages on some tasks. On Hopper all FOCUS variants and OOPS approach expert performance within the interaction budget, but FOCUS learns substantially faster in wall-clock time. On Ant, OOPS is competitive with the OT variants in environment interactions with overlapping endpoint confidence intervals but learns much more slowly in wall-clock time, and ILD attains the highest final wall-clock mean. On Humanoid, FOCUS-OT-Cos attains the highest mean return under both budgets, followed by FOCUS-OT-L2 and ILD, with overlapping wall-clock confidence intervals.

• The two OT configurations are environment-dependent. FOCUS-OT-Cos leads on Ant and Humanoid; the authors state the relative strengths cannot be attributed solely to the distance metric (ℓ2 vs. cosine ground cost) because reward shaping also differs.

• Critic bootstrapping helps. Including the critic improves mean return on all four tasks for FOCUS-OT-Cos and on three of four for FOCUS-OT-L2; the exception is Ant, where intervals overlap. Both configurations benefit substantially on the humanoid tasks.

• Reward-shaping effects differ by ground cost. With a critic, exponential shaping gives higher mean returns than logarithmic shaping for the L2 cost on three of four tasks, while logarithmic shaping gives higher mean returns for the cosine cost on all four.

• First-order training beats black-box optimization of the same reward. ZO-FOCUS variants (SAC or PPO optimizing the same imitation reward without differentiating through the simulator) attain lower mean returns than the corresponding first-order OT variants on all four tasks at matched interaction budgets; the authors note this comparison also changes the full actor–critic algorithm, so it does not isolate simulator derivatives alone.

• Best-of-K mitigates phase error in a controlled diagnostic. In a paired diagnostic on held-out expert windows at a half-gait-cycle mismatch, adding seven random alternatives yields a phase-compatible minimum-OT-cost selection in 46–72% of query–pool pairs and reduces mean phase error by approximately 3–5× across four environments.

• Backbone and data ablations favor SAPO over SHAC, and FOCUS remains effective with one expert trajectory.

Methodology in Plain English

FOCUS assumes a known, differentiable simulator: you can push actions in and get gradients back through the physics.

At each training iteration the agent is rolled out for a short horizon H (H much smaller than the full episode length) inside the simulator. This produces a small window of agent states/transitions. Separately, the algorithm samples K short segments of length L from the expert demonstrations — random start times, so no temporal correspondence is needed.

For each of those K expert segments, the method computes an entropy-regularized optimal transport plan (via a fixed number of Sinkhorn iterations) between the agent window and the expert segment, using either an ℓ2 or a cosine ground cost on state–state pairs normalized by expert-data statistics. This yields both a per-segment scalar loss and per-timestep costs for the agent states.

The per-segment losses are then converted into softmin weights, so segments that are locally compatible with the agent's current phase get more weight and badly phase-shifted segments get less. The per-step costs are aggregated with those weights (gradients flow through the weights too) and transformed into rewards. The transformed reward drives a standard bootstrapped short-horizon actor–critic objective — discounted reward over the H rollout steps plus a learned value estimate at the final state — with the main experiments using SAPO as the backbone (SHAC in ablations). Policy and critic are updated by backpropagating through the simulator, the Sinkhorn iterations, and the weighting.

No expert actions, no discriminator, no inverse-dynamics model, no pretraining, and no alignment tricks are used. Evaluation uses expert-normalized return R_norm = R_π / R_E, where the environment's task reward is used only for evaluation, never for training; demonstrations come from 8 state-only expert trajectories generated by SAPO-trained policies.

Why This Matters

Impact on research. The paper argues that simulator derivatives are a viable path for high-dimensional state-only imitation, not just for reward-driven control. It bridges two literatures that had been mostly separate — OT-based imitation (PWIL, SIL, OOPS, LWAIL) and first-order model-based RL (SHAC, AHAC, SAPO) — and shows that the alignment ambiguity of ILfO can be handled locally rather than globally. It also provides an ILfO comparison point for differentiable imitation baselines such as ILD (Chamfer-α), and explains why DiffMimic was excluded (it relies on expert-aligned initialization and teacher forcing).

Real-world applications:

• Musculoskeletal motor control and rehabilitation: the paper's own motivation comes from a Swiss National Science Foundation project on modelling pathological gait for treatment selection support (METATHESIS, grant no. 220435); the 152-muscle SNU Humanoid is a proxy for learning human movement from motion-capture-like state data.
• Robot learning from human demonstrations where only motion/pose data is captured and joint torques or actions are unavailable.
• Character animation and physically simulated agents, where reference motions supply states but not the control signals that produced them.
• Sim-to-real controller design for legged and humanoid locomotion, where demonstrations come from reference gaits or prior controllers.

Industry relevance. Any pipeline that already has a differentiable simulator (Brax, Warp, DFlex) but lacks a well-specified reward function can, in principle, substitute expert state trajectories for reward engineering. That is relevant to robotics simulation teams, game/animation studios using physics engines, and biomechanics or digital-human companies. Because the method avoids adversarial training, it sidesteps a notoriously tuning-sensitive component.

Future Directions

• Extending to learned world models. FOCUS assumes a known differentiable simulator; the authors leave to future work the case where dynamics are learned, which would inject model-Jacobian error into the policy update.

• Broadening the task suite. Validation is restricted to periodic locomotion in DFlex; manipulation, deformable-object control, sparse-goal tasks, non-periodic tasks, and real-world transfer are all open.

• Making alignment global rather than local. The best-of-K objective is local and can underrepresent rare expert phases or favor fragmentary matches when candidate segments are unrepresentative; the authors call for alignment mechanisms that enforce long-horizon consistency.

• Understanding the composite design choices. The two OT variants differ in both distance metric and reward transformation, so the paper explicitly does not isolate either factor; disentangling ground cost, reward shaping, K, and horizon remains open.

Target Audience

Researchers and graduate students working on imitation learning, model-based RL, or optimal transport for control; practitioners with access to differentiable simulators who want to replace reward engineering with state-only demonstrations; and roboticists or biomechanics researchers interested in high-dimensional contact-rich control (for example, musculoskeletal humanoids). Readers without background in OT and differentiable simulation will find the method sections demanding, though the problem framing and results are accessible.

Authors’ abstract

Imitation learning from observations must recover behavior from state-only demonstrations without known temporal correspondence. Differentiable simulators provide pathwise policy gradients, but differentiating through long, contact-rich trajectories can make optimization unstable. We introduce FOCUS, which combines local optimal-transport alignment with short-horizon differentiable actor-critic learning. FOCUS matches agent rollout windows to multiple sampled expert segments using a soft best-of-$K$ objective that emphasizes locally compatible matches. The resulting transport-based costs provide differentiable imitation rewards, while a learned critic bootstraps value beyond the rollout horizon. FOCUS requires neither temporally aligned trajectories nor expert-aligned initialization or teacher forcing. Across contact-rich locomotion tasks in DFlex, including a musculoskeletal humanoid controlled by 152 muscle activations, FOCUS improves sample efficiency and training stability over adversarial imitation-from-observation methods and prior differentiable imitation baselines.