PaperlayerView exact revision on arXiv

Source-linked AI summary

FlexiWorld: Learning and Planning via Flexible Action Chunks Across Multiple Time Scales

Shidu Ren, Qilin Gu, Zhenghao Ni, Junhan Sun, Jiaqi Wang, Damien Scieur, Yunze Liu

arXiv:2609.35138v2cs.LG

TL;DR

Latent world models support goal-directed planning, but prior approaches often use fixed-length chunks and omit goal-conditioned action generation or supervise it only over short spans. FlexiWorld jointly learns latent prediction and goal-conditioned action generation across mixed goal spans and variable-length chunks, with ARCEM adding residual search. Across four benchmarks, FlexiWorld with ARCEM reaches 89.29% mean success versus 83.98% for INTACT with Guarded-A.

  • Problem

    Prior methods use fixed-length action chunks and either omit goal-conditioned action generation or limit its supervision to short goal spans.

    Source: §1 · p. 2

  • Method

    FlexiWorld jointly trains latent prediction and a goal-conditioned autoregressive actor over mixed-span goals and variable-length chunks, with ARCEM refining action chunks through residual search.

    Source: §1 · p. 2

  • Results

    89.29% mean success with ARCEM exceeds INTACT with Guarded-A at 83.98% across four tasks and four goal distances.

    Source: §4.2 · p. 7

  • Takeaways & Limitations

    A single trained model supports distant-goal control and adjustable planning granularity without retraining, balancing planning cost and control performance.

    Source: §5 · p. 9

    Student Forcing trains on generated action prefixes but retains recorded states and expert targets, leaving recovery from perturbed physical states untested.

    Source: §4.4 · p. 9

Abstract

from arXiv · show

Latent world models predict future states for goal-directed planning using action chunks spanning multiple primitive steps. Existing methods typically use fixed-length chunks and either omit goal-conditioned action generation or limit their supervision to short goal spans. We introduce FlexiWorld, a JEPA-based world model that combines mixed-span goal supervision with variable-length action chunks to improve long-horizon control. During training, we sample varying goal spans and randomly partition the actions into variable-length chunks. We jointly train the world model with a causal action encoder that embeds variable-length chunks and an autoregressive actor that generates primitive actions sequentially. Student Forcing reduces exposure bias by training on generated action prefixes. For planning, Actor-Residual Cross-Entropy Method (ARCEM) combines action-residual search with within-chunk autoregressive feedback and chunk-boundary latent prediction. Across four benchmarks and goal distances, FlexiWorld with ARCEM achieves 89.29% mean success, compared with 83.98% for the strongest baseline. PushT ablations show improved direct control from mixed-span supervision, variable-length chunks, and Student Forcing. Without retraining, FlexiWorld supports different planning chunk lengths: longer chunks accelerate ARCEM by approximately $1.3\times$ on average while maintaining comparable average success.

1 INTRODUCTION

FlexiWorld combines mixed-span goal supervision and variable-length action chunks with jointly trained latent prediction and autoregressive action generation. ARCEM refines generated chunks through action-residual search, and evaluations report stronger success than INTACT and faster planning with longer chunks.

1 INTRODUCTION

  • Prior JEPA world models leave a gap between goal-conditioned action generation and flexible primitive-action chunk boundaries for distant-goal planning.LeWM uses search without learned action generation; INTACT limits goal supervision to short spans; VLWM and Fast-LeWM retain fixed-size base chunks and use CEM without a jointly trained goal-conditioned actor.Source: §1 · p. 2
  • FlexiWorld jointly trains a causal chunk encoder and autoregressive actor with the world model across variable-length chunks and mixed-span goals.Student Forcing trains the actor on expert and self-generated action prefixes to address the exposure bias created when planning uses generated prefixes.Source: §1 · p. 2
  • ARCEM extends action-residual search to autoregressive chunks, feeding perturbed actions into subsequent within-chunk outputs and predicting latent states only at chunk boundaries.This combines policy feedback within chunks with latent prediction across chunks.Source: §1 · p. 2
  • The study evaluates PushT, Cube, Reacher, and TwoRoom across multiple goal distances, with CEM-only success similar for FlexiWorld and INTACT.Distant-goal action-prediction tests favor FlexiWorld, indicating the clearest reported benefit is in goal-conditioned action generation.Source: §1 · p. 2
  • Using ten-action rather than five-action chunks makes ARCEM approximately 1.3× faster at 50- and 100-step goals, with comparable four-benchmark mean success.The same trained model supports longer chunks; task-level success can show gains or losses, and lower expert-action latent rollout error does not consistently improve control.Source: §1 · p. 3

