PaperlayerView exact revision on arXiv

Source-linked AI summary

IRASim: A Fine-Grained World Model for Robot Manipulation

Fangqi Zhu, Hongtao Wu, Song Guo, Yuxiao Liu, Chilam Cheang, Tao Kong

arXiv:2406.14540v2cs.ROcs.AIcs.CV

TL;DR

Robot-manipulation world models must represent fine-grained interactions and align each action with its corresponding video frame. IRASim uses a diffusion transformer with frame-level action conditioning to generate trajectory-conditioned videos. It outperforms video-generation baselines, correlates with ground-truth simulator policy evaluations, and improves Push-T planning IoU from 0.637 to 0.961.

  • Problem

    Robot-manipulation world models must capture delicate robot-object interactions, while trajectory-conditioned video generation requires precise alignment between each action and its corresponding predicted frame.

    Source: §1 · pp. 1–2

  • Method

    IRASim is a diffusion-transformer world model that conditions video generation on historical observations and robot action trajectories, using frame-level action conditioning to align actions with frames.

    Source: §1 · p. 2, §3.3 · p. 5

  • Results

    IRASim outperforms video-generation baselines across four datasets, policy evaluations strongly correlate with the ground-truth simulator, and Push-T planning raises vanilla diffusion policy IoU from 0.637 to 0.961.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    IRASim is reported as a model for policy evaluation and model-based planning, and it supports virtual-robot control from keyboard- or VR-collected trajectories.

    Source: §1 · p. 3, §5 · p. 14

    Video generation with IRASim is not real-time.

    Source: §5 · p. 14

Abstract

from arXiv · show

World models allow autonomous agents to plan and explore by predicting the visual outcomes of different actions. However, for robot manipulation, it is challenging to accurately model the fine-grained robot-object interaction within the visual space using existing methods which overlooks precise alignment between each action and the corresponding frame. In this paper, we present IRASim, a novel world model capable of generating videos with fine-grained robot-object interaction details, conditioned on historical observations and robot action trajectories. We train a diffusion transformer and introduce a novel frame-level action-conditioning module within each transformer block to explicitly model and strengthen the action-frame alignment. Extensive experiments show that: (1) the quality of the videos generated by our method surpasses all the baseline methods and scales effectively with increased model size and computation; (2) policy evaluations using IRASim exhibit a strong correlation with those using the ground-truth simulator, highlighting its potential to accelerate real-world policy evaluation; (3) testing-time scaling through model-based planning with IRASim significantly enhances policy performance, as evidenced by an improvement in the IoU metric on the Push-T benchmark from 0.637 to 0.961; (4) IRASim provides flexible action controllability, allowing virtual robotic arms in datasets to be controlled via a keyboard or VR controller.

1 Introduction

Robot manipulation world models must capture delicate robot-object interactions and align action trajectories with individual predicted frames. IRASim addresses this with frame-level action conditioning and supports video generation, policy evaluation, planning, and flexible control.

1 Introduction

  • The task is challenging because subtle interaction variations can cause failure, while action chunking specifies trajectories rather than single actions.Each action describes robot movement for a corresponding frame, unlike high-level text prompts that provide general context.Source: §1 · pp. 1–2
  • IRASim uses a diffusion transformer with a frame-level action-conditioning module to align each action with its corresponding video frame.It can be rolled out autoregressively to generate long-horizon videos while maintaining temporal consistency across clips.Source: §1 · pp. 2–3
  • IRASim outperforms comparison methods on all four datasets, is preferred in human evaluation, and scales with model size and computation.It generates videos up to 288×512 resolution and more than 150 frames long.Source: §1 · pp. 2–3
  • On Push-T, IRASim-based planning raises vanilla diffusion policy IoU from 0.637 to 0.961; policy evaluations also strongly correlate with the ground-truth simulator.The paper reports planning gains in simulation and real-world settings, with improvement scaling with test-time computation.Source: §1 · pp. 2–3
  • IRASim also demonstrates flexible action control by generating videos from trajectories collected with a keyboard or VR controller.The demonstrations control virtual robots in the datasets.Source: §1 · p. 3

2 Related Work

Related work applies world models to prediction, planning, and learning across domains, while robot-manipulation studies use action-conditioned video models and video-generation controls. IRASim is positioned around trajectory-conditioned, fine-grained interaction modeling and generative rollouts for planning and evaluation.

