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

João A. Cândido Ramos; Lionel Blondé; Hugues Vinzant; Jacopo Castellini; Alexandros Kalousis

cs.LGOct 30, 2026

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

Authors: João A. Cândido Ramos, Lionel Blondé, Hugues Vinzant, Jacopo Castellini, Alexandros Kalousis

Organizations: University of Geneva, Switzerland (UNIGE) · University of Applied Sciences and Arts Western Switzerland (HES-SO)

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.

Figures & tables

Figure 1: Aggregate performance of selected methods across four DFlex tasks. Curves show the interquartile mean (IQM) of expert-normalized return; shading indicates 95% confidence intervals. Progress is expressed as a percentage of each task’s wall-clock budget (left) or interaction budget (right); the full comparison, budgets, and per-task results are reported in Appendix C .

Figure 2: Computational graph of FOCUS. A short agent rollout (green) is matched to sampled expert segments (orange, blue, yellow) using OT. Soft best-of-

K

aggregation emphasizes well-aligned segments, and the resulting supervision is backpropagated through the simulator and policy (red arrows).

Panel of Figure 3 — Figure 3: Representative local OT couplings: (a) diffuse, (b) near-diagonal, and (c) shifted matching patterns. Rows index agent transitions and columns index expert transitions; darker entries indicate greater transport mass. Panels show selected training snapshots with different candidate expert segments. Additional examples appear in Appendix D .

Panel of Figure 4 — Figure 4: Expert-normalized return of FOCUS variants and baselines over both environment interactions and wall-clock training time. Curves show the mean across six seeds, and shaded regions show bootstrap 95% confidence intervals.

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

Environment	Observation Space ${\mathcal{S}}$	Action Space $\mathcal{A}$	Expert Trajectories
			# Trajectories $\times$ Length
Hopper	$\mathbb{R}^{11}$	$\mathbb{R}^{3}$	8 $\times$ 1000
Ant	$\mathbb{R}^{37}$	$\mathbb{R}^{8}$	8 $\times$ 1000
Humanoid	$\mathbb{R}^{76}$	$\mathbb{R}^{21}$	8 $\times$ 1000
SNU Humanoid	$\mathbb{R}^{53}$	$\mathbb{R}^{152}$	8 $\times$ 1000

Appendix

Table 1: Specifications of DFlex environments and expert data used in the reported experiments.

Environment	Mean return	Standard deviation
Hopper	4812.97	7.32
Ant	9329.51	37.59
Humanoid	8225.22	81.05
SNU Humanoid	6498.92	49.13

Appendix

Table 2: Expert returns in the original environment reward units. Values are rounded to two decimal places; standard deviations describe return variability, not confidence intervals.

	Hopper	Ant	Humanoid	SNU Humanoid
Actor	(128, 64, 32)	(128, 64, 32)	(256, 128)	(512, 256)
Critic	(64, 64)	(64, 64)	(128, 128)	(256, 256)

Appendix

Table 4: Actor and critic MLP hidden-layer widths for each DFlex environment. All algorithms share these widths per network, except OOPS, which uses the reference-implementation sizes (actor

(256,256)

, critic

(1024,1024)

, ReLU, no normalization) in every environment (Table 7 ). Algorithm-specific settings for PPO, SAC, SHAC, and SAPO, including the number of critics, are reported in Table 3 .

	SAMfO/ DACfO	OPOLO	GAIfO	MAAD	ILD	FOCUS- L2	FOCUS- OT
Expert seg. length						32	32
Imitation loss	Disc	Disc	Disc	Disc	Chamfer- $\alpha$	L2	OT- {L2, Cos}
Matching	$(s,s^{\prime})$	$(s,s^{\prime})$	$(s,s^{\prime})$	$(s,s^{\prime})$	$s$	$(s,s^{\prime})$	$(s,s^{\prime})$
Expert seg. count $K$						8	8
Discriminator	Critic †	Critic	Critic	Critic
Inverse model	–	Critic	–	Critic

Appendix

Table 5: Imitation-learning objectives and architectures; selected per-environment settings are in Table 6 . “Critic” denotes the critic architecture; † except SAMfO/DACfO on SNU Humanoid, whose discriminator uses

(128,128)

against a

