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

cs.LGSep 28, 2026

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

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

Organizations: University of Toronto · Zhejiang University · Tencent Jarvis Lab · Mila & Université de Montréal · Samsung SAIL · Tsinghua University

Abstract

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.

Figures & tables

Figure 1: FlexiWorld: flexible action chunks across multiple time scales. Training and planning interfaces (top), four-benchmark mean success (bottom right), and selected PushT configurations (bottom left). Tables 1 and 4 detail the comparisons.

Figure 2: Joint prediction and action learning in FlexiWorld. Training samples different goal spans and partitions each window into variable-length action chunks paired with their boundary observations. Chunk embeddings condition the predictor on the actions leading to the next boundary latent state. The actor is supervised by expert actions, conditioned on local or final-goal intents and expert or generated action prefixes. Training combines prediction MSE, SIGReg, and action negative log-likelihood (NLL). All modules, including the observation encoder and predictor, are trained jointly.

Figure 3: ARCEM refinement. Residuals modify actions and subsequent actor conditions. Chunk embeddings and predicted states propagate these changes; arrival costs select elites to refit the residual distribution.

Figure 4: Direct planning at increasing goal distances. Direct success with one replan across four tasks. Goal distance is in primitive steps; bar labels round success to the nearest percent.

Variant	Variable chunks	Mixed spans	SF	Success (%) $\uparrow$
Baseline	–	–	–	$46.58\!\pm\!1.81$
Baseline + 75-step span	–	–	–	$40.92\!\pm\!0.80$
Baseline + mixed spans	–	$\checkmark$	–	$54.42\!\pm\!0.38$
New architecture	–	–	–	$46.08\!\pm\!2.43$
Fixed chunks + SF	–	–	$\checkmark$	$48.42\!\pm\!0.38$
Variable chunks	$\checkmark$	–	–	$50.83\!\pm\!1.38$

Table 3: PushT component study. We compare INTACT baselines with our encoder and actor variants. Direct success averages four distances with up to one replan; results are mean

\pm

SD across three evaluation seeds (training seed 0).

\checkmark

/– mark enabled/disabled; SF denotes Student Forcing (

p_{\mathrm{SF}}=0.5

). Mixed spans use 35/55/75-step goals for two epochs; single spans use 35-step goals (75 where indicated) for six. Appendix A.2 details evaluation provenance.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Variant	Goal span	Action interface	$p_{\mathrm{SF}}$	Epochs	Success (%)
New architecture	35	Fixed, AR	0	6	$46.08\!\pm\!2.43$
Fixed chunks + SF	35	Fixed, AR	0.5	6	$48.42\!\pm\!0.38$
Variable chunks	35	Variable, AR	0	6	$50.83\!\pm\!1.38$
Variable chunks + SF	35	Variable, AR	0.5	6	$52.00\!\pm\!0.50$
No action feedback	35	Variable, parallel	0	6	$48.83\!\pm\!1.94$
INTACT	35	Fixed	–	6	$46.58\!\pm\!1.81$

Appendix

Table 4: Extended PushT component study. Success (%) averages four goal distances with one replan. Sample SD is across three evaluation seeds; training settings follow Appendix A.2 . Goal spans are in primitive steps. AR denotes autoregressive generation; SF denotes Student Forcing, with generated-prefix probability

p_{\mathrm{SF}}

Planner	Task	$D=25$	$D=50$	$D=75$	$D=100$	Mean
Direct	PushT	$89.67\!\pm\!1.67$	$71.56\!\pm\!3.98$	$46.89\!\pm\!5.36$	$33.44\!\pm\!5.35$	$60.39\!\pm\!3.92$
Direct	Cube	$100.00\!\pm\!0.00$	$84.67\!\pm\!2.40$	$87.00\!\pm\!3.06$	$93.78\!\pm\!0.84$	$91.36\!\pm\!1.56$
Direct	Reacher	$98.11\!\pm\!0.19$	$99.89\!\pm\!0.19$	$99.56\!\pm\!0.19$	$98.67\!\pm\!0.33$	$99.06\!\pm\!0.05$
Direct	TwoRoom	$100.00\!\pm\!0.00$	$99.11\!\pm\!0.51$	$96.44\!\pm\!3.36$	$89.89\!\pm\!3.89$	$96.36\!\pm\!1.92$
ARCEM	PushT	$94.33\!\pm\!2.08$	$78.11\!\pm\!4.55$	$55.00\!\pm\!8.01$	$48.11\!\pm\!5.35$	$68.89\!\pm\!4.80$
ARCEM	Cube	$100.00\!\pm\!0.00$	$88.78\!\pm\!2.52$	$86.22\!\pm\!3.02$	$92.78\!\pm\!1.02$	$91.94\!\pm\!1.50$