2 Related Work

  • World models support prediction and planning across driving, gaming, and manipulation, while IRASim uses predicted future videos to assess action rewards and select actions for execution.Prior work also combines action-conditioned video prediction with model-predictive control and uses physical simulators or real robots for policy evaluation.Source: §2 · p. 3
  • Prior robot-manipulation models predict videos from actions, while iVideoGPT uses an autoregressive transformer and VLP and UniSim prompt text-to-video models with language and action information.IRASim is distinguished by high-resolution, long-horizon videos conditioned on an initial observation and robot trajectory, with fine-grained interaction details.Source: §2 · p. 3
  • Video-generation controls include strokes, bounding boxes, and human poses; IRASim instead models complex 2D and 3D actions over time with frame-level action conditioning.The cited approaches use these control types to condition video synthesis.Source: §2 · p. 3

3 Methods

The methods define trajectory-to-video prediction from historical robot observations and an action sequence, then implement it with latent-space diffusion. IRASim combines historical-frame conditioning with frame-level action conditioning to support consistency and precise action-frame alignment.

3.1 Problem Statement

  • The trajectory-to-video task predicts a robot executing actions a_t:t+n from historical observation images I_t−h:t; a typical robot-arm action has seven degrees of freedom.Those degrees of freedom comprise three translations, three rotations, and one gripper action.Source: §3.1 · p. 4

3.2 Preliminaries

  • Diffusion adds Gaussian noise over forward timesteps and reverses the process from sampled noise to recover the data, with conditional generation controlled by c.The paper performs diffusion in VAE-compressed latent space to reduce the computation required for long, high-resolution videos.Source: §3.2 · pp. 4–5
  • A pretrained SDXL VAE encodes each video frame into a latent representation and decodes latents back into pixels.The VAE remains frozen during training, and predicted-frame latents are decoded after denoising.Source: §3.2 · p. 5, §3.3 · p. 6

3.3 IRASim

  • IRASim predicts future video frames from historical-frame latents and a robot action trajectory using a conditional diffusion model in VAE latent space.The predicted frames depict the robot executing the supplied trajectory.Source: §3.1 · p. 4, §3.3 · p. 5
  • Historical frames remain intact while noise is added only to predicted-frame tokens, allowing attention to support consistency with the observed history.The diffusion loss is computed only on predicted frames.Source: §3.3 · p. 5
  • Each action is encoded separately to condition its corresponding generated frame, strengthening frame-level alignment rather than using one trajectory embedding.The action and diffusion-timestep embeddings set scale and shift parameters for each frame's spatial block; temporal-block conditioning is shared across frames.Source: §3.3 · pp. 5–6

4 Experiments

The experiments evaluate trajectory-conditioned video prediction on four real-robot datasets, then test IRASim for policy evaluation, model-based planning, and flexible action control. The evaluation covers short and long trajectories, human preferences, scaling, simulated policy evaluation, and planning in simulation and on a real robot.

Video Generation of Short Trajectories.

Video Generation of Long Trajectories.

Human Preference Evaluation.

Scaling.

  • Across three test datasets, IRASim scales effectively with larger model sizes and more training steps.The tested IRASim-Frame-Ada models range from 33M to 679M parameters.Source: Figure 5 · p. 9, Scaling · p. 9

4.2 Policy Evaluation

  • 0.99 Pearson correlation links success-rate evaluations in IRASim and the ground-truth Mujoco simulator.IRASim also simulates successful and failed rollouts, including a bowl slipping from the gripper.Source: Table 4 · p. 9, §4.2 · p. 10, Figure 6 · p. 10

4.3 Model-based Planning for Policy Improvement

  • With K = 50 and P = 1000, IRASim raises Push-T IoU from 0.637 to 0.961.Policy performance improves with more sampled trajectories when P > 0; the results indicate that training-data size and test-time computation should scale together.Source: Table 5 · p. 11, §4.3 · p. 13
  • On real-robot planning, MSE value functions significantly outperform ResNet features, and both significantly outperform planning without a value function.The no-planning policy randomly selects a trajectory for rollout.Source: Table 6 · p. 12, §4.3 · p. 14

4.4 Flexible Action Controllability

  • IRASim follows keyboard- and VR-controller trajectories in Language-Table and RT-1, including trajectories whose distributions deviate from the original datasets.The generated videos simulate robot-object interactions realistically and reasonably, and handle multimodal generation robustly.Source: Figure 9 · p. 13, §4.4 · p. 14

5 Conclusion, Limitation and Future Work

IRASim generates detailed manipulation videos and is used for policy evaluation and planning, but video generation is not real-time. Appendix results describe short- and long-trajectory generation and quality improvements with scaling.