(256,256)

critic. MAAD augments GAIfO with inverse-model regularization. Baseline settings were selected from a small grid over expert batch size, subsampling factor, discriminator update frequency and iterations, discriminator width, and inverse-model regularization. We use the same six random seeds (100, 110, 120, 1000, 1100, 1200) across all training experiments, including hyperparameter selection and ablations.

Method	Setting	Hopper	Ant	Humanoid	SNU Humanoid
SAMfO/DACfO	Disc. update freq.	5	5	10	10
	Expert batch size	512	512	512	512
	Subsampling factor	20	1	1	20
OPOLO	Disc. update freq.	5	5	5	5
	Expert batch size	512	512	512	512
	Subsampling factor	20	20	20	20

Appendix

Table 6: Selected hyperparameters for the discriminator-based imitation baselines in each environment. For these methods, subsampling expert trajectories often improved performance. SAMfO/DACfO and GAIfO have no inverse-model regularizer. Table 5 reports the architectures and tuning-seed protocol.

	PWIL	LWAIL	OOPS
RL backbone	SAC	DDPG (TD3-style)	DDPG (TD3-style)
Actor widths	Table 4	Table 4	$(256,256)$
Critic widths	Table 4	Table 4	$(1024,1024)$
Activation / norm	SiLU / LayerNorm	SiLU / LayerNorm	ReLU / none
Actor LR	5e-4	3e-4	3e-4
Critic LR	5e-4	1e-3	3e-4

Appendix

Table 7: Policy-learning settings for PWIL, LWAIL, and OOPS. All use 64 environments,

\gamma=0.99

, and the same 8 expert trajectories as the other baselines. Their imitation-specific settings are in Table 8 .

	PWIL	LWAIL	OOPS
Agent–expert matching	$(s,s^{\prime})$ atoms	$(s,s^{\prime}-s)$	$(s,s^{\prime})$ atoms
Reward	$5\exp(-\beta\,\text{cost})$ , $\beta=1/2/1/1$	Wasserstein critic	Sinkhorn, scale 5
Expert subsampling	20	1	1
Discriminator widths	–	Critic widths (Table 4 )	–
Disc. optimizer / LR	–	Adam / 1e-3	–
Gradient penalty	–	10	–

Appendix

Table 8: Imitation and pretraining settings for PWIL, LWAIL, and OOPS. PWIL’s

\beta

is given as Hopper/Ant/Humanoid/SNU Humanoid, selected from

\{1,2,3\}

on the evaluation seeds. LWAIL’s random-rollout ICVF data (approximately 9,984 transitions) are included in the plotted interaction budget. The plotted wall-clock time includes random-data collection, ICVF pretraining, and discriminator pretraining.

Figure 5: Aggregate performance of all methods across four DFlex tasks. Curves show the interquartile mean (IQM) of expert-normalized return; shading indicates 95% confidence intervals. Progress is expressed as a percentage of each task’s wall-clock budget (left) or interaction budget (right), as specified in Tables 10 and 9 .

Method	Hopper	Ant	Humanoid	SNU Humanoid
SAMfO/DACfO	$0.320\;[-0.465,0.720]$	$0.093\;[0.061,0.120]$	$0.572\;[0.456,0.711]$	$0.057\;[0.024,0.089]$
OPOLO	$0.419\;[0.213,0.634]$	$0.123\;[0.104,0.140]$	$0.265\;[0.087,0.459]$	$0.064\;[0.042,0.081]$
GAIfO	$0.792\;[0.624,0.881]$	$0.108\;[0.094,0.133]$	$0.122\;[0.030,0.191]$	$0.012\;[-0.006,0.033]$
MAAD	$0.723\;[0.664,0.796]$	$0.100\;[0.088,0.109]$	$0.234\;[0.128,0.332]$	$0.010\;[0.007,0.013]$
PWIL	$0.778\;[0.697,0.833]$	$0.349\;[0.323,0.375]$	$0.398\;[0.317,0.468]$	$0.093\;[0.073,0.112]$
OOPS	$0.934\;[0.899,0.961]$	$0.883\;[0.838,0.927]$	$0.010\;[0.007,0.013]$	$0.002\;[0.000,0.004]$

Appendix