Appendix

Table 5: FlexiWorld success by goal distance. Success (%) allows one replan. Entries show mean and SD across three training seeds after averaging evaluation seeds.

D

is measured in primitive steps. The last column averages distances within each training seed before computing SD.

Planner	$k$	PushT	Cube	Reacher	TwoRoom	Average
Direct	5	$60.39\!\pm\!3.92$	$91.36\!\pm\!1.56$	$99.06\!\pm\!0.05$	$96.36\!\pm\!1.92$	$86.79\!\pm\!1.65$
Direct	10	$54.39\!\pm\!1.35$	$91.50\!\pm\!1.32$	$99.25\!\pm\!0.25$	$96.22\!\pm\!1.88$	$85.34\!\pm\!0.92$
ARCEM	5	$68.89\!\pm\!4.80$	$91.94\!\pm\!1.50$	$99.72\!\pm\!0.17$	$96.61\!\pm\!2.45$	$89.29\!\pm\!1.96$
ARCEM	10	$67.14\!\pm\!1.82$	$92.11\!\pm\!1.93$	$99.69\!\pm\!0.05$	$97.72\!\pm\!1.42$	$89.17\!\pm\!1.08$
ARCEM	5 †	$68.50\!\pm\!4.13$	$91.72\!\pm\!1.88$	$99.69\!\pm\!0.17$	$96.83\!\pm\!2.17$	$89.19\!\pm\!1.85$

Appendix

Table 6: Planning chunk length and control success.

k

is the chunk length in primitive actions. Success (%) with one replan averages four goal distances. Mean and sample SD across three training seeds, after averaging evaluation seeds and distances. All rows use the matched chunk-length study; they do not replace the main-result evaluation.

\dagger

: common scoring grid, retaining five-action rollouts but scoring only the ten-action boundary times.

Planner	$D$	SR, $k=5$	SR, $k=10$	ms, $k=5$	ms, $k=10$	Speedup
Direct	25	$96.94\!\pm\!0.38$	$96.86\!\pm\!0.25$	64.2	55.1	1.17 $\times$
Direct	50	$88.81\!\pm\!1.18$	$87.28\!\pm\!0.84$	115.6	92.8	1.25 $\times$
Direct	75	$82.47\!\pm\!2.70$	$81.00\!\pm\!1.30$	167.2	135.1	1.24 $\times$
Direct	100	$78.94\!\pm\!2.50$	$76.22\!\pm\!1.92$	218.4	172.6	1.27 $\times$
ARCEM	25	$98.53\!\pm\!0.55$	$98.86\!\pm\!0.38$	268.8	221.5	1.21 $\times$
ARCEM	50	$91.36\!\pm\!1.99$	$90.75\!\pm\!0.58$	510.9	393.6	1.30 $\times$

Appendix

Table 7: Matched success and planning time by goal distance. SR denotes success rate (%), averaged over four tasks with training-seed SD.

D

is goal distance and

k

is chunk length, both in primitive steps. Planning time averages per-input median solve times on identical inputs, including encoding but excluding environment execution. Speedup is the ratio of mean planning times, not an episode-runtime ratio.

Task	$D$	Direct, $k=5$	Direct, $k=10$	ARCEM, $k=5$	ARCEM, $k=10$
PushT	25	$89.67\!\pm\!1.67$	$88.00\!\pm\!1.20$	$94.33\!\pm\!2.08$	$95.44\!\pm\!1.50$
PushT	50	$71.56\!\pm\!3.98$	$64.56\!\pm\!1.39$	$78.11\!\pm\!4.55$	$74.44\!\pm\!2.83$
PushT	75	$46.89\!\pm\!5.36$	$41.22\!\pm\!0.69$	$55.00\!\pm\!8.01$	$56.44\!\pm\!1.50$
PushT	100	$33.44\!\pm\!5.35$	$23.78\!\pm\!3.91$	$48.11\!\pm\!5.35$	$42.22\!\pm\!5.42$
Cube	25	$100.00\!\pm\!0.00$	$100.00\!\pm\!0.00$	$100.00\!\pm\!0.00$	$100.00\!\pm\!0.00$
Cube	50	$84.67\!\pm\!2.40$	$85.67\!\pm\!2.40$	$88.78\!\pm\!2.52$	$89.22\!\pm\!4.60$