2 RELATED WORK

Related work spans latent world models, temporal abstraction, and policy-guided planning. FlexiWorld differs by jointly learning variable-duration latent dynamics and goal-conditioned action generation, then using residual search with autoregressive chunks.

Latent world models for goal-conditioned control

  • JEPA-based models predict representations without image reconstruction; prior methods include LeWM and Sub-JEPA, while INTACT jointly learns prediction and shared intent-to-action control using fixed-length chunks.FlexiWorld varies both transition durations and goal spans.Source: §2 · p. 3

Temporal abstraction and action chunking

  • Existing temporal-abstraction methods vary prediction horizons or organize action chunks, but do not jointly learn variable-duration latent dynamics and a goal-conditioned actor.FlexiWorld varies chunk boundaries at primitive-action resolution and jointly trains its chunk encoder, autoregressive actor, and latent predictor.Source: §2 · p. 3

Policy-guided model predictive control

  • PETS and iCEM optimize action sequences, while POPLIN uses action-residual search and INTACT preserves a Direct reference during Guarded-A search.ARCEM applies action-residual replanning to autoregressive chunks without training a residual policy.Source: §2 · p. 3

Figure 2

  • Figure 2 depicts joint training: variable-length chunk embeddings condition boundary-state prediction, while actor supervision uses local or final-goal intents and expert or generated prefixes.Training combines prediction MSE, SIGReg, and action NLL across jointly trained modules.Source: Figure 2 · p. 4

3 METHODOLOGY

FlexiWorld trains a latent world model and goal-conditioned actor across varied goal spans and action-chunk lengths. It supports direct planning and ARCEM search, which refines autoregressive chunks with action residuals and boundary-level latent predictions.

3.1 METHOD OVERVIEW

  • The method jointly trains latent prediction and action generation on sampled chunks, then uses the actor and predictor to construct Direct plans that ARCEM can refine.The overview describes actor and predictor alternation at deployment.Source: §3.1 · p. 4, §3.3 · p. 5

3.2 MULTI-TIME-SCALE TRAINING

  • Training samples 35-, 55-, or 75-step goal windows and partitions each into variable-duration chunks, with different partitions sharing the final goal.The alternative partitions provide different intermediate transitions.Source: §3.2 · pp. 4–5

3.3 FLEXIBLE ACTION CHUNKS

  • A causal action encoder embeds chunks of different lengths for latent prediction, while the same encoder supplies the actor’s previous-chunk context.The encoder processes each chunk and adds a learned chunk-length embedding.Source: §3.3 · p. 5
  • The autoregressive actor generates primitive actions for variable-length chunks, and Student Forcing mixes expert and generated prefixes during training.At each position, the expert action remains the prediction target when a prefix is supplied.Source: §3.3 · p. 5

3.4 DIRECT PLANNING AND ACTOR-RESIDUAL CEM

  • Direct planning alternates actor and predictor calls at chosen chunk lengths; longer chunks reduce predictor calls while still generating all actions autoregressively.The final chunk may be shorter to reach the requested plan length.Source: §3.4 · p. 6
  • ARCEM applies action residuals with within-chunk autoregressive feedback, updating predicted latents at chunk boundaries to refine candidate plans.It refits the residual distribution using low-cost elite plans.Source: §3.4 · p. 6, Figure 3 · p. 6

4 EXPERIMENTS

The experiments evaluate FlexiWorld against JEPA-based baselines on four visual goal-reaching tasks and multiple goal distances, then test its planning methods, training components, and deployment chunk lengths. FlexiWorld leads in mean success, while ablations and planning analyses examine the contributions and trade-offs behind that result.

4.2 MAIN RESULTS

  • At D = 100, FlexiWorld Direct reaches 33.44% on PushT versus 12.67% for INTACT, and 93.78% on Cube versus 84.67%.Direct success is similar at D = 25; D = 100 exceeds the longest training span of 75 steps.Source: §4.2 · p. 7, Figure 4 · p. 8
  • 89.29% mean success for FlexiWorld with ARCEM exceeds INTACT with Guarded-A at 83.98% across four tasks and four goal distances.FlexiWorld leads on PushT, Cube, and Reacher, with its largest gains over INTACT on PushT and Cube.Source: §4.2 · p. 7, Table 1 · p. 8
  • 86.79% mean success for FlexiWorld Direct rises to 89.29% with ARCEM, exceeding INTACT Direct at 81.62%.ARCEM’s largest gain over Direct is on PushT, from 60.39% to 68.89%.Source: §4.2 · p. 7, Table 2 · p. 8

