Single-Pass Multimodal Policies over Action Chunks
Overview
Multimodal imitation learning requires diverse executable futures under the same observation and consistent behavior across replanning cycles. We present Conditional Trajectory Peaks (CTP), a single-pass policy framework that jointly predicts complete action-chunk candidates, probability masses, and trajectory scales. Distribution-Aware Peak Specialization (DAPS) specializes trajectory peaks using trajectory-level posterior responsibilities and mass- and scale-modulated overlap constraints. Evidence-Gated Trajectory Belief Transport (ETBT) maintains cross-chunk consistency through geometric correspondence between exchangeable candidate sets, while allowing current policy evidence to override historical constraints. CTP achieves a coverage score of 91.40% on Push-T; success rates of 100.0%, 79.72%, and 84.44% on D3IL Avoiding, Aligning, and Sorting-2, respectively. On LIBERO, CTP achieves an average success rate of 97.25%. In real-world dual-arm experiments, CTP preserves both placement modes in a two-plate task, succeeding in all 50 trials. On bottle uprighting and pen placement into a holder, it maintains success rates comparable to π0.5 while reducing policy inference latency from 218.24 ms to 75.80 ms. These results demonstrate that single-pass trajectory modeling can combine multimodal behavior, closed-loop consistency, and efficient inference.
Motivation
Method
Given an observation history, CTP models the conditional distribution of an entire H-step action chunk rather than predicting each action independently. A shared policy backbone produces K complete trajectory centers, their probability masses, and bounded residual scales in a single forward pass. One discrete component therefore explains the full planning horizon, giving every candidate a trajectory-level behavioral identity while retaining NFE = 1.
CTP separates the two challenges of multimodal closed-loop control. During training, DAPS prevents exchangeable candidates from collapsing onto the same dominant solution. During receding-horizon execution, ETBT associates candidates geometrically across consecutive replanning cycles and preserves a behavior only while it remains compatible with current observations.
CTP defines a finite conditional mixture over complete action chunks. Each peak predicts a trajectory center μ, probability mass π, and residual scale σ. The representation supports either direct time-domain chunks or a compact trajectory basis, without predefined behavioral anchors.
Trajectory-level posterior responsibilities assign demonstrations using fit, probability mass, and uncertainty together. A distribution-aware overlap constraint then separates high-probability peaks that redundantly explain the same data region, while unsupported candidates may reduce their mass instead of being forced into artificial modes.
After executing part of a chunk, ETBT matches every previous unexecuted suffix to current candidate prefixes. It transports behavioral belief through these permutation-equivariant correspondences, fuses history as a soft prior with current policy evidence, and releases that prior when correspondence becomes unreliable.
Quantitative results
| Benchmark | CTP | Evaluation summary |
|---|---|---|
| Push-T | 91.40% | 200 fixed initial conditions; IMLE Policy 88.74, Diffusion Policy-C 87.90 |
| D3IL | 100 / 79.72 / 84.44 | Avoiding / Aligning / Sorting-2 success rates |
| LIBERO | 98 / 100 / 98 / 93 | Spatial / Object / Goal / Long-Horizon |
Component analysis
| Variant | Peaks | DAPS | ETBT | Push-T |
|---|---|---|---|---|
| Single peak | K = 1 | — | — | 77.34 |
| Multi-peak | K = 4 | — | — | 80.30 |
| Peak specialization | K = 4 | Yes | — | 85.67 |
| CTP | K = 4 | Yes | Yes | 91.40 |
Matched Push-T training and evaluation protocol; higher is better.
Qualitative results
D3IL
Avoiding, Aligning and Sorting-2 are the three tasks reported in the main quantitative table; Sorting-4 is included as an additional D3IL rollout. Every video is generated by its corresponding evaluated checkpoint on a formally successful context.
Removing ETBT causes a collision under the matched initial state; full CTP preserves the route and succeeds.
Push-T
A matched initial condition compares a single-peak policy, iterative Diffusion Policy and full CTP.
LIBERO
CTP is integrated into π0.5, predicts four 50-step candidates, executes 10 actions and replans with trajectory belief decoding. Representative successful rollouts are shown for all four LIBERO suites.
Efficiency
Real-robot experiments
All eight displayed rollouts begin from the same object placement. The opposite successful outcomes therefore reflect conditional policy multimodality rather than variation in the initial state; the side counters summarize the full 50-trial evaluation.
Successful executions on bottle uprighting and pen placement into a holder, using the deployed CTP policy.
