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

Beyond the Next Step: Variable-Length Latent World Models for Long-Horizon Planning
Jun 19, 2026Tianqi Du, Qi Zhang, Yifei Wang +1Latent World Models Long-Horizon Planning
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
FF-JEPA: Long-Horizon Planning in World Models with Latent Planners
Jun 8, 2026Sergi Masip, Jonathan Swinnen, Yutong Hu +2Long-Horizon Planning Low-Latency Latent Planning
PSI, ESAT, KU Leuven · RAM, MECH, KU Leuven
Beyond a single latent space: a dual-latent world model for long-horizon planning
Sep 29, 2026Delin Zhao, Zhengrong Yue, Shaobin Zhuang +6Latent World Models Long-Horizon Planning
Nanjing University · Shenzhen University of Advanced Technology · Shanghai Jiao Tong University +1

Jun 19, 2026, cs.LGJ/K move · Enter open · S save

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

Tianqi Du, Qi Zhang, Yifei Wang, Yisen Wang

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

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

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