Appendix

Table 8: Chunk-length success by task and goal distance. Matched Direct and ARCEM evaluations with one replan.

D

is goal distance and

k

is chunk length, both in primitive steps. Entries give mean success (%) and sample SD across three training seeds after averaging evaluation seeds.

Figure 5: Planning time by goal distance. FlexiWorld with Direct, ARCEM (three iterations), and Guarded-A (three or six iterations). Curves average per-input median times over both solve stages, including image encoding and planning. The six-iteration run uses the same GPU and input cases. The vertical axis shows per-solve time on a logarithmic scale.

Planner	Candidates	PushT	Cube	Reacher	TwoRoom	Average	Time (ms)
Direct	0	$60.17\!\pm\!1.04$	$89.67\!\pm\!1.84$	$99.08\!\pm\!0.52$	$96.42\!\pm\!1.01$	$86.33\!\pm\!0.55$	141.3
Guarded-A	384	$63.50\!\pm\!2.14$	$88.67\!\pm\!2.27$	$98.08\!\pm\!1.01$	$97.83\!\pm\!0.38$	$87.02\!\pm\!0.49$	414.1
Guarded-A †	768	$65.83\!\pm\!2.36$	$88.75\!\pm\!1.56$	$97.17\!\pm\!1.04$	$98.00\!\pm\!0.50$	$87.44\!\pm\!0.76$	620.5
ARCEM	384	$69.67\!\pm\!2.38$	$90.42\!\pm\!2.04$	$99.58\!\pm\!0.38$	$96.17\!\pm\!1.38$	$88.96\!\pm\!0.67$	631.8

Appendix

Table 9: Planning success and time with FlexiWorld. Success (%) averages four goal distances with one replan; Average also averages four tasks. Success evaluations share the seed-0 checkpoints, the math backend of scaled dot-product attention (Math SDPA), and batch size one; mean and SD are across the three evaluation seeds. Candidates gives the total search budget per solve across all iterations. Planning time uses identical inputs on one RTX PRO 6000.

\dagger

denotes six iterations of 128 candidates.

$T$	PushT	TwoRoom	$T$	PushT	TwoRoom
0.10	$68.81\!\pm\!3.81$	$96.36\!\pm\!2.33$	0.55	$62.14\!\pm\!3.89$	$92.00\!\pm\!3.57$
0.15	$68.42\!\pm\!3.37$	$96.42\!\pm\!2.25$	0.60	$62.56\!\pm\!2.38$	$90.28\!\pm\!3.43$
0.20	$68.89\!\pm\!4.80$	$96.61\!\pm\!2.45$	0.65	$61.56\!\pm\!3.26$	$88.64\!\pm\!2.77$
0.25	$67.39\!\pm\!4.48$	$96.47\!\pm\!2.51$	0.70	$60.72\!\pm\!3.17$	$87.06\!\pm\!3.18$
0.30	$67.06\!\pm\!3.38$	$95.92\!\pm\!3.06$	0.75	$60.83\!\pm\!2.92$	$86.14\!\pm\!3.67$
0.35	$65.53\!\pm\!3.73$	$95.75\!\pm\!2.93$	0.80	$59.78\!\pm\!3.13$	$84.83\!\pm\!3.00$

Appendix

Table 10: Dense ARCEM temperature sweep.

T

scales action residuals in normalized coordinates. Success (%) with one replan at

T=0.10,0.15,\ldots,0.90

. Mean and sample SD across three training seeds after averaging three evaluation seeds and four goal distances; 100 episodes per cell. All 1,224 cells are complete, using additive residuals without actor-standard-deviation scaling and the main candidate-scoring rule.

Figure 6: ARCEM temperature sensitivity at 0.05 resolution. Markers show measured settings from

T=0.10

0.90

, with lines connecting the means and bands showing training-seed SD after averaging evaluation seeds and goal distances. Dashed lines mark the main setting

T=0.2

. Both panels use the main candidate-scoring rule and one replan.

Figure 7: Goal-conditioned action correspondence on PushT. Lines show means and shaded bands show sample SD across three training seeds. Hollow points show individual checkpoints, slightly offset horizontally for visibility. Both models predict the next five expert actions from matched observations, goals, and preceding expert chunks. Arrows indicate the preferred direction for each metric. INTACT is better at