Table 9: Expert-normalized return at the final interaction budget, reported as mean [bootstrap 95% confidence interval] across six seeds. Budgets are 10M interactions for Hopper and Ant and 15M for Humanoid and SNU Humanoid. The highest mean in each environment is bold.

Panel of Figure 6 — Figure 6: Expert-normalized return of FOCUS variants and baselines as a function of environment interactions. Curves show the mean across six seeds, and shaded regions show bootstrap 95% confidence intervals.

Method	Hopper	Ant	Humanoid	SNU Humanoid
SAMfO/DACfO	$0.608\;[0.470,0.763]$	$0.198\;[0.126,0.333]$	$0.221\;[0.047,0.361]$	$0.034\;[-0.002,0.077]$
OPOLO	$0.475\;[0.292,0.634]$	$0.104\;[0.075,0.131]$	$0.251\;[0.071,0.448]$	$0.029\;[0.002,0.063]$
GAIfO	$0.639\;[0.346,0.845]$	$0.215\;[0.071,0.500]$	$0.119\;[0.067,0.154]$	$0.025\;[0.007,0.040]$
MAAD	$0.788\;[0.740,0.832]$	$0.107\;[0.103,0.111]$	$0.242\;[0.131,0.353]$	$0.010\;[0.007,0.013]$
PWIL	$0.876\;[0.834,0.918]$	$0.362\;[0.305,0.420]$	$0.661\;[0.603,0.720]$	$0.310\;[0.226,0.384]$
OOPS	$0.851\;[0.780,0.916]$	$0.234\;[0.180,0.284]$	$0.010\;[0.008,0.011]$	$0.002\;[0.001,0.004]$

Appendix

Table 10: Expert-normalized return at the final wall-clock budget, reported as mean [bootstrap 95% confidence interval] across six seeds. Budgets are 2 hours for Hopper, 3 hours for Ant, and 6 hours for Humanoid and SNU Humanoid. The highest mean in each environment is bold.

Panel of Figure 7 — Figure 7: Expert-normalized return of FOCUS variants and baselines as a function of wall-clock training time. Curves show the mean across six seeds, and shaded regions show bootstrap 95% confidence intervals.

Method	Hopper	Ant	Humanoid	SNU Humanoid
critic-exp-l2	$0.988\;[0.981,0.994]$	$0.818\;[0.647,0.984]$	$0.816\;[0.784,0.847]$	$0.912\;[0.902,0.920]$
no-critic-exp-l2	$0.887\;[0.834,0.940]$	$0.850\;[0.764,0.923]$	$0.015\;[0.012,0.017]$	$0.265\;[0.065,0.515]$
critic-logexp-l2	$0.993\;[0.989,0.996]$	$0.662\;[0.341,0.906]$	$0.432\;[0.150,0.721]$	$0.458\;[0.158,0.769]$
no-critic-logexp-l2	$0.799\;[0.587,0.950]$	$0.610\;[0.341,0.854]$	$0.013\;[0.010,0.016]$	$0.281\;[0.063,0.535]$
critic-neg-l2	$0.816\;[0.499,0.983]$	$0.008\;[0.005,0.011]$	$0.007\;[0.005,0.009]$	$0.008\;[0.007,0.010]$
no-critic-neg-l2	$0.870\;[0.768,0.942]$	$0.602\;[0.354,0.782]$	$0.012\;[0.008,0.016]$	$0.155\;[0.008,0.339]$

Appendix

Table 11: Reward-shaping and critic ablation. Values are final returns normalized by each environment’s expert return, reported as mean with a 95% bootstrap confidence interval across six seeds, shown as [lower, upper]. The configurations critic-exp-l2 and critic-logexp-cos correspond to FOCUS-OT-L2 and FOCUS-OT-Cos, respectively.

Figure 8: Learning curves for the reward-shaping and critic ablation. Returns are normalized by each environment’s expert return (dotted line at $1$ ). Solid and dashed curves denote runs with and without a critic, respectively. Lines show the mean across six seeds, and shaded bands show 95% bootstrap confidence intervals.

Setting	L2	OT $K=1$	OT $K=4$	OT $K=8$	OT $K=16$	OT $K=32$	OT $K=64$
Matching step time	0.162%	0.668%	0.689%	0.691%	0.724%	1.025%	1.850%

