Mixture of Horizons in Action Chunking

In plain words

A robot policy that predicts a list of future actions at once, instead of one action per step, has to pick how far ahead it plans. Plan too far and the fine motions near the start get sloppy. Plan too little and the robot can't see the end of a long task. Today's approaches usually fix one length and accept the trade-off.

Mixture of Horizons does not pick one length. It has the same network predict several lengths side by side, from short to long, using the same underlying computation. A small learned gate then decides, for each future step, how much to trust each length. Short predictions dominate where precision matters, and long ones where planning matters. The added cost is small, because the heavy visual and language part is computed once and shared.

The same spread of predictions also gives a signal for how far ahead the robot can safely commit before it looks again. When the predictions agree, the robot can run more steps before replanning. When they disagree, it replans sooner.

TL;DR

MoH is a plug-in for action heads that predict action chunks. It splits a chunk into prefixes of several lengths, runs them in parallel through a shared action transformer, and fuses them with a per-step softmax gate plus a balance loss §3.2. The authors report gains on four VLA base models across LIBERO, RoboTwin, RoboCasa and a small real-robot study, with pi0.5 reaching 99.0% average on LIBERO tab. 1. A dynamic-inference variant uses horizon agreement to choose how much of each chunk to execute §3.3. All numbers are author-reported.

Problem

Modern VLAs put a full-attention action transformer on top of a pretrained vision-language backbone and predict a chunk of future actions §3.1. The chunk length, called the horizon, strongly affects performance §1. On pi0, the authors observed that longer horizons help long-horizon LIBERO suites but hurt short-horizon ones, and that no single horizon is best across suites §1 fig. 1. Fixed execution length at inference also ignores how confident the policy is at each moment §3.3.

Method

Setting §3.1. The policy takes multi-view images, an optional history, a language instruction and an optional proprioceptive state, and outputs an action chunk of length H. Two action-head families are covered: flow-matching policies that integrate a learned velocity field, and one-step policies that output the chunk directly, either as classification bins or as \ell_1 regression. The paper's MoH experiments use flow-matching (pi0, pi0.5, StarVLA's GR00T-style head) and a one-step regression head (pi_reg) §4.1.

Chunk rearrangement §3.2. Pick a maximum horizon H and candidate horizons h_1 < \dots < h_N = H. For each candidate h, take the first h actions of the ground-truth chunk. All candidates share one observation context, which the VLM encodes once. Each candidate is padded to length H, and a mask removes positions beyond h, so one batched action-transformer pass produces all horizons in parallel.

Gate §3.2 fig. 3. The action transformer's hidden states go through an action head that outputs a prediction per horizon, and a linear gate head that outputs one logit per step and horizon. At step k only horizons with h \ge k are valid; a softmax over the valid ones gives the weights. The fused prediction at step k is the weighted sum of the horizon predictions for that step. The authors call this the simplest effective gate and note it leaves the backbone unchanged.

Loss §3.2. The total loss is the fused-prediction loss plus a weighted sum of the per-horizon losses plus a weighted balance loss. Per-horizon and fused losses use the same objective as the base policy, which is velocity matching for flow-matching models and the classification or regression loss for one-step models §3.1. The balance loss splits steps into intervals by the horizon boundaries. Within each interval it takes the average gate usage of each active horizon and penalises the squared coefficient of variation of those averages. This stops the gate from collapsing onto one horizon.

# MoH training loss for one batch  [§3.1] [§3.2]
# Inputs:  obs; actions (B, H, d_a) with H = max(horizons)
# Output:  scalar loss

def moh_loss(obs, actions, horizons):
    # Run the VLM prefix once; all horizons reuse its KV cache  [§3.2]
    prefix_cache = vlm_prefix(obs)
    # Flow time per example and a Gaussian noise chunk  [§3.1]
    tau = sample_time(actions.shape[0])      # (B,), in [0, 1]
    noise = sample_normal(actions.shape)     # (B, H, d_a)
    # Point on the reference path and its velocity target  [§3.1]
    noisy = (1 - tau) * noise + tau * actions
    target = actions - noise                 # (B, H, d_a)
    # One batched suffix pass; per-horizon preds and gate logits  [§3.2]
    preds, gate_logits = action_expert(prefix_cache, noisy, tau)
    # preds: (N, B, H, d_a); gate_logits: (B, H, N)
    individual = 0.0
    for i, h in enumerate(horizons):
        # Each horizon is supervised on its first h steps  [§3.2]
        individual = individual + mse(preds[i, :, :h], target[:, :h])
    # Softmax over the horizons still valid at each step k  [§3.2]
    weights = masked_softmax(gate_logits, horizons)  # (B, H, N)
    fused = 0.0
    for i in range(len(horizons)):
        # Gate-weighted sum of horizon predictions  [§3.2]
        fused = fused + weights[:, :, i, None] * preds[i]
    mix = mse(fused, target)
    # Usage balance across horizons per step interval  [§3.2]
    balance = balance_loss(weights, horizons)
    return mix + 1.0 * individual + 1e-3 * balance