4.3 ABLATION STUDIES

  • On PushT, variable chunks raise success from 46.08% to 50.83% without Student Forcing and from 48.42% to 52.00% with it.At D = 25, variable chunks without Student Forcing reach 88.67%, below INTACT’s 90.33%; adding Student Forcing raises success to 92.67%.Source: §4.3 · pp. 7–8, Table 3 · p. 9
  • Mixed-span training raises PushT success to 54.42% for fixed-chunk INTACT and from 52.00% to 60.17% for FlexiWorld’s full recipe.For INTACT, the mixed-span result exceeds both the 35-step and 75-step single-span results; removing Student Forcing from the FlexiWorld setting reduces success to 54.42%.Source: §4.3 · p. 8

4.4 PLANNING GRANULARITY AND FURTHER ANALYSIS

  • Using k = 10 instead of k = 5 makes ARCEM approximately 1.3× faster at D ∈ {50, 100}, with similar four-task mean success.Predictor calls halve while primitive actions remain autoregressive; k = 10 also reduces endpoint latent prediction error at all four tested PushT and Cube distances.Source: §4.4 · pp. 8–9
  • With actors disabled and the same CEM planner, FlexiWorld and INTACT achieve 40.75% and 40.17% success, respectively.Frozen-feature temporal-gap R2 rises from 0.565 to 0.603, while FlexiWorld’s next-five-action prediction improves over INTACT at D ∈ {50, 75, 100}.Source: §4.4 · p. 9
  • ARCEM can lose some Direct successes because candidate ranking uses predicted latent costs rather than actual outcomes, and recovery from perturbed physical states remains untested.Student Forcing uses generated action prefixes while retaining recorded states and expert targets.Source: §4.4 · p. 9

5 CONCLUSION

FlexiWorld jointly learns latent dynamics and goal-conditioned action generation across time scales using variable-length action chunks. Across four benchmarks, it supports search-free control, ARCEM further improves success, and longer chunks reduce planning time with comparable mean success.

5 CONCLUSION

  • A single trained FlexiWorld model supports distant-goal control and adjustable planning granularity without retraining, balancing planning cost and control performance.Future work will investigate reliable candidate ranking and evaluate FlexiWorld on physical systems.Source: §5 · p. 9

SUPPLEMENTARY MATERIAL

The supplementary material details FlexiWorld’s architecture and training, including a variable-length action encoder, autoregressive actor, and Student Forcing. It also specifies the model’s objectives and training setup.

A.1 ARCHITECTURE AND TRAINING

  • The predictor uses a shared visual backbone and six-layer causal Transformer, with actions conditioning each predictor layer through adaptive normalization and residual gates.The visual encoder is a randomly initialized ViT-Tiny; the predictor uses width 192 and learned positional embeddings over at most three latent states.Source: §A.1 · p. 13
  • A causal action encoder maps variable-length chunks to a fixed-size predictor input, while the actor generates primitive actions sequentially from state, intent, prior-chunk context, and within-chunk prefixes.The encoder combines the last valid action token with a learned chunk-length embedding; the actor predicts one primitive at a time.Source: §A.1 · p. 13
  • Training uses AdamW with a 3 × 10^-4 learning rate, and the model is trained without a pretrained visual backbone, target EMA, or auxiliary temporal loss.The prediction loss has unit weight; regularization, local action, and goal action losses have weights 0.02, 0.10, and 0.05, respectively.Source: §A.1 · pp. 13–14
  • Student Forcing trains the actor using either greedy-generated or expert within-chunk prefixes, selected with probability pSF, while stopping gradients through generated prefixes.The supervised pass updates the actor, action encoder, and attached visual representation.Source: §A.1 · p. 14
  • The objective retains SIGReg, normalizes each chunk’s action NLL by its valid length and action dimension, and averages regularization over boundary latent batches.The regularization weight is 0.02, and the per-chunk weighting does not increase with chunk length.Source: §A.1 · p. 14

A.2 SAMPLING AND TRAINING SCHEDULES

The sampling schedule mixes goal spans and variable-length action partitions, and the PushT component study compares training choices under specified evaluation conditions. The reported controls distinguish mixed-span supervision, action feedback, and baseline evaluation settings.

