RIG: Synergizing Reasoning and Imagination in End-to-End Generalist Policy

🔼 Figure 2 illustrates the process of creating the training dataset for the RIG model. It begins with existing datasets: S0 uses refined data from MineRL-V0, while S1 collects vision-action pairs from STEVE-1. In S2, reasoning is added to the data using GPT-40. Crucially, S3 involves a comparative analysis. The trained RIG-basic and STEVE-1 are run in parallel, and only instances where RIG-basic underperforms STEVE-1 are retained for the dataset. This ensures that the model learns from its mistakes and improves its performance. Finally, S4 aligns the trajectories temporally, ensuring consistency.

read the caption

Figure 2: Illustration of the data collection pipeline (S0–S4). Note that at S3 (Vision-Reviewing), we run the trained RIG-basic and policy model (STEVE-1 [20]) in parallel, keeping instances where RIG-basic performs poorly compared to STEVE-1.

🔼 Figure 3 illustrates the multi-step reasoning process within RIG, the proposed model. RIG doesn’t just react to immediate observations; instead, it engages in a structured conversation with itself. This is represented by a series of interconnected steps. The key element is the use of the token ‘Imagine:’ to explicitly distinguish between real-world observations and internally generated scenarios. By employing this distinction, RIG maintains coherence in its reasoning, action prediction, and visual imagination processes. This methodology allows for a more thorough evaluation of prospective actions before taking real-world action.

read the caption

Figure 3: Inference process in RIG. RIG follows a structured conversation flow through multi-turn interactions. It consistently uses the fixed word Imagine: to clearly separate internally imagined scenarios from real observations, thereby guiding coherent reasoning, action prediction, and visual imagination.

🔼 Figure 4 presents a comparison of RIG-basic’s performance against other state-of-the-art methods in terms of both sample efficiency and overall performance on embodied tasks. RIG-basic demonstrates significantly better performance using only 111 hours of training data, significantly less than the others. The figure shows that RIG-basic requires far fewer training iterations to achieve comparable or superior performance compared to models like VPT, STEVE-1, Jarvis-1, and MineDreamer. The reduced training time is particularly noteworthy considering the superior performance obtained by RIG-basic. Note that MineDreamer is a hybrid system, requiring separate training for a visual generation model in addition to utilizing the VPT policy model, thereby increasing its total training data requirements.

read the caption

Figure 4: Performance and data-efficiency comparison. RIG-basic significantly outperforms other baselines with higher sample efficiency and achieves superior performance using only 111 hours of training data (42h S0 MineRL-V0 and 69h S1-S4). MineDreamer [48], a hybrid-system model, separately trains a visual generation model (139 hours) but also relies on VPT for the policy model, increasing total data requirements. Duration of VPT [2] reflects only the IDM data used, measured as video frames, while STEVE-1 [20] and Jarvis-1 [35] also leverage the VPT dataset.

🔼 Figure 5 presents a comparative analysis of RIG (Reasoning and Imagination Generalist) against several baseline models across four key aspects: embodied tasks, image generation, understanding, and reasoning. The results showcase RIG’s superior performance. RIG-basic, the simpler model using reasoning alone, already demonstrates improvements over the baselines. However, RIG-lookahead, incorporating both reasoning and a review mechanism, achieves the best results across all four evaluation categories. This visualization highlights the substantial improvements in performance achieved by synergizing reasoning and imagination.

read the caption

Figure 5: Comparison with various baselines across embodied tasks, generation, understanding, and reasoning. RIG-basic incorporates reasoning without reviewing, while RIG-lookahead integrates both reasoning and reviewing capabilities.

🔼 Figure 6 shows the scalability of the RIG model across training data, iterations, and inference steps. The top row displays the number of samples collected during material-gathering tasks for different scalability settings (e.g., varying training data, iterations, or lookahead steps), while the bottom row shows the success rate in exploration-based tasks under the same settings. Shaded areas indicate variance. The figure excludes the initial 42 hours of MineRL-V0 pretraining from the total 111 training hours. The training ratio only considers data before lookahead reasoning is applied.

read the caption

Figure 6: Scalability Evaluation Across Training, Iteration, and Inference. We evaluate the scalability of RIG by testing its performance across three different aspects: training scalability, iteration scalability, and inference scalability. Each column corresponds to a different scalability setting. The top row presents the number of collected samples in material-gathering tasks, while the bottom row reports the success rate in exploration-based tasks. Shaded regions represent variance. We exclude 42h MineRL-V0 pretraining from the total 111h in Figure. The training ratio is only counted before the lookahead reasoning.

🔼 The figure illustrates RIG’s inference process, which involves generating imagined visual states and corresponding reasoning to simulate multiple action trajectories. This allows the agent to conduct self-review and make corrective predictions before executing real-world actions. The process begins with an instruction, current visual observation, and the history of previous actions. RIG uses this to produce textual reasoning steps. Importantly, to separate imagined scenarios from real observations, RIG uses the token <Imagine:> before generating predicted future visual states (frames) and associated reasoning. These imagined scenarios are then used to refine the decision-making process. Finally, the next action and a prediction of the resulting frame are output.

read the caption

Figure A1: Detailed inference pipeline. RIG generates imagined visual states and corresponding reasoning to simulate multiple action trajectories, enabling self-review and corrective prediction.

🔼 This figure details the training pipeline of the RIG model, showing how it progressively incorporates reasoning and imagination. S0 and S1 stages pretrain the model using real and imagined trajectories. S2 and S3 further improve the model by adding reasoning and reviewing capabilities via GPT-40 relabeling. Finally, S4 aligns the model’s imagined future trajectories (dream flow) with actual environmental observations (environment-grounded traces), allowing the model to refine its predictions based on reality.