# Balance loss over step intervals  [§3.2]
# Interval i covers steps (h_{i-1}, h_i]; active horizons are h > h_{i-1}

def balance_loss(weights, horizons):
    # Boundaries start at 0, then each candidate horizon  [§3.2]
    bounds = [0] + horizons
    terms = []
    for start, end in zip(bounds[:-1], bounds[1:]):
        active = [i for i, h in enumerate(horizons) if h > start]
        if len(active) < 2:
            # A single active horizon has no balance term  [§3.2]
            continue
        # Mean gate usage of each active horizon over batch and steps  [§3.2]
        usage = weights[:, start:end, active].mean(axis=(0, 1))
        # Squared coefficient of variation, eps is a small constant  [§3.2]
        terms.append(var(usage) / (mean(usage) ** 2 + eps))
    return mean(terms)

# Inference with cross-horizon consensus  [§3.3] [alg. 1]
# Inputs:  obs; horizons; n (min steps), m (min active horizons), r (ratio)
# Output:  executable prefix of the fused chunk

def moh_act(obs, horizons, n, m, r):
    # Denoise all horizons together and get the gate weights  [§3.2]
    fused, preds, weights = sample_all_horizons(obs)  # fused: (H, d_a)
    # Per-step disagreement: gate-weighted l1 to the fused action  [alg. 1]
    disagree = weighted_l1(preds, fused, weights)     # (H,)
    # Threshold is r times the mean disagreement of the first n steps  [alg. 1]
    threshold = mean(disagree[:n]) * r
    k_exec = n
    for k in range(n, H):
        # Stop when fewer than m horizons are still active at step k  [alg. 1]
        if active_count(k, horizons) < m:
            break
        # Stop when disagreement at step k exceeds the threshold  [alg. 1]
        if disagree[k] > threshold:
            break
        k_exec = k + 1
    return fused[:k_exec]

Pseudocode above is our condensation. The exact step-accounting for the disagreement term and the boundary comparisons differ from the code in ways listed under Paper vs code.

Dynamic inference §3.3 alg. 1. Each step's fused action is compared with the horizon predictions that are still active at that step, using the gate weights. The first n steps set a threshold, which is a data-dependent mean times a ratio r. Execution then extends through every step whose disagreement stays under the threshold, while at least m horizons remain. Steps past the first failure are deferred to the next replanning call.

Hyperparameters

Results

Author-reported. Setup: LIBERO (four suites, 500 trials each, same seeds), mixed-suite training, 30k iterations, batch 32, 3B-class base models §4.1 §4.2. The baseline for each row is the same model without MoH.

LIBERO average tab. 1

pi0.5 per suite tab. 1: Spatial 98.8 to 98.8, Object 99.0 to 100, Goal 97.6 to 98.8, Long 95.4 to 98.4. The paper calls 99.0 a new state of the art on this benchmark §4.2.

RoboCasa, GR00T base tab. 2: average success 28.0 plus or minus 1.2 to 31.4 plus or minus 1.2, across five single-task settings with three seeds.

RoboTwin 2.0, pi0 base fig. 4: 7 tasks, easy and hard modes, 100 trials each. The paper reports the highest average with MoH in the figure. Per-task numbers are only in the figure, not in the text.

Real robot §5: pi0.5 on the fine-grained towel fold (T4) goes from 15/20 to 18/20 successes. The text reports gains on all tasks, but gives numbers only for T4.

Horizon density, pi0.5 on LIBERO tab. 3: single horizon 30 gives 97.7 average; d=10 gives 98.3; d=5 gives 98.3; d=3 gives 99.0; d=2 gives 98.3; d=1 gives 98.3.

Fusion and balance ablation, pi0.5 on LIBERO tab. 4: ten identical horizons of 30 give 97.9; loss reweighting with no MoH gives 98.1; uniform mean fusion gives 98.4; MoH without the balance loss gives 98.5; full MoH gives 99.0.

Dynamic inference §4.4 fig. 6: throughput is the average number of executed steps per chunk. The abstract and §1 claim 2.5 times the default throughput while still beating the fixed-length pi0.5 baseline on LIBERO-Long. The sweep is only plotted; the numbers are not in the text.

Temporal ensembling tab. 5: for pi0.5, adaptive temporal ensemble gives 96.6 average and drops Long from 95.4 to 90.6. MoH gives 99.0, and MoH with ensemble gives 98.4.

Paper vs code

Code: official repository linked from the authors' project page, at commit 5da35004d27a86ac1a7d70ce93f7986eff16c6a7. The repository contains no LICENSE file, so code_license is recorded as not stated. Model code is under src/openpi/, built on openpi (its files carry Apache-2.0 headers). We cite code only by path and line.

