Hybrid Imitation Learning: Teleoperation Augmentation Primitives that Policies Learn to Trigger

Jonne Van Haastregt∗,1, Bastian Orthmann∗,1, Michael C. Welle∗,1,2, Yuchong Zhang2, Danica Kragic2

  • ∗ These authors contributed equally.
  • 1 INCAR Robotics AB, Stockholm, Sweden
  • 2 KTH Royal Institute of Technology, Stockholm, Sweden

Abstract. What an operator can demonstrate bounds what imitation learning can learn. Teleoperation interfaces map the human body to the robot, so motions that are hard for a human, such as holding an exact orientation, returning to the same viewpoint, or turning a wrist joint several full revolutions, are difficult to demonstrate on most teleoperation interfaces, even when they are trivial for the robot. We introduce Teleoperation Augmentation Primitives (TAPs): axis locks, perching waypoints, pose anchors, and embodiment-specific routines that the operator triggers during a demonstration by speech, an AR menu, or simply with a controller button. TAPs are themselves recorded in the demonstration and can therefore also be learned and invoked by the policy itself. In simulation, where benchmarks provide human demonstrations, we show that augmenting them with TAPs after the fact yields better policies on a wrist-camera-only peg insertion (0.372 to 0.531 success) and on a mug-cleanup task with a learned “remember this pose” anchor (0.203 to 0.323). On a real robot, three tasks show the same pattern: primitives that help the operator do not hurt the policy, and a policy triggering an unscrewing routine succeeds where a plain policy cannot resolve the multi-turn motion from images alone (67% vs. 38% success). We close with an industrial proof-of-concept case leveraging this idea.

Preprint: coming soon. The paper is currently under review.

Overview. During teleoperation the operator triggers a Teleoperation Augmentation Primitive (TAP) by speech or (AR) controller input; the press and the primitive's motion are recorded with the demonstration. The policy is trained with a discrete primitive channel next to its continuous action and learns to trigger the same primitives at inference time. The mechanism is policy-agnostic: in this work the channel is carried by Diffusion Policy, a behavioural-cloning MLP (BC) and a causal-Transformer behavioural-cloning policy (BC-Transformer).

Overview of Hybrid Imitation Learning: the operator triggers TAPs during demonstration; the policy learns to trigger the same TAPs

Teleoperation Augmentation Primitives

A TAP is a tuple (trigger, routine, parameters): a trigger (the operator's press during demonstration; the policy's channel at execution), a routine that takes over the robot from the trigger until it terminates, and parameters that come from static configuration, from anchors recorded by the operator, or from the robot's proprioception at the trigger. We use four types of TAPs:

  • Axis locks. Adding rotation to translation is hard for a human hand. An axis lock fixes one or more of {x, y, z, roll, pitch, yaw} in the tool or base frame: the operator reaches a desired orientation, locks the rotational axes, and completes the translation without holding the orientation by hand; locking x and y gives a pure vertical insertion.
  • Perching waypoints. A pre-set pose the operator returns to when useful: it brings the scene back into the eye-in-hand camera, so no stationary camera is needed, and splits a long task into shorter segments with a common start.
  • Anchors. A zero-motion primitive (“remember this pose”): the press records the current end-effector pose, and a later primitive executes relative to it. Anchors are how a learned policy can dynamically integrate information into a routine or a perching waypoint.
  • Routines. Embodiment-specific motions the operator's body cannot perform through the interface, such as rotating a wrist joint through several full revolutions. The unscrewing routine is close → rotate → open → rotate back → close → rotate → lift, parameterised by the repetitions and the rotation per step. Routines are open-loop in this paper; a closed-loop controller or a sub-policy can be triggered the same way.
The four TAP types on the four tasks: axis locking (vial aspiration), perching waypoints (open-container liquid transfer), open-loop routine (container unscrewing), anchor (mug cleanup, simulated)