A.2 SAMPLING AND TRAINING SCHEDULES

  • Training samples 35-, 55-, or 75-step spans partitioned into chunks of one to ten primitives, using mixed-span and single-span schedules.Two mixed-span epochs and six single-span epochs balance exposure across span groups, but do not exactly match optimizer updates or compute.Source: §A.2 · p. 14
  • The no-Student-Forcing mixed-span variant and Full method share loader and optimization settings, differing only in pSF; the single-span variants form a fixed/variable-chunk by pSF comparison.The single-span comparisons use a 35-step span and six-epoch budget.Source: §A.2 · p. 15
  • Removing within-chunk action feedback lowers success from 50.92% to 48.83% in an independently matched re-evaluation.This control tests feedback within this architecture, not autoregressive versus parallel actors generally.Source: §A.2 · p. 15
  • 54.42% success with mixed spans exceeds 40.92% when training only at 75 steps in the fixed-interface control.The extended component study reports success averaged over four goal distances with one replan.Source: §A.2 · p. 15, Table 4 · p. 15

A.3 BASELINES AND EVALUATION PROTOCOL

  • The main evaluation permits one reobservation and replan after the first D actions, while the strict setting uses only the initial D-step open-loop plan.The external baseline models retain their official configurations and are not retrained under a common compute budget.Source: §A.3 · p. 15
  • Reported standard deviations describe variability across specified repeats, not confidence intervals or tests of statistical significance.The paper notes that run-to-run variability and aggregation choices can affect empirical comparisons.Source: §A.3 · p. 16

B FLEXIWORLD RESULTS BY GOAL DISTANCE

Table 5 expands FlexiWorld’s Direct and ARCEM results by goal distance.

B FLEXIWORLD RESULTS BY GOAL DISTANCE

  • Table 5 reports FlexiWorld success rates by goal distance for Direct and ARCEM.Success allows one replan, and goal distance is measured in primitive steps.Source: B FLEXIWORLD RESULTS BY GOAL DISTANCE · p. 16, Table 5 · p. 16
  • The table reports means and standard deviations across three training seeds after averaging evaluation seeds.The last column averages distances within each training seed before calculating the standard deviation.Source: Table 5 · p. 16
  • The caption specifies success (%) with one replan as the reported outcome for the Direct and ARCEM distance breakdown.It does not provide the individual cell values in the supplied passage.Source: Table 5 · p. 16, B FLEXIWORLD RESULTS BY GOAL DISTANCE · p. 16

C.1 PLANNING WITH FIVE- AND TEN-ACTION CHUNKS

The supplementary experiments examine planning chunk length, search budget and temperature, goal-conditioned action prediction, and actor-free controls. Longer chunks preserve ARCEM’s average success in the matched study, while the other diagnostics report timing, sensitivity, and prediction or control comparisons.

C.1 PLANNING WITH FIVE- AND TEN-ACTION CHUNKS

  • 89.17% mean success with ARCEM at k = 5 is close to 89.29% at k = 10, though task-level changes are not uniform.PushT decreases from 68.89% to 67.14%, while TwoRoom increases from 96.61% to 97.72%.Source: §C.1 · pp. 16–17, Table 6 · p. 17

C.2 SEARCH BUDGET AND PLANNING TIME

  • 88.96% average success for ARCEM exceeds six-iteration Guarded-A at 87.44%; their matched mean solve times are 631.8 ms and 620.5 ms, respectively.The two planning times are within 1.8% on the matched input cases.Source: §C.2 · p. 18, Table 9 · p. 18, Figure 5 · p. 18

C.3 TEMPERATURE SENSITIVITY

  • At T = 0.9, PushT success is 58.14 ± 1.85%, below 68.89 ± 4.80% at T = 0.2; TwoRoom shows the same broad pattern.The authors note that adjacent settings do not establish a sharp optimum.Source: §C.3 · p. 19, Figure 6 · p. 19, Table 10

D.1 GOAL-CONDITIONED ACTION PREDICTIONS

  • INTACT is better at D = 25, while FlexiWorld is better at D ∈ {50, 75, 100} on all four mean action-prediction metrics.Both models predict the next five expert actions from matched observations, goals, and preceding expert chunks.Source: Figure 7 · p. 20, §D.1 · p. 20

D.2 FROZEN PROBES AND ACTOR-FREE PLANNING

  • The frozen temporal-distance probe has R2 of 0.603 ± 0.029 for FlexiWorld and 0.565 ± 0.026 for INTACT, predicting demonstrated time gaps rather than minimum control times.The probes measure information accessible to ridge regression, not representation quality for every downstream use.Source: §D.2 · p. 21, Figure 8 · p. 21, Table 11 · p. 21
  • With both learned actors disabled, pure CEM yields similar one-replan PushT success: 40.75% for FlexiWorld and 40.17% for INTACT.After the first plan, the corresponding success rates are 35.67% and 35.83%.Source: §D.2 · p. 21, Figure 8 · p. 21