read the caption

Figure A2: Training pipeline of RIG. S0/S1 pretrain the model by aligning real and imagined flows. S2/S3 enhance reasoning and reviewing via GPT-4o relabeling. S4 aligns temporally predicted trajectories (dream flow) with environment-grounded traces.

🔼 Figure A3 illustrates the decision-making process of RIG, highlighting the benefits of lookahead reasoning and internal review. The agent first perceives and interprets the environment (steps 1 and 2). Then, it simulates potential future states (step 3) by using its imagination of how actions may affect the world. By leveraging its internal predictive model, it can anticipate a hidden hazard or undesirable outcome. After reviewing the simulated future (step 3), the agent refines its planned action and modifies the decision to avoid the anticipated hazard before acting (step 4). This mechanism enhances the robustness and reliability of the agent’s actions.

read the caption

Figure A3: Qualitative example of lookahead and review. The agent understands the environment (1–2), simulates future states (3), and refines its decision through internal review before acting (4), successfully avoiding a hidden hazard.

🔼 Figure A4 is a visualization showing the distribution of different types of embodied tasks present in the datasets used to train the RIG model. The tasks are categorized by complexity, ranging from simpler tasks like collecting resources to more complex tasks requiring strategic planning and higher-level reasoning. The diversity in task complexity is highlighted to emphasize the model’s ability to generalize well across a wide range of challenging scenarios in the Minecraft environment.

read the caption

Figure A4: Task distribution. Our datasets include various embodied tasks with varying complexity, ensuring strong generalization across downstream goals.

🔼 This figure compares the performance of RIG and GPT-4 in a Minecraft tree-chopping task. Both models receive the same visual input (the game screen) and the same instruction: ‘chop a tree.’ RIG demonstrates a more robust approach. It first reasons about the scene, identifies a suitable tree, and plans actions to reach and chop it, adjusting its camera position as needed. Then, it simulates the result of its plan (through visual imagination). GPT-4, while producing visually appealing output, makes a critical error in distance judgment, selecting an unreachable tree and attempting an invalid action. Because GPT-4 doesn’t use visual imagination, it fails to correct its plan after executing the incorrect action. The comparison highlights RIG’s ability to combine reasoning, visual imagination, and planning for reliable decision-making in dynamic environments.

read the caption

Figure A5: Case study comparison with GPT-4o. Given the same input and prompt (chop a tree), RIG reasons and imagines future states to choose a reachable tree and adjust position before acting. GPT-4o, despite high visual quality, misjudges the distance, executes an invalid action, and fails to revise its plan.

🔼 This table presents an ablation study analyzing the impact of different model components on the performance of the RIG model across various aspects. Specifically, it examines the effects of incorporating action prediction, image generation, explicit reasoning, and lookahead reasoning capabilities. Performance is evaluated using metrics such as FID and PSNR for image generation, Score-Stc. for static understanding on the STEVE-21K dataset, and Score-Env. for environment-specific and reasoning-specific understanding. The table allows for a systematic comparison of how each component contributes to the overall performance of the model.

read the caption

Table 2: Ablation study on Generation, Understanding, and Reasoning performance. We compare different combinations of Action, Generation (Gen.), Reasoning (Reason), and Lookahead capabilities. “Score-Env.” represents the environment-specific evaluation score from online understanding testing, while “Score-Env.” denotes reasoning-specific evaluation. “Score-Stc.” is computed on the static dataset STEVE-21K [43], and “FID” / “PSNR” measure image generation quality.

🔼 Table A1 compares different unified multimodal models, categorized by whether they use an autoregressive approach or a combination of autoregressive and diffusion methods. The table details key characteristics of each model, including the type of vision encoder used (e.g., VQGAN, VAE), the model’s parameter size (in billions), vision generation quality as measured by GenEval SD3 8B, the quality of multimodal understanding, and the evaluation benchmarks used to assess the model’s performance. This provides a comprehensive overview of the capabilities and performance of various state-of-the-art multimodal models.

read the caption

Table A1: Comparison of various unified multimodal methods, categorized by their training approach (Autoregressive and AR + Diffusion), detailing vision encoder type, parameter scale, vision generation quality (GenEval SD3 8B), multimodal understanding quality, and evaluation benchmarks.

🔼 Table A2 presents a comparison of various tokenizers’ performance across multiple image benchmark datasets. It shows the Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM), and Fréchet Inception Distance (FID) scores for each tokenizer on the MS-COCO, ImageNet-1K, FFHQ, and CelebA-HQ datasets. A separate PSNR score is also provided for Minecraft images. This allows for a comprehensive evaluation of the tokenizers’ ability to reconstruct and represent images effectively, highlighting their relative strengths and weaknesses for various image types and visual fidelity.

read the caption

Table A2: Comparison of Tokenizers across different benchmarks. PSNR, SSIM, and rFID are measured on MS-COCO, ImageNet-1K, FFHQ, and CelebA-HQ datasets. PSNR for Minecraft images is provided separately.

🔼 Table A3 provides a detailed comparison of RIG against existing methods for embodied AI, focusing on key architectural differences. It contrasts data sources used for training (e.g., videos, simulated experience), reward structures (sparse vs. dense), the type of observations provided to the agent (raw pixels, additional metadata), how actions are represented (discrete actions, keyboard inputs), the presence of planning mechanisms (iterative vs. non-iterative), and whether the agent uses a skill database for pre-trained actions. This comparison highlights RIG’s unique features and relative strengths among similar approaches.

read the caption

Table A3: Comparison between RIG (Ours), and existing works. This system-level comparison of LLM-based and RL-based methods focuses on data sources, reward setup, observation type, action representation, iterative planning, and skill database usage.