Appendix

Table 12: Runtime scaling of the matching step as a percentage of total epoch time on SNU Humanoid. All settings other than

K

are held fixed (horizon

H=32

, segment length

L=32

, parallel environments, batch size, actor/critic updates, hardware), so only the matching step varies.

Panel of Figure 9 — Figure 9: Full sequence of saved Sinkhorn transport-plan visualizations. Each panel reports the training step and candidate expert-segment index $k$ for the corresponding plan. The two plans at step $768$ k show different candidate expert segments sampled for the same training stage, including the aligned case highlighted in Fig. 3 .

Panel of Figure 10 — Figure 10: Expert-normalized return of FOCUS-OT-Cos with SHAC and SAPO as the underlying RL algorithm, evaluated across 1, 4, 8, and 16 expert trajectories.

Panel of Figure 11 — Figure 11: Expert-normalized return of FOCUS-OT-L2 with SHAC and SAPO as the underlying RL algorithm, evaluated across 1, 4, 8, and 16 expert trajectories.

Panel of Figure 12 — Figure 12: Expert-normalized return of FOCUS-L2 with SHAC and SAPO as the underlying RL algorithm, evaluated across 1, 4, 8, and 16 expert trajectories.

Method	Hopper	Ant	Humanoid	SNU Humanoid
FOCUS-L2	$0.951\;[0.938,0.966]$	$0.459\;[0.315,0.653]$	$0.596\;[0.479,0.714]$	$0.582\;[0.419,0.723]$
FOCUS-OT-L2	$\mathbf{0.988\;[0.981,0.994]}$	$0.818\;[0.647,0.984]$	$0.816\;[0.784,0.847]$	$\mathbf{0.912\;[0.902,0.920]}$
FOCUS-OT-Cos	$0.983\;[0.979,0.986]$	$\mathbf{0.899\;[0.884,0.914]}$	$\mathbf{0.934\;[0.895,0.965]}$	$0.894\;[0.876,0.912]$
ZO-FOCUS-SAC-L2	$0.702\;[0.098,0.872]$	$0.431\;[0.291,0.458]$	$0.220\;[0.191,0.242]$	$0.083\;[0.035,0.102]$
ZO-FOCUS-SAC-OT-L2	$0.338\;[-0.017,0.879]$	$0.202\;[0.112,0.258]$	$0.215\;[0.189,0.238]$	$0.118\;[0.109,0.126]$
ZO-FOCUS-SAC-OT-Cos	$0.844\;[0.810,0.878]$	$0.216\;[0.158,0.236]$	$0.243\;[0.217,0.318]$	$0.071\;[0.060,0.208]$

Appendix

Table 13: First-order (FOCUS) vs. matched zeroth-order (ZO-FOCUS) optimization at the final interaction budget. Values are means [bootstrap 95% confidence intervals] across six seeds. ZO-FOCUS results are reported separately for SAC and PPO. Within the FOCUS rows, the highest mean in each environment is bold.

Figure 13: Effect of candidate-pool size on training performance across Hopper, Ant, Humanoid, and SNU Humanoid. Each point averages 24 final returns, with six seeds per environment and each return normalized by its environment’s expert return. Error bars show 95% bootstrap confidence intervals. Both OT variants use SAPO; a normalized return of $1$ corresponds to expert performance.

Figure 14: Interaction between candidate-pool size and rollout horizon on Hopper. Points show final expert-normalized returns averaged across six seeds at $H\in\{16,32\}$ ; error bars show 95% bootstrap confidence intervals. The dotted line marks expert-level return.

Environment	$K$	Coverage (%)	Rescue (%)	Mean phase error
Hopper	4	47.3 [44.9, 49.7]	32.6 [30.2, 35.0]	0.200 [0.190, 0.211]
	8	78.2 [76.5, 79.9]	46.3 [43.6, 48.9]	0.147 [0.139, 0.156]
	16	95.8 [94.9, 96.6]	56.6 [54.1, 58.8]	0.118 [0.112, 0.124]