D.3 WITHIN-MODEL ROLLOUT ERROR ACROSS CHUNK LENGTHS

Ten-action chunks reduce expert-action endpoint prediction error on PushT and Cube, but lower error does not consistently improve control. The section also frames chunk length and reobservation as planning and execution trade-offs.

D.3 WITHIN-MODEL ROLLOUT ERROR ACROSS CHUNK LENGTHS

  • The diagnostic compares k = 5 and k = 10 within the same checkpoint using identical expert actions and no intermediate observations.It measures endpoint latent MSE over the initial D-step horizon, not the later replan.Source: §D.3 · p. 22, Table 12 · p. 22
  • Ten-action chunks lower expert-action endpoint MSE at every tested distance on PushT and Cube, but PushT Direct success decreases while Cube success stays similar.On PushT, MSE changes from 0.1604 to 0.1541 at D = 50 and from 0.5797 to 0.5297 at D = 100.Source: §D.3 · p. 22, Table 12 · p. 22

E EXECUTION OUTCOMES AND PLANNING LIMITATIONS

E.1 SEPARATING REOBSERVATION FROM SEARCH

ARCEM and Direct differ in first-stage success, recovery after reobservation, and final outcomes across PushT and Cube. ARCEM gains substantially on PushT, while Cube results vary by goal distance.

E.1 SEPARATING REOBSERVATION FROM SEARCH

  • ARCEM has a net gain of 114 successes out of 1,200 PushT evaluations, succeeding in 183 cases where Direct fails and losing 69 Direct successes.At D = 25, ARCEM needs fewer second-stage rescues (5 versus 39) because it already succeeds in 283 rather than 230 first-stage evaluations.Source: §E.1 · p. 22, Figure 9 · p. 23
  • At PushT D = 100, first-stage successes rise from 67 to 91, and subsequent rescues rise from 32 to 52.Source: §E.1 · p. 23, Figure 9 · p. 23
  • On Cube, ARCEM succeeds in all 300 first-stage evaluations at D = 25, matching Direct’s 100% final success; across distances it gains 43 cases and loses 34.At D = 75, 16 evaluations improve and 18 regress despite the small aggregate change.Source: §E.1 · p. 23, Figure 9 · p. 23

E.2 A CONTROLLED SEARCH FAILURE IN TWOROOM

In the TwoRoom stress test, search can select plans that fail despite successful candidates being available, and bounded-action scoring improves two-stage success. The section also discusses training–deployment differences and illustrates varied rollout outcomes.

E.2 A CONTROLLED SEARCH FAILURE IN TWOROOM

  • Successful candidates exist in 1,775 of 1,800 evaluations, but selected plans succeed in only 878; 605 of 667 lost Direct successes involve wall collisions.The findings localize much of the failure to candidate ranking rather than a lack of successful proposals.Source: §E.2 · pp. 23–24
  • 94.44% versus 78.11% at D = 75, and 90.89% versus 69.33% at D = 100, are the two-stage success rates with bounded-action scoring versus the original score.The comparison keeps temperature and search budget fixed.Source: §E.2 · p. 24
  • In four cases where Direct succeeds and ARCEM fails, all selected ARCEM plans have lower predicted cost but worse physical proximity, with wall collisions on 9–20 execution steps.Figure 10 compares predicted costs and replayed distances; episode 55’s closest approach is 57.44 pixels, outside the 16-pixel success radius.Source: Figure 10 · p. 24, §E.2 · p. 24

E.3 TRAINING AND DEPLOYMENT CONDITIONS

  • Student Forcing trains the actor on generated within-chunk action prefixes while keeping latent state, intent, preceding-chunk context, and expert targets fixed.Training on model-generated contexts and corrective targets after physical perturbations are described as complementary ways to address training–deployment differences.Source: §E.3 · p. 24
  • Chunk length controls imagined rollout resolution, while reobservation supplies physical feedback; the proposed uncertainty-adaptive choices remain future work.The text suggests comparing fixed schedules on success, observation cost, and planning cost.Source: §E.3 · p. 24

E.4 QUALITATIVE ROLLOUTS ON PUSHT AND CUBE

  • The PushT and Cube examples at D = 75 and k = 5 include first-plan success, recovery after reobservation, and persistent failure after both plans.At replanning, the actor receives the last five executed action commands.Source: §E.4 · p. 24

Related papers

Loading 2609.35138v2…