D=25

, whereas FlexiWorld is better at

D\in\{50,75,100\}

on all four mean metrics.

Model	$D$	Probe $R^{2}$	Effective rank	Latent std.
INTACT	25	$0.332\!\pm\!0.016$	$81.3\!\pm\!9.5$	$0.981\!\pm\!0.015$
INTACT	50	$0.346\!\pm\!0.025$	$79.3\!\pm\!9.3$	$0.980\!\pm\!0.017$
INTACT	75	$0.373\!\pm\!0.003$	$75.1\!\pm\!9.0$	$0.976\!\pm\!0.015$
INTACT	100	$0.463\!\pm\!0.013$	$72.2\!\pm\!8.4$	$0.990\!\pm\!0.015$
FlexiWorld	25	$0.365\!\pm\!0.017$	$53.0\!\pm\!12.2$	$0.976\!\pm\!0.011$
FlexiWorld	50	$0.359\!\pm\!0.005$	$52.0\!\pm\!12.1$	$0.973\!\pm\!0.013$

Appendix

Table 11: Frozen-feature probes and representation statistics on PushT. Entries report mean and SD across three training seeds on matched samples. The action probe predicts five expert primitives. Effective rank and latent standard deviation (std.) describe the current-state representations.

R^{2}

is the coefficient of determination;

D

is goal distance in primitive steps.

Figure 8: Frozen-feature probes on PushT. Left: temporal-gap MAE. Right: action

R^{2}

from frozen-feature ridge probes. Curves show means and sample SD across three training seeds, using each diagnostic’s matched sample set. The temporal probe predicts demonstrated time gaps, not minimum time to a goal. Neither probe updates the world model.

Task	$D$	MSE, $k=5$	MSE, $k=10$	Relative change
PushT	25	$0.0408\!\pm\!0.0017$	$0.0391\!\pm\!0.0023$	-4.13%
PushT	50	$0.1604\!\pm\!0.0180$	$0.1541\!\pm\!0.0120$	-3.90%
PushT	75	$0.3532\!\pm\!0.0227$	$0.3108\!\pm\!0.0358$	-12.00%
PushT	100	$0.5797\!\pm\!0.0243$	$0.5297\!\pm\!0.0388$	-8.61%
Cube	25	$0.0315\!\pm\!0.0017$	$0.0177\!\pm\!0.0013$	-43.70%
Cube	50	$0.0356\!\pm\!0.0027$	$0.0278\!\pm\!0.0005$	-21.93%

Appendix

Table 12: Expert-action endpoint prediction under two chunk schedules. Latent mean squared error (MSE) after a

D

-step rollout (

k

: chunk length in primitive actions), with identical expert primitive actions and no intermediate observations. Mean and sample SD across three training seeds; each seed averages 300 matched examples per task and distance. Changes compare

k=10

with

k=5

within the same model; latent errors are not pooled across tasks.

Figure 9: When reobservation and ARCEM help. Top: solid bars show first-stage success and hatched extensions show additional success after reobservation. Within each distance, the left bar is Direct and the right bar is ARCEM. Bottom: positive counts are ARCEM-only successes and negative counts are Direct-only successes after both stages. Each distance contains 300 paired evaluations.

Figure 10: Search can discard a successful Direct plan. Predicted costs and closest replayed distances for four

D=75

cases where Direct succeeds and ARCEM at

T=0.8

fails. Lower predicted cost is preferred. Distances use complete first-plan replay without reobservation. The dashed line marks the 16-pixel success radius.

Figure 11: PushT: success, recovery, and remaining error. FlexiWorld Direct trajectories. The middle row places the object near its goal in the first plan, then moves the agent closer to its target before success. The bottom row remains unsuccessful after replanning. A trajectory that ends early retains its terminal frame.

Figure 12: Cube: successful placement and incomplete recovery. FlexiWorld Direct trajectories. The middle row succeeds after reobservation, while the bottom row leaves the cube away from its goal after both plans. Terminal frames are held as in Figure 11 .

Explore similar work

Jun 19, 2026cs.LG

Beyond the Next Step: Variable-Length Latent World Models for Long-Horizon Planning

