MAC-v2: Long-horizon model-based value expansion from pixels

AI summary

Overview

Research area: Offline reinforcement learning (RL) from image observations, model-based RL, latent world models, robotic manipulation.

Technical level: Advanced. The paper assumes familiarity with offline RL, value functions, latent dynamics models, flow-matching policies and rejection sampling.

Scope: MAC-v2 is an offline, pixel-based RL framework that combines a reconstruction-free latent world model with action chunking to make long-horizon model-based value expansion practical, and it is evaluated on visual OGBench, MMBench2, Robomimic and LIBERO-10.

What This Paper Is About

Offline RL agents must learn a policy from a fixed dataset without further environment interaction, and model-based methods try to help by imagining rollouts in a learned world model. Prior pixel-based offline model-based methods only imagine short rollouts (for example, LOMPO, OfflineDV2 and SeMOPO use an imagination horizon of 5 steps), which limits how far useful value information can propagate. The paper asks whether long-horizon model-based value expansion can work directly from pixels, where the agent must learn both the representation and its dynamics and where errors from either can compound recursively.

Key Contributions

1. A reconstruction-free latent world model for offline pixel-based RL. MAC-v2 encodes each image into a compact latent state and trains the encoder with dynamics-prediction, reward-prediction and behavior-cloning (flow-matching) gradients, rather than with image reconstruction. This avoids encoding control-irrelevant visual detail that hinders dynamics prediction.

2. Action chunking applied to both the policy and the world model. With chunk size n, an imagination horizon of Hn environment steps requires only H recursive model predictions, which reduces compounding prediction error. The paper reports that this enables stable value expansion over more than 100 environment steps, versus the 5-step horizons of prior pixel-based model-based methods.

3. A full long-horizon latent value-expansion procedure. The method builds on MAC (state-based) and adapts it to latent space: a flow-matching behavioral policy with rejection sampling over M candidate action chunks, a distributional action-chunked reward model with two-hot encoding, a distributional value function trained on nH-step targets from model rollouts, and an action-chunked Q-function used for rejection sampling.

4. Broad empirical evaluation and analysis. Evaluation covers 25 visual OGBench tasks, 10 MMBench2 tasks, 3 Robomimic tasks and 10 LIBERO-10 tasks, plus ablations on rollout horizon, chunk size, encoder-gradient choices and alternative representation objectives.

Main Findings

• Strong OGBench results: MAC-v2 outperforms all baselines in 4 of 5 OGBench task families and, in the authors' description, matches QC on the fifth. On visual-puzzle-3x3 MAC-v2 reaches 95±4 where the best prior method scores 22±1, and prior model-based RL methods score near zero across OGBench tasks. Other MAC-v2 scores: visual-cube-single 74±1, visual-cube-double 53±3 (best baseline 13±1), visual-scene 72±12 (best baseline 57±2), visual-puzzle-4x4 46±5 (best baseline 20±2).

• Mixed but competitive MMBench2 results: MAC-v2 is best on the ogbench (33±4 vs. FQL 30±4) and pygame (66±2 vs. FQL 58±1) domains, but below baselines on dmcontrol (31±4 vs. FQL 53±1) and maniskill (4±2 vs. LOMPO 23±3).

• Best RL method on Robomimic: MAC-v2 scores 100±0 on lift, 95±3 on can and 72±10 on square, versus BC at 100±0, 97±3 and 68±21. The authors describe MAC-v2 as the only RL method that matches or exceeds behavior cloning across all evaluated tasks; the table shows it improving on the challenging square task.

• Competitive with behavior cloning on LIBERO-10: MAC-v2 scores 80±1, equal to BC (80±1) and above QC (77±1) and QC-IFQL (78±3), while LOMPO (1±2) and OffDV2 (0±0) score near zero.

• Horizon matters: On puzzle-4x4-task5 with chunk size 10, performance improves as the rollout horizon increases across {10, 50, 100} environment steps. On puzzle-4x4-task1, performance also improves as the action chunk size increases across {1, 5, 10}.

• Baselines do not benefit from longer horizons: Extending the value-estimation horizon of QC (via chunk size) and OfflineDV2 (via rollout length) on puzzle-4x4-task5 does not produce comparable improvements, suggesting that simply lengthening the horizon is insufficient without MAC-v2's design choices.

• Both encoder objectives are needed: Removing the behavior-cloning gradient or the world-modeling gradient to the encoder on cube-double-task3 substantially degrades performance.

• Conventional representation objectives fall short: Replacing MAC-v2's encoder objective with pixel reconstruction, SimCLR or LeWorldModel on cube-double-task3 substantially degrades performance.