TAP types. (a) Vial aspiration: the operator locks the rotational axes once the needle is aligned. (b) Open-container liquid transfer: perching waypoints bring the source and the target container into the eye-in-hand view. (c) Container unscrewing: an unscrewing routine grasps, unscrews and lifts the cap. (d) MugCleanup (simulated): an anchor press stores the drawer frame while the open gripper sits behind the handle; after the mug is grasped, a carry routine transports it into that frame and shuts the drawer.

Triggering during teleoperation

TAPs are triggered on a Quest-based teleoperation interface by speech (the operator names one of the predefined commands, an embedded Whisper model transcribes it and the closest command executes, with vibration patterns confirming start, end and rejection), through an AR menu over the controller in the headset, or by haptic buttons: experienced operators trigger the same menu blind, with vibration patterns telling them which menu is active. Every press is logged with the demonstration as a channel per primitive.

Three ways to trigger a TAP: speech command, AR interface, controller only with haptic feedback

The primitive channel

With K primitives available, a channel pt ∈ {−1, +1}K is appended to the robot action and the policy is trained on the augmented action [xt, pt]. While a routine runs, the executed motion is recorded as ordinary actions, so a policy trained without the channel on the same data sees the routine as demonstrated motion. Any policy that regresses an action chunk can regress the channel. At execution a primitive is triggered on a rising edge of its channel while no primitive is active; the primitive takes over until it terminates while the policy keeps predicting, then the policy takes control back.

Simulation study

The simulation study asks whether demonstrations with TAPs produce a better policy than without. Since the demonstrations come from benchmarks that predate TAPs, the unassisted control is unbiased. We start from human teleoperation datasets in robosuite: the proficient-human Square demonstrations of robomimic and a seeded subset of the MugCleanup D1 demonstrations of MimicGen. A demonstration is replayed from its recorded state up to the point an operator should have activated a TAP, the primitive is executed from that state with observations rendered online, and the demonstration is stitched back into the unassisted one afterwards. The unassisted demonstrations are therefore an identical control: same demonstration names, same split; the only difference is the TAP.

  • Square (wrist camera only, no proprioception): the robot grasps a square nut and threads it onto a peg that is out of view while the nut is grasped. The TAP is a perching waypoint above the peg; three waypoints cover the three ways operators grasp the handle. After the perch, the policy performs the insertion itself.
  • MugCleanup (wrist camera + proprioception): open a drawer, place a mug upright inside, shut the drawer. Two TAPs: an anchor press with the open gripper behind the handle, the last step before the drawer moves, and a carry routine that, once the mug is grasped, transports it into the anchored drawer frame and shuts the drawer.

Rollouts

Left: the policy trained on the unassisted demonstrations. Right: the policy trained on the same demonstrations collected with TAPs, triggering the primitive itself (green pill: who controls the robot; orange thumbnail on MugCleanup: the wrist image at the anchor press). The large view is for the viewer; the policy sees only the inset wrist camera (plus proprioception on MugCleanup).

Square (wrist camera only in both arms), Diffusion Policy, the same initial state on both sides. Paper: unassisted 0.372 vs. TAP 0.531 success, pooled over 3 training seeds × 500 episodes. This reel: 12 consecutive episodes from one training seed, 2× speed; 6 of 12 (unassisted) and 6 of 12 (TAP) succeed, which does not reproduce the table and is not meant to.

MugCleanup (wrist camera + proprioception in both arms), Diffusion Policy, the same initial state on both sides. Paper: unassisted 0.203 vs. TAP 0.323 success, pooled over 3 training seeds × 500 episodes. This reel: 12 consecutive episodes from one training seed, 2× speed; 3 of 12 (unassisted) and 6 of 12 (TAP) succeed, which does not reproduce the table and is not meant to.

Rollouts of every condition in the tables below (no-channel control, BC, BC-Transformer): All simulation rollouts

Square: trajectories of unassisted demonstrations vs TAP-collected demonstrations MugCleanup: trajectories of unassisted demonstrations vs TAP-collected demonstrations