Recently, world models have emerged as a promising paradigm for building intelligent agents by learning predictive models that estimate future environment states conditioned on observations and actions. In particular, JEPA-style latent world models provide an efficient alternative to pixel space prediction by learning action-conditioned dynamics in compact representation spaces. However, existing latent world models typically rely on one-step prediction and must be recursively rolled out for long-horizon planning, which leads to compounding errors and a mismatch between training objectives and downstream planning tasks. To address this limitation, we propose Variable-length Latent World Models (VLWMs), a framework that learns to predict future latent states conditioned on action sequences of variable lengths. Instead of training only on one-step transitions, VLWMs directly model temporally extended dynamics, allowing the same predictor to evaluate action plans over different horizons. We further introduce a curriculum training strategy that progressively expands the action horizon, stabilizing optimization from short-range dynamics to long-range prediction. At test time, we design planning methods tailored to VLWMs to better exploit their variable-length predictive capabilities. Experiments on long-horizon control tasks show that VLWMs significantly improve latent space world models, achieving 13% average improvement over the state-of-the-art LeWM across different datasets, with especially large gains on tasks requiring extended planning. These results suggest that VLWM provides a simple yet effective paradigm for improving long-horizon prediction and planning in latent world models.

Tianqi Du, Qi Zhang, Yifei Wang +1

State Key Lab of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · Amazon AGI SF Lab · Institute for Artificial Intelligence, Peking University

Latent World Models Long-Horizon Planning Intermediate Latent States Long-Horizon Task Planning Next

Ask PDF

Jun 8, 2026cs.AI

FF-JEPA: Long-Horizon Planning in World Models with Latent Planners

Joint Embedding Predictive Architectures (JEPAs) have shown promising world modeling capabilities, enabling planning in latent space by optimizing action trajectories using methods like the Cross-Entropy Method (CEM). These methods are, however, too computationally expensive and ineffective for long-horizon planning. Furthermore, these methods typically require an explicit image of the goal state, which is not always possible in real-world tasks. In this work, we tackle these limitations by proposing Forward-Forward-JEPA (FF-JEPA), a hierarchical approach leveraging two forward dynamics models. Alongside a standard action-conditioned forward model, we introduce an action-free latent planner that predicts the next subgoal given the current state. This approach removes the need for goal images and enables long-horizon planning by decomposing complex trajectories into a sequence of tractable, short-term optimization problems. Preliminary results on PushT demonstrate that FF-JEPA successfully overcomes flat world models' long-horizon collapse, highlighting this approach as a promising direction for goal-free planning.

Sergi Masip, Jonathan Swinnen, Yutong Hu +2

PSI, ESAT, KU Leuven · RAM, MECH, KU Leuven

Long-Horizon Planning Low-Latency Latent Planning Latent World Models Planning World Models S-Jepa

Ask PDF

Sep 29, 2026cs.AI

Beyond a single latent space: a dual-latent world model for long-horizon planning

Latent world models often struggle with long-horizon planning despite accurate short-term predictions. Recursive rollouts accumulate errors, while distance concentration in high-dimensional latent spaces can weaken goal discrimination. We introduce the Dual-Latent World Model (Dual-WM), which separates local execution and long-range planning through distinct state representations and dynamics models. The low-level model predicts action-conditioned transitions, while the high-level model uses learned macro-actions to plan over longer temporal spans. We also propose Long-Horizon Representation Learning with Weighted Rollout (LoRe), which supervises self-generated predictions at both levels. An analysis of recursive error propagation motivates exponential horizon weights with separate decay rates for the two temporal scales. During planning, the high-level model generates latent subgoals that the low-level model refines into actions for precise execution. We evaluate from-scratch Dual-WM on five goal-conditioned visual control tasks against the task-wise strongest baselines without actor-guided proposals. At goal offsets of 50 and 100 environment steps, mean success increases from 75.9% to 84.4% and from 61.4% to 69.5%, respectively. At offset 100, Dual-WM outperforms these baselines on all five tasks and improves mean success over LeWM by 30.8 percentage points. Ablations and supporting analyses provide evidence of more informative representations for goal evaluation and greater consistency under recursive prediction. These results highlight the value of separating temporal roles and training across multiple horizons for reliable latent planning. Our core implementation is available at https://github.com/DeLin1001/Dual-WM-Official.

Delin Zhao, Zhengrong Yue, Shaobin Zhuang +6

Nanjing University · Shenzhen University of Advanced Technology · Shanghai Jiao Tong University +1

Latent World Models Long-Horizon Planning Latent Space Planning

Ask PDF