• Default horizon set differs. Paper default is \{3, 6, \dots, 30\} §4.1. The training script uses the same set run_train_moh.sh. The JAX config default is [5, 10, \dots, 30] pi0moh_config.py, and the PyTorch config default is [10, 20, 30] config.py.
• Gate noise is absent from the paper. The PyTorch forward pass adds a learned, input-dependent Gaussian noise to the gate logits during training, on by default pi0_pytorch_moh.py pi0_pytorch_moh.py pi0_pytorch_moh.py. The paper's gate is a plain linear layer followed by a softmax §3.2. The JAX model has no such noise pi0_moh.py.
• Dynamic-inference disagreement is summed over denoising steps. The paper defines disagreement at the final chunk §3.3 alg. 1. The PyTorch sampler accumulates a gate-weighted L1 term at every Euler step and uses that sum for the threshold pi0_pytorch_moh.py and pi0_pytorch_moh.py.
• Break condition is off by one. The paper stops when fewer than m horizons remain alg. 1. The code stops when the count is less than or equal to min_active_horizons pi0_pytorch_moh.py.
• Default m depends on the entry point. The sampler default is 3 pi0_pytorch_moh.py. The LIBERO example script defaults to 5, which matches the paper pi_moh_dynamic.py.
• Gradient clipping is not in the paper. The code clips the gradient norm at 1.0 config.py train_pytorch_moh.py. The paper's hyperparameter table omits this row; a commented-out line in the source carries it.
• Time convention is reversed. The paper's reference path puts noise at \tau = 0 and data at \tau = 1 §3.1. The code puts noise at t = 1 and integrates with a negative step pi0_moh.py pi0_moh.py. Both are valid, but the two conventions cannot be mixed.
• Flow-time distribution and sampler steps. The paper states neither. The code samples time from a Beta(1.5, 1) distribution rescaled into [0.001, 0.999] pi0_moh.py and defaults to 10 Euler steps pi0_moh.py.
• Loss weights match. \lambda_{\text{bal}} = 10^{-3} and the per-horizon sum with weight 1 match the paper pi0moh_config.py pi0_moh.py.
• Schedule matches. Peak learning rate 5 \times 10^{-5}, warmup 1k, and decay to 10^{-6} over 30k steps match the paper's table config.py.
• Citation note. The paper cites Shi et al. 2025 [arXiv:2502.19417] for pi0.5 §1 §4.1, and that bibliography entry is titled "Hi robot". The paper does not state which is the intended reference for the pi0.5 model.

Pitfalls

• Horizon set is ambiguous. The paper's default is a stride-3 set up to 30, but the code has three different defaults. Use the set from the training script and the paper, and check the value your checkpoint was trained with.
• Compute is not free. Each sample runs N suffix passes, with the KV cache repeated N times along the batch dimension pi0_moh.py. The "negligible overhead" claim §1 app. overhead depends on tensor parallelism and a lightweight action expert. Measure it on your hardware.
• Two masks, not one. Horizons are padded to H. The per-horizon loss covers only the first h steps, and the gate softmax must exclude horizons with h < k (1-indexed). The padding mask alone does not do either pi0_moh.py.
• Flow-time sampler and convention. The paper gives neither the time distribution nor the sampler step count. Copying the code's Beta sampler is one choice; a uniform sampler is another. Keep the time convention consistent between training and sampling (see Paper vs code).
• Balance epsilon. The paper says only "a small constant". The code does not expose the value in the material we read.
• Dynamic inference needs a gate. The consensus rule uses gate weights and per-horizon predictions. It does not apply to a plain temporal ensemble, and the authors found stacking ATE on top of MoH slightly hurt Long §4.5.
• Gains are small on a saturated benchmark. Most LIBERO gains are about one point and some suites are already near 98 to 100%. Seed variance is reported only for RoboCasa tab. 2. Treat single-seed LIBERO differences of about one point with caution.
• Real-robot evidence is small. Real-robot results are 10 to 20 rollouts per task on one platform §5.
• Max horizon 50 is worse than 30 on LIBERO. The appendix's H_{\max} = 50 MoH variant reaches 98.2 average, below the 99.0 at H_{\max} = 30 app. maximum horizon.
• The two code paths differ. The PyTorch and JAX implementations are not equivalent: only the PyTorch path adds gate noise (see Paper vs code). Say which one you follow.

Lineage

• Builds on: pi0 [arXiv:2410.24164] and pi0.5, which the paper cites as Shi et al. 2025 [arXiv:2502.19417], as base policies. Regression head pi_reg is built from pi0 with a learned query token §3.1 app. pi_reg design. Action chunking follows prior work such as ACT (Zhao et al., 2023) §1.
• Compared against: OpenVLA-OFT, CogACT with adaptive temporal ensemble [arXiv:2411.19650] tab. 5, Spatial Forcing, X-VLA, UniVLA, GR00T-N1, SmolVLA, Octo, OpenVLA, CoT-VLA, and pi0-FAST tab. 1.
• Related idea: mixture-of-experts load balancing, which the balance loss borrows §3.2.
• Superseded by: not known from the content in this wiki.