The two simulated tasks with visualised trajectories. Left of each pair: unassisted demonstrations, the operator throughout. Right: the same demonstrations TAP-collected: the operator (blue) up to the hand-over, the primitive (orange), and on Square the stitching back to the operator's path after the perch (green); on MugCleanup the star is the anchor press whose pose the policy learns to remember.

Main result

Diffusion Policy, three training seeds per condition, 500 seeded initial states per seed, i.e. 1500 episodes per condition; both arms share the demonstrations, split, seeds, architecture, observation, training budget, horizon and success test. “Stage not reached” is the nut never lifted (Square) and the mug never grasped (MugCleanup).

TaskDemonstrationsSuccessStage not reachedSuccess | stage
Square
(wrist only)
unassisted0.3720.2070.469
TAP0.5310.2250.685
MugCleanup
(wrist + proprio.)
unassisted0.2030.5090.412
TAP0.3230.5340.692

TAP-collected demonstrations raise pooled success from 0.372 to 0.531 on Square (Fisher's exact test, p = 2.1×10−18) and from 0.203 to 0.323 on MugCleanup (p = 9.3×10−14); all paired comparisons on identical initial states are significant. On Square the gap is entirely downstream of the grasp: the fraction of episodes in which the nut is never lifted is unchanged, while success given a lift rises from 0.469 to 0.685. On MugCleanup both arms grasp about equally often (never grasped 0.534 vs. 0.509, p = 0.18); converting a grasp into a success is where the TAP arm gains (0.412 to 0.692).

Square, Diffusion Policy: pooled success and result by stage, three seeds as points
MugCleanup, Diffusion Policy: pooled success and result by stage, three seeds as points

Left: Square. Right: MugCleanup. Pooled success (a) and the same result by stage (b), with the three training seeds as points.

Data effect or trigger effect?

To separate what the TAP data contributes from what the learned trigger contributes, we also train on the TAP demonstrations without the channel: the policy sees the routine's motion as demonstrated motion and runs alone at execution, like the unassisted policy.

TaskunassistedTAP data, no channelTAP data + channel
Square0.3720.2970.531
MugCleanup0.2030.2430.323

On MugCleanup about a third of the gain is in the data, the regular carry being easier to imitate than the operators' own transport, and two thirds come from the learned hand-over. On Square the TAP data on its own costs: the no-channel policy must imitate the perch routine from the wrist image alone, without the end-effector pose the routine itself servos on, and falls below the unassisted control; with the channel the same demonstrations give the best policy.

TAPs across policy classes

The primitive channel is an extra regression target, so any policy class whose loss can treat it as one can carry it. On the same two datasets, with the same image encoder, epochs, batch size, seeds and evaluation, we also train a feed-forward behavioural-cloning MLP (BC) and a causal-Transformer behavioural-cloning policy (BC-Transformer). Pooled success over three seeds × 500 episodes:

TaskPolicyunassistedTAP + channelp
Square
(wrist only)
Diffusion Policy0.3720.5312.1×10−18
BC0.0850.2502×10−34
BC-Transformer0.2630.3297×10−5
MugCleanup
(wrist + proprio.)
Diffusion Policy0.2030.3239.3×10−14
BC0.0250.0743×10−10
BC-Transformer0.1060.1532×10−4

Every class improves on both tasks: the channel carries across an MLP, a causal Transformer and a diffusion head, in both observation regimes.

Real-robot validation

The real-robot experiments validate the full system end to end on a ViperX 300 S arm with a single eye-in-hand camera and no stationary camera: demonstrations collected with the operator triggering TAPs, a policy that learns to trigger the same TAPs, and rollouts on novel initial conditions. Unlike in simulation, every real demonstration was collected with TAPs, and both policies are trained on that data. Ours is Hybrid Imitation Learning instantiated on Diffusion Policy (Hybrid-Diffusion), triggering the primitives it was given; the baseline (DP) is the same Diffusion Policy on the same demonstrations without the primitive channel, so it sees each primitive's motion as demonstrated actions and must reproduce it from its observations. The comparison therefore isolates the learned trigger. For each task we sample seven initial configurations with the same seed for both policies and execute each three times, 21 rollouts per policy and task.

Task (TAP)PolicySuccessFail
Vial aspiration
(axis lock)
Ours13 (62%)8
DP12 (57%)9
Liquid transfer
(perching waypoints)
Ours15 (71%)6
DP13 (62%)8
Unscrewing
(routine)
Ours14 (67%)7
DP8 (38%)13
Per-configuration outcomes of the TAP-triggering policy (HD) and the baseline Diffusion Policy (D) on the three real-robot tasks

Outcome of every rollout per initial configuration (pos 1–7, three repetitions each). HD: Ours (Hybrid-Diffusion); D: baseline Diffusion Policy.

In no task does the TAP-triggering policy perform worse than its counterpart. On vial aspiration and liquid transfer the two perform alike: the axis lock and the perching waypoints made the demonstrations clean for both policies, and a policy trained on clean demonstrations can learn to hold an orientation, or to return to a vantage point, without an explicit trigger. On unscrewing the upside appears: the baseline turns the cap once or twice and then lifts while the cap is still threaded, since a closed cap and a cap turned once look nearly identical in the images, so the policy faces a multimodal choice between turning again and lifting; the routine turns a fixed number of times, so the hybrid policy removed the cap in every rollout in which it triggered the routine from the right position.

Vial aspiration (axis lock)

A syringe-equipped arm draws liquid from a vial mounted on a sphere that permits arbitrary orientation, up to the point where no liquid would flow. Aligning the needle with the vial is the hard part, so the operator locks all rotational axes once the orientation is reached.

All rollouts: Vial aspiration

Open-container liquid transfer (perching waypoints)

With the same syringe, draw from the right container and deposit into the left. The operator starts from a perching waypoint to draw and triggers a second waypoint that brings the target container into clear view to deposit.

All rollouts: Open-container liquid transfer

Container unscrewing (routine)

A second arm, not controlled by the policy, holds a container with a threaded cap. The operator positions the gripper over the cap and triggers the unscrewing routine, which grasps the cap, unscrews it and lifts it off, after which the cap is placed on the table. The baseline removes the cap in 8 of 21 rollouts and the hybrid policy in 14.

All rollouts: Container unscrewing

Deployment: policy-triggered routines for clutch assembly

We applied the same decomposition to an industrial clutch-assembly task as a proof of concept: stacking the discs onto the clutch shaft of a multi-plate clutch. One clutch assembly is 35 discs of three types (a spring disc that can face up or down, a friction disc, and a steel disc), stacked onto a shaft whose rotation is not fixed, mounted on a tray whose position is only roughly known. A UR16e with a compliant two-finger gripper and two eye-in-hand cameras (no scene camera) retrieves the discs from three magazines, which an operator loads in random orientations.

Four routines can be triggered to retrieve the discs. Each routine retrieves its disc from the magazine and brings it to a perch point over the workspace from which the tray is in view; the policy then performs the insertion and retreats, and the next routine is triggered. The demonstrations were collected with the routines triggered by the operator, who teleoperated only the insertion of each disc: 15 complete assemblies, i.e. 525 demonstrated insertions. The policy's best runs assembled two complete clutches, 70 consecutive insertions, without an interruption.

Clutch assembly proof of concept: a UR16e with a compliant two-finger gripper and two eye-in-hand cameras, three disc magazines above the tray with the clutch shaft

Clutch assembly proof of concept. The three magazines (top) hold the disc types; a routine retrieves a disc and brings it to a perch pose over the tray (bottom), where the shaft is mounted and the stack accumulates.

Clutch-assembly proof of concept: policy-triggered routines retrieve the discs and the learned policy performs the insertions.