Ant	4	51.8 [49.8, 53.9]	40.1 [38.0, 42.3]	0.168 [0.162, 0.174]
	8	82.0 [80.2, 83.8]	63.4 [60.9, 65.9]	0.102 [0.098, 0.107]
	16	97.7 [97.2, 98.3]	81.4 [79.6, 83.3]	0.073 [0.069, 0.078]

Appendix

Table 14: Nested-pool phase-rescue comparison at a target half-cycle mismatch. Entries are point estimates [95% trajectory-cluster bootstrap CI], using 512 held-out query windows and four pools per window. All conditions share the same mismatched baseline; percentages use all 2048 pairs, and mean phase errors are measured in gait cycles.

Environment	Period $P$ (steps)	Phase error $K=1$	Phase error $K=8$	Reduction factor
Hopper	37	0.486 [0.486, 0.486]	0.147 [0.139, 0.156]	$3.31\times$
Ant	23	0.478 [0.478, 0.478]	0.102 [0.098, 0.107]	$4.69\times$
Humanoid	20	0.500 [0.500, 0.500]	0.121 [0.116, 0.125]	$4.13\times$
SNU Humanoid	30	0.477 [0.477, 0.478]	0.095 [0.092, 0.099]	$5.02\times$

Appendix

Table 15: Mean circular phase errors in gait cycles [95% trajectory-cluster bootstrap CI]. Phase labels use a single shared reference-cycle template per environment. The reduction factor is the ratio of the reported mean errors (

K=1

divided by

K=8

), computed from the rounded point estimates.

Imitation Learning

Imitation learning trains agents to perform tasks by learning from expert demonstrations, feedback, or observed behavior.

See all 268 papers

Estimate, Don't Imitate: Reusing Differentiable State-Based Policies for Visuomotor Control
Sep 27, 2026Denis Shcherba, Adrian Abel, Eckart Cobo-Briesewitz +3Visuomotor Policy Learning Robotic Manipulation
Fraunhofer Heinrich-Hertz-Institut, Berlin, Germany · Technische Universität Berlin, Germany
Diffeomorphic Hidden-State Steering: Improving Imitation Learning Policy Robustness through Expert Hidden-State Geometry
Oct 15, 2026KwangBin LeeImitation Learning Robot Imitation Learning
Beyond State-as-Action: Exploiting Command-State Discrepancy for Robot Imitation Learning
Sep 27, 2026Peiyan Li, Yueran Tao, Enhao Zhang +10Robot Policy Learning Robot Imitation Learning
Tsinghua University · SEEN·E Robotics · Imperial College London +3

Sep 27, 2026J/K move · Enter open · S save

Estimate, Don't Imitate: Reusing Differentiable State-Based Policies for Visuomotor Control

Denis Shcherba, Adrian Abel, Eckart Cobo-Briesewitz, Paul Mattes+2

Fraunhofer Heinrich-Hertz-Institut, Berlin, Germany · Technische Universität Berlin, Germany

Visuomotor Policy Learning Robotic Manipulation Robot State Estimation Policy Learning

Open paper AskPDF arXiv

Simulation-trained manipulation policies can exploit privileged state information to learn effective contact-rich behaviours, but deployment requires acting from partial observations such as noisy camera images. A common solution is teacher-student distillation, in which a visuomotor policy is trained to reproduce the actions of the privileged expert. This requires the student to jointly infer the task-relevant state and relearn the expert's action mapping that is already available. An alternative is to reuse the state-based expert and learn only a perceptual interface that reconstructs its missing state inputs. However, minimising the state estimate error alone does not necessarily minimise the downstream control error induced by these estimates. To bridge this gap, we train a visual state estimator using both direct state supervision and an action-consistency loss backpropagated through the frozen, differentiable expert. A scheduled objective first establishes a physically meaningful state estimate and progressively emphasises errors that affect the expert's actions. Across five goal-conditioned manipulation tasks, retaining the expert consistently outperforms direct pixel-to-action imitation from the same expert demonstration corpus. We further demonstrate sim-to-real transfer on a physical Panda robot, achieving 76% success without retraining the underlying expert.

Abstract

Imitation Learning

Explore similar work

Estimate, Don't Imitate: Reusing Differentiable State-Based Policies for Visuomotor Control

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

Beyond State-as-Action: Exploiting Command-State Discrepancy for Robot Imitation Learning