• Honest limitation noted: On LIBERO-10, the authors attribute the lack of improvement over BC to the data being largely unimodal demonstrations, leaving limited room for rejection sampling to improve the behavioral policy.

Methodology in Plain English

MAC-v2 treats a picture as an input to a learned encoder that compresses it into a small latent vector, deliberately without ever asking the model to redraw the image. That latent vector is trained to be useful in three ways at once: it must let a dynamics model predict the latent state n steps ahead, it must let a reward model predict the discounted reward over those n steps, and it must let a flow-matching policy produce the actions that the dataset actually took. Because all three losses send gradients into the encoder, the representation keeps what is needed for prediction and control while dropping visual detail that is irrelevant to the task.

The policy does not output one action at a time; it outputs a chunk of n actions. The dynamics model also jumps n steps at a time. So to imagine 100 environment steps with n = 10, the model only needs 10 recursive predictions instead of 100, which sharply limits how much error can compound. To turn the behavioral policy into a better policy, the method draws several candidate action chunks from the flow policy using different noise samples and picks the one the learned Q-function scores highest (rejection sampling). Values are then learned entirely from imagined latent trajectories: starting from an encoded dataset observation, the policy, dynamics model and reward model generate a rollout of length H, and the value and Q functions are trained on the multi-step returns from that rollout, using distributional (two-hot) heads and an exponentially averaged target network.

Why This Matters

Impact on research: The paper shows that the barrier to long-horizon model-based value expansion from pixels is not the idea itself but the representation and the prediction granularity. It provides evidence that a reconstruction-free latent encoder trained with control objectives, plus action chunking, is sufficient to extend imagination horizons from 5 steps to more than 100, and it documents that neither ingredient alone is enough. This connects state-based long-horizon model-based RL (MAC) with reconstruction-free latent world models from online RL.

Real-world applications (as suggested by the paper's framing of robotic manipulation):

• Learning manipulation skills from previously collected camera footage without risky trial-and-error on hardware.
• Multi-step pick-and-place and rearrangement tasks where the consequence of an action appears many steps later.
• Settings where expert demonstrations exist but are expensive, so value-based selection must improve on pure imitation.
• Deploying policies from visual observations in labs or factories where additional data collection is costly.

Industry relevance: Offline learning from pixels is attractive for robotics and automation, where online exploration is slow or unsafe. MAC-v2's use of action chunks also matches how many modern robot policies are deployed, and the authors point to vision-language-action models as a natural extension. The reported compute cost is modest: the paper states roughly 9 hours for OGBench and MMBench2 benchmarks, 12 hours for Robomimic and 14 hours for LIBERO-10 on a single A5000 GPU.

Future Directions

• Adaptive or hierarchical action chunking. Larger chunks reduce recursive predictions, but excessively long chunks reduce the policy's reactivity; the authors propose adaptive or hierarchical chunking to trade off long-horizon prediction against closed-loop control.
• Beyond simulated single-task settings. The experiments are limited to simulated, single-task environments with relatively limited visual variation, leaving a gap to realistic multi-task domains.
• Scaling to large vision-language-action models (VLAs) or world action models (WAMs) as a route to long-horizon model-based value expansion in more diverse and realistic settings.
• Improving gains on near-expert, unimodal datasets. The LIBERO-10 results suggest that when demonstration data is largely unimodal, rejection sampling has little room to improve over the behavioral policy, which is an open problem for value-based methods.

Target Audience

Researchers and graduate students working on offline RL, model-based RL and robot learning from images; practitioners who want to add value-based policy improvement on top of imitation-learned visuomotor policies; and readers already comfortable with latent dynamics models, flow-matching policies and rejection sampling who want to understand how to make imagined rollouts stay accurate over long horizons.

Authors’ abstract

We study whether \emph{long-horizon model-based value expansion} can serve as an effective framework in pixel-based offline reinforcement learning (RL). Long-horizon model-based value expansion from pixels is substantially harder than state-based environments, as the agent must learn both the representation and its dynamics, and errors from either can accumulate recursively over imagined trajectories. We introduce MAC-v2, which enables long-horizon model-based expansion from pixels by training a reconstruction-free latent world model with action chunking. By learning representations from prediction and control objectives, MAC-v2 does not need to encode control-irrelevant visual details that can hinder the dynamics prediction. Action chunking further mitigates the prediction error by reducing the number of recursive model predictions required for value expansion. Across diverse visual benchmarks including visual OGBench, MMBench2, Robomimic, and LIBERO-10, MAC-v2 achieves strong performance against model-free and model-based offline RL baselines. Our analyses show that long-horizon model-based value expansion is critical for the performance of MAC-v2, and that both action chunking and latent learning objectives are essential for reliable long-horizon latent imagination.