5 Conclusion, Limitation and Future Work

  • IRASim generates fine-grained robot-object interaction videos from historical observations and action trajectories, but its video generation is not real-time.The authors plan to explore diffusion distillation to accelerate generation and reinforcement learning to improve policies within the world model.Source: §5 · p. 14

A.1 Video Generation of Short Trajectories

  • IRASim-Frame-Ada better aligns frames with actions and models robot-object interactions than other methods in short-trajectory results, including flexible-object manipulation.The RoboNet experiments use two context frames for prediction.Source: §A.1 · p. 19
  • On Bridge, using one versus two context frames yields the same PNSR of 25, and the authors hypothesize that trajectories provide sufficient velocity information.The authors report that adding context frames brings no significant improvement.Source: §A.1 · p. 19

A.2 Video Generation of Long Trajectories

  • IRASim-Frame-Ada generates consistent long-horizon videos that simulate entire trajectories while maintaining the short-setting advantages in frame-action alignment and interaction modeling.The reported long-trajectory evidence is qualitative.Source: §A.2 · p. 19

A.3 Scaling

  • Generated-video reality and accuracy improve as IRASim's model size increases.The scaling results are qualitative.Source: §A.3 · p. 19

A.4 Robustness to Physically Implausible Trajectories

The robustness experiment tests IRASim with a trajectory that commands downward motion after the robot touches the table. The generated video keeps the robot on the table.

A.4 Robustness to Physically Implausible Trajectories

  • The test trajectory commands downward motion after table contact, when the robot cannot physically penetrate the table.The authors use this physically implausible trajectory to evaluate IRASim's handling of such inputs.Source: §A.4 · p. 19
  • IRASim generates a physically accurate video in which the robot stays on the table despite the downward command.The outcome is reported for the physically implausible trajectory in Fig. 14.Source: §A.4 · p. 19

B Datasets

The appendix describes the four datasets, their action representations, and video-conditioning designs. Frame-level conditioning uses each action to condition its corresponding spatial attention frame, while video-level conditioning uses a trajectory embedding.

B Datasets

  • RT-1, Bridge, and Language-Table training samples are 16-frame clips; RoboNet uses two context frames to predict the next 10 frames.Bridge videos are resized to 256 × 320 and Language-Table videos to 288 × 512; RoboNet predictions use 256 × 256 resolution.Source: B Datasets · p. 20
  • RT-1 and Bridge use 7-dimensional relative actions, Language-Table uses 2D translation, and RoboNet uses a 5-dimensional universal action vector.RoboNet's universal vector represents delta XYZ position, delta yaw, and gripper value, with zeros padding unavailable dimensions.Source: B Datasets · p. 20

C.1 Video-Level Conditioning

  • Video-level conditioning adds the diffusion-timestep embedding to the trajectory embedding to regress scale and shift parameters for spatial and temporal blocks.The block computations apply these parameters around attention and feed-forward layers, with layer normalization before scaling and shifting.Source: §C.1 · p. 20

C.2 Frame-Level Condition

  • Frame-level conditioning encodes each trajectory action separately for spatial attention, while temporal attention uses the video-level conditioning embedding.The per-frame embeddings regress scale and shift parameters for the corresponding frame's spatial blocks.Source: §C.2 · pp. 20–21

D Baselines Details

The appendix describes implementation choices for VDM and LVDM and contrasts iVideoGPT and MaskViT with IRASim in tokenization, finetuning, and pretraining.

D Baselines Details

  • VDM uses a 3D U-Net and pixel-space diffusion, whereas LVDM uses the same architecture with latent-space diffusion; both baselines encode trajectories with an MLP.VDM and LVDM share IRASim's initial-frame conditioning method, and LVDM also shares its VAE and training setting.Source: D Baselines Details · p. 21
  • iVideoGPT and MaskViT use VQGAN image tokenizers and require RoboNet finetuning, while IRASim uses the SDXL VAE without extra finetuning.The passage also reports that iVideoGPT is pretrained on OpenX-Embodiment, while IRASim achieves better video prediction performance training only on RoboNet.Source: D Baselines Details · p. 21

E Training Details

The appendix reports training settings, stability observations, inference cost, and configurations used for model scaling.

E Training Details

  • IRASim was trained with AdamW, a constant learning rate of 1e-4, for 300k steps at batch size 64, with gradient clipping set to 0.1; training was very stable.Source: E Training Details · p. 21
  • IRASim training showed no loss spikes even without gradient clipping, whereas LVDM and VDM often did.The models use EMA with decay 0.9999; other hyperparameters follow Peebles and Xie [35].Source: E Training Details · p. 22
  • IRASim generates a 16-frame video in about 30 seconds on an A100 GPU using 8GB of memory.Evaluation uses PNDM with 50 sampling steps; the authors note room to improve latency.Source: E Training Details · p. 22
  • Scaling experiments use four IRASim-Frame-Ada model sizes and focus on the best-performing variant.The configurations are listed in Table 11, and training-resource information is listed in Table 12.Source: E Training Details · p. 22

F Evaluation Details

Evaluation compares generated videos with ground truth under short- and long-trajectory settings, using reconstruction and video-quality metrics. The authors omit FID and FVD for long videos because they consider those metrics poorly suited to this reconstruction task.

F Evaluation Details

  • Latent L2 loss and PSNR measure latent- and pixel-space distances, SSIM measures image similarity, and FID and FVD compare video-feature distributions.Metric calculations exclude the initial frame.Source: F Evaluation Details · p. 22
  • Short-trajectory evaluation generates 15 frames from one initial frame and 15 actions; long trajectories use a complete episode trajectory.Average long-trajectory lengths are 42.5 frames for RT-1, 33.4 for Bridge, and 23.7 for Language-Table; generation is autoregressive.Source: F Evaluation Details · p. 22
  • Long-video FID and FVD are omitted because they poorly reflect human preferences, even for short trajectories, in this reconstruction task.The authors consider reconstruction loss more suitable than distribution similarity for trajectory-to-video evaluation.Source: F Evaluation Details · p. 22

G Real-Robot Model-based Planning Details

The real-robot planning experiment uses IRASim to predict outcomes for candidate trajectories, then selects trajectories by similarity to a goal image. Its success rates exceed a random-trajectory baseline, while cost-function quality and planner choice remain relevant constraints.

G Real-Robot Model-based Planning Details

  • The planner samples 50 candidate trajectories, scores their predicted videos against a goal image, and executes the five lowest-cost trajectories.Costs use either mean squared error or ResNet50-feature cosine similarity; each task is repeated three times.Source: G Real-Robot Model-based Planning… · p. 23
  • Using IRASim to select trajectories significantly increases real-world task success compared with randomly choosing among 50 candidates.The experiment reports quantitative results in Table 5 and qualitative results in Figure 7.Source: G Real-Robot Model-based Planning… · p. 23
  • MSE generally outperforms the ResNet cost but can select incorrect predicted videos, so outcomes depend on both prediction and cost-function accuracy.The authors identify better cost functions as a direction for future work.Source: G Real-Robot Model-based Planning… · p. 23
  • Although the experiment uses a simple sampling-based planner, IRASim can pair with trajectory-sampling policies, and more advanced policies may improve performance and task range.The authors point to policies generating more precise and complex trajectories.Source: G Real-Robot Model-based Planning… · p. 23

H Human Preference Evaluation

The human evaluation compares generated clips against a ground-truth reference, while supplementary figures show qualitative comparisons across datasets and trajectory lengths. Additional examples address scaling, flexible-object deformation, and physically implausible actions.

H Human Preference Evaluation

  • Five participants each evaluated 90 video pairs, comparing IRASim-Frame-Ada with VDM, LVDM, or IRASim-Video-Ada against ground truth.For each participant, the clips included 10 randomly sampled test videos from each of three datasets; pair order was randomized.Source: H Human Preference Evaluation · pp. 23–24
  • Supplementary short- and long-trajectory figures compare methods on RT-1, Bridge, and Language-Table, highlighting differences in colored boxes.For the long-trajectory examples, the input trajectory is an entire episode.Source: Figure 10 · p. 25, Figure 11 · p. 26
  • Scaling examples compare IRASim-Frame-Ada model sizes across RT-1, Bridge, and Language-Table.The figure presents qualitative video results for the different sizes.Source: Figure 12 · p. 27
  • RoboNet examples show clothes being dragged, illustrating simulation of flexible-object deformation.Source: Figure 13 · p. 27
  • When a commanded robot trajectory pokes the table, IRASim depicts the arm blocked by the table rather than passing through it.The caption presents this as robustness to physically implausible trajectories.Source: Figure 14 · p. 28
  • The GUI placed the two generated videos in the upper row in random left-right order and the ground-truth video below.Participants judged which generated video was more realistic and accurate, with a similar option when differences were unclear.Source: H Human Preference Evaluation · p. 24, Figure 15 · p. 28
Loading 2406.14540v2…