修复混合精度vae相关的配置错误，确保在推理阶段正确使用了混合精度模型，并且导出了正确精度的检查点文件。

所有case的baseline，amd版本的ground truth都上传了
2026-02-08 12:35:59 +00:00 · 2026-02-08 09:42:14 +00:00
71 changed files with 3074 additions and 2147 deletions
--- a/.claude/settings.local.json
+++ b/.claude/settings.local.json
@@ -1,15 +0,0 @@
-{
-  "permissions": {
-    "allow": [
-      "Bash(conda env list:*)",
-      "Bash(mamba env:*)",
-      "Bash(micromamba env list:*)",
-      "Bash(echo:*)",
-      "Bash(git show:*)",
-      "Bash(nvidia-smi:*)",
-      "Bash(conda activate unifolm-wma)",
-      "Bash(conda info:*)",
-      "Bash(direnv allow:*)"
-    ]
-  }
-}
--- a/.envrc
+++ b/.envrc
@@ -1,2 +0,0 @@
-eval "$(conda shell.bash hook 2>/dev/null)"
-conda activate unifolm-wma
--- a/.gitignore
+++ b/.gitignore
@@ -55,6 +55,7 @@ coverage.xml
 *.pot

 # Django stuff:
+
 local_settings.py
 db.sqlite3

@@ -120,7 +121,6 @@ localTest/
 fig/
 figure/
 *.mp4
-
 Data/ControlVAE.yml
 Data/Misc
 Data/Pretrained
@@ -129,8 +129,4 @@ Experiment/checkpoint
 Experiment/log

 *.ckpt
-
-*.0
-ckpts/unifolm_wma_dual.ckpt.prepared.pt
-trt_engines/video_backbone.engine
-trt_engines/video_backbone.onnx
+*.0
--- a/case4_run.log
+++ b/case4_run.log
@@ -0,0 +1,135 @@
+nohup: ignoring input
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 07:38:45.572744: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 07:38:45.576864: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:38:45.624825: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 07:38:45.624883: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 07:38:45.627150: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 07:38:45.638316: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:38:45.638803: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 07:38:46.426363: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/7 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+ 14%|█▍        | 1/7 [01:38<09:52, 98.73s/it]
+ 29%|██▊       | 2/7 [03:17<08:14, 98.85s/it]
+ 43%|████▎     | 3/7 [04:56<06:35, 98.80s/it]
+ 57%|█████▋    | 4/7 [06:35<04:56, 98.94s/it]
+ 71%|███████▏  | 5/7 [08:14<03:17, 98.93s/it]
+ 86%|████████▌ | 6/7 [09:53<01:38, 98.89s/it]
+100%|██████████| 7/7 [11:31<00:00, 98.81s/it]
+100%|██████████| 7/7 [11:31<00:00, 98.85s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
--- a/ckpts/configuration.json
+++ b/ckpts/configuration.json
@@ -0,0 +1 @@
+{"framework": "pytorch", "task": "robotics", "allow_remote": true}
--- a/configs/inference/world_model_interaction.yaml
+++ b/configs/inference/world_model_interaction.yaml
@@ -222,7 +222,7 @@ data:
    test:
      target: unifolm_wma.data.wma_data.WMAData
      params:
-        data_dir: '/home/qhy/unifolm-world-model-action/examples/world_model_interaction_prompts'
+        data_dir: '/mnt/ASC1637/unifolm-world-model-action/examples/world_model_interaction_prompts'
        video_length: ${model.params.wma_config.params.temporal_length}
        frame_stride: 2
        load_raw_resolution: True
--- a/env.sh
+++ b/env.sh
@@ -0,0 +1,21 @@
+# Note: This script should be sourced, not executed
+# Usage: source env.sh
+#
+# If you need render group permissions, run this first:
+#   newgrp render
+# Then source this script:
+#   source env.sh
+
+# Initialize conda
+source /mnt/ASC1637/miniconda3/etc/profile.d/conda.sh
+
+# Activate conda environment
+conda activate unifolm-wma-o
+
+# Set HuggingFace cache directories
+export HF_HOME=/mnt/ASC1637/hf_home
+export HUGGINGFACE_HUB_CACHE=/mnt/ASC1637/hf_home/hub
+
+echo "Environment configured successfully"
+echo "Conda environment: unifolm-wma-o"
+echo "HF_HOME: $HF_HOME"
--- a/record.md
+++ b/record.md
@@ -1,208 +0,0 @@
- 好，下面是完整的 decode 优化讲解。                                                                                                    
-                                                                                                                                        
-  ---                                                                                                                                   
-  背景：原始 decode 流程的问题
-                                                                                                                                        
-  原始代码中，每轮迭代的 world model 交互会调用 image_guided_synthesis_sim_mode，内部对DDIM 采样得到的 latent 做全量 VAE decode：     
-
-  # ddpms.py decode_core, perframe_ae=True
-  for index in range(z.shape[0]):  # 逐帧循环，16帧就调16次
-      frame_z = 1. / self.scale_factor * z[index:index+1]
-      frame_result = self.first_stage_model.decode(frame_z)
-      results.append(frame_result)
-
-  假设 8 轮迭代、每轮 16 帧，WM 路径 decode 128 次，policy 路径再 decode 128 次，共 256 次 VAE decode。每次都是同步阻塞的。
-
-  decode 完的视频还会被逐轮写成 mp4 + tensorboard，产生大量磁盘 I/O。最后还要把所有轮的 decoded video 在内存中torch.cat
-  拼接，再写一次完整视频。
-
-  ---
-  优化1：decode_video 开关——按需跳过 decode
-
-  文件： world_model_interaction.py函数 image_guided_synthesis_sim_mode
-
-  改动： 给函数加decode_video 参数（默认 False），返回值增加 raw samples：
-
-  def image_guided_synthesis_sim_mode(...,
-          decode_video: bool = False,  # 新增
-          ...) -> tuple[Tensor | None, Tensor, Tensor, Tensor | None]:
-
-      samples = None
-      if ddim_sampler is not None:
-          samples, actions, states, intermedia = ddim_sampler.sample(...)if decode_video:  # 条件 decode
-              batch_images = model.decode_first_stage(samples)
-              batch_variants = batch_images
-
-      return batch_variants, actions, states, samples# 多返回 samples
-
-  调用侧：
-  - Policy 路径：由 CLI 参数 --fast_policy_no_decode 控制，只需要 action 时可跳过 decode
-  - WM 交互路径：传decode_video=False，只拿 raw latent
-
-  效果： WM 路径每轮省掉 16 帧全量 decode。
-
-  ---
-  优化2：只decode observation 需要的帧
-
-  问题： WM 跳过了全量 decode，但下一轮的CLIP embedding 需要 pixel-space 图像做 observation。
-
-  改动： 只decode exe_steps 帧（通常 1帧），而不是全部 16 帧：
-
-  # WM 调用，不做全量 decode
-  pred_videos_1, _, pred_states, wm_samples = image_guided_synthesis_sim_mode(
-      ..., decode_video=False)
-
-  # 只 decode exe_steps 帧给 observation
-  obs_pixels = model.decode_first_stage(
-      wm_samples[:, :, :args.exe_steps, :, :])
-
-  for idx in range(args.exe_steps):
-      observation = {
-          'observation.images.top':obs_pixels[0, :, idx:idx + 1].permute(1, 0, 2, 3),
-          ...
-      }
-      cond_obs_queues = populate_queues(cond_obs_queues, observation)
-
-  关键细节： 必须逐帧填充 observation queue（idx:idx+1），不能全用最后一帧，否则 CLIP embedding 输入变了会影响精度。
-
-  效果： 每轮从 decode 16 帧降到 decode exe_steps 帧（省15 帧/轮）。
-
-  ---
-  优化3：decode stream——GPU 上并行 decode 和 UNet
-
-  问题： 写入最终视频仍需要完整 segment 的 pixel，这部分 decode 还是要做。
-
-  思路： 用独立 CUDA stream 做 segment decode，和下一轮 UNet 推断在 GPU 上并行。
-
-  改动：
-
-  初始化：
-  decode_stream = torch.cuda.Stream(device=device)
-  pending_decode = None
-
-  循环尾部：
-  # 收集上一轮 decode 结果
-  if pending_decode is not None:
-      decode_stream.synchronize()
-      write_q.put(pending_decode.cpu())
-      pending_decode = None
-
-  # 在 decode stream 上启动当前轮 segment decode（不阻塞主线程）
-  latent_slice = wm_samples[:, :, :args.exe_steps]
-  decode_stream.wait_stream(torch.cuda.current_stream())  # 确保 latent 就绪
-  with torch.cuda.stream(decode_stream):
-      pending_decode = model.decode_first_stage(latent_slice)
-  # 主线程立即进入下一轮 UNet
-
-  循环结束后收集最后一轮：
-  if pending_decode is not None:
-      decode_stream.synchronize()
-      write_q.put(pending_decode.cpu())
-
-  原理： decode_stream.wait_stream() 建立 stream间依赖，确保 latent 产出后才开始 decode。两个 stream 的 kernel 可以被GPU
-  调度器交错执行。
-
-  效果： segment decode 时间被下一轮 UNet 推断掩盖。
-
-  ---
-  优化4：Writer 进程——CPU 工作跨进程并行
-
-  问题： decode 完的tensor 需要转numpy + cv2 编码写盘，这是 CPU 密集型操作，Python GIL 限制线程并行。
-
-  改动：
-
-  辅助函数（主进程和子进程都能调用）：
-  def _video_tensor_to_frames(video: Tensor) -> np.ndarray:
-      video = torch.clamp(video.float(), -1., 1.)
-      n = video.shape[0]
-      video = video.permute(2, 0, 1, 3, 4)
-      frame_grids = [
-          torchvision.utils.make_grid(f, nrow=int(n), padding=0) for f in video
-      ]
-      grid = torch.stack(frame_grids, dim=0)
-      grid = ((grid + 1.0) / 2.0 * 255).to(torch.uint8).permute(0, 2, 3, 1)
-      return grid.numpy()[:, :, :, ::-1]  # RGB → BGR
-
-  Writer 进程：
-  def _video_writer_process(q: mp.Queue, filename: str, fps: int):
-      vwriter = None
-      while True:
-          item = q.get()
-          if item is None:  # sentinel，退出
-              break
-          frames = _video_tensor_to_frames(item)
-          if vwriter is None:
-              h, w = frames.shape[1], frames.shape[2]
-              fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-              vwriter = cv2.VideoWriter(filename, fourcc, fps, (w, h))
-          for f in frames:
-              vwriter.write(f)
-      if vwriter is not None:
-          vwriter.release()
-
-  主进程启动 writer：
-  write_q = mp.Queue()
-  writer_proc = mp.Process(target=_video_writer_process,
-                           args=(write_q, sample_full_video_file, args.save_fps))
-  writer_proc.start()
-
-  主进程通过 write_q.put(tensor.cpu()) 发送数据，循环结束发None sentinel 并join()。
-
-  效果：
-  - tensor→numpy 转换和cv2 编码不占主进程 CPU 时间
-  - 不受 GIL 限制
-  - cv2.VideoWriter 增量写入，不攒内存，不做最终 torch.cat
-  - Queue 自带背压，writer 处理不过来时 put 自然阻塞，不会 OOM
-
-  ---
-  同时删除的冗余 I/O
-
-  - 每轮迭代的 WM 中间tensorboard log（log_to_tensorboard_async）
-  - 每轮迭代的 WM 中间 mp4（save_results_async）
-  - 每轮迭代的 policy 中间 mp4
-  - 最终的 torch.cat + 一次性 save_results_async 写完整视频
-
-  ---
-  总结：优化前后对比
-
-  ┌─────────────────────┬────────────────────────┬───────────────────────────┐
-  │                     │          原始          │          优化后           │
-  ├─────────────────────┼────────────────────────┼───────────────────────────┤
-  │ WM decode帧数/轮    │ 16                     │ exe_steps（通常1）        │
-  ├─────────────────────┼────────────────────────┼───────────────────────────┤
-  │ segment decode 阻塞 │ 同步阻塞主线程         │ decode stream，GPU并行    │
-  ├─────────────────────┼────────────────────────┼───────────────────────────┤
-  │ tensor→numpy+写盘   │ 主线程同步/线程池(GIL) │ 独立进程，无GIL           │
-  ├─────────────────────┼────────────────────────┼───────────────────────────┤
-  │ 视频写入方式        │ 攒内存→cat→一次写      │ cv2增量追加写             │
-  ├─────────────────────┼────────────────────────┼───────────────────────────┤
-  │ 中间I/O             │ 每轮2×mp4+2×tb         │ 无                        │
-  ├─────────────────────┼────────────────────────┼───────────────────────────┤
-  │ 内存占用            │ 随迭代数线性增长       │ 恒定（不攒decoded video） │
-  └─────────────────────┴────────────────────────┴───────────────────────────┘
-
-✻ Cooked for 1m 31s                                                                                                                     
-  
-❯ 目前哪些优化是默认开启的                                                                                                              
-                                                                                                                                      
-● 让我检查各个优化的开关状态。                                                                                                          
-                                                                                                                                        
-● Searched for 2 patterns (ctrl+o to expand)                                                                                            
-                                                                                                                                      
-● 当前各优化的默认状态：                                                                                                                
-                                                                                                                                        
-  ┌───────────────────────────────────┬────────────┬────────────────────────────────────────────────────────┐                           
-  │               优化                │ 默认开启？ │                        控制方式                        │                         
-  ├───────────────────────────────────┼────────────┼────────────────────────────────────────────────────────┤                           
-  │ WM路径跳过全量decode              │ ✅ 是      │ 硬编码 decode_video=False（789行）                     │                           
-  ├───────────────────────────────────┼────────────┼────────────────────────────────────────────────────────┤
-  │ 只decode exe_steps帧给observation │ ✅ 是      │ 无开关，直接生效（792行）                              │
-  ├───────────────────────────────────┼────────────┼────────────────────────────────────────────────────────┤
-  │ decode stream GPU并行             │ ✅ 是      │ 无开关，直接生效（824-827行）                          │
-  ├───────────────────────────────────┼────────────┼────────────────────────────────────────────────────────┤
-  │ Writer进程增量写盘                │ ✅ 是      │ 无开关，直接生效（672-675行）                          │
-  ├───────────────────────────────────┼────────────┼────────────────────────────────────────────────────────┤
-  │ 删除中间wm mp4/tb保存             │ ✅ 是      │ 代码已删除                                             │
-  ├───────────────────────────────────┼────────────┼────────────────────────────────────────────────────────┤
-  │ Policy路径跳过decode              │ ❌ 否      │ 需传--fast_policy_no_decode（默认False，即默认decode） │
-  └───────────────────────────────────┴────────────┴────────────────────────────────────────────────────────┘
--- a/run.log
+++ b/run.log
@@ -0,0 +1,150 @@
+nohup: ignoring input
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 08:15:49.934949: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 08:15:49.937974: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:15:49.969069: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 08:15:49.969100: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 08:15:49.970909: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 08:15:49.979005: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:15:49.979255: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 08:15:50.597743: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/12 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  8%|▊         | 1/12 [01:37<17:51, 97.37s/it]
+ 17%|█▋        | 2/12 [03:14<16:13, 97.31s/it]
+ 25%|██▌       | 3/12 [04:51<14:35, 97.26s/it]
+ 33%|███▎      | 4/12 [06:29<12:58, 97.25s/it]
+ 42%|████▏     | 5/12 [08:06<11:20, 97.24s/it]
+ 50%|█████     | 6/12 [09:43<09:43, 97.24s/it]
+ 58%|█████▊    | 7/12 [11:20<08:06, 97.27s/it]
+ 67%|██████▋   | 8/12 [12:58<06:29, 97.36s/it]
+ 75%|███████▌  | 9/12 [14:36<04:52, 97.49s/it]
+ 83%|████████▎ | 10/12 [16:13<03:15, 97.52s/it]
+ 92%|█████████▏| 11/12 [17:51<01:37, 97.47s/it]
+100%|██████████| 12/12 [19:28<00:00, 97.35s/it]
+100%|██████████| 12/12 [19:28<00:00, 97.35s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 6: generating actions ...
+>>> Step 6: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 8: generating actions ...
+>>> Step 8: interacting with world model ...
--- a/run_all_case.sh
+++ b/run_all_case.sh
@@ -1,114 +0,0 @@
-#!/bin/bash
-
-# 自动执行所有场景的所有case
-# 总共5个场景，每个场景4个case，共20个case
-# 设置环境变量（离线模式）
-export HF_HUB_OFFLINE=1
-export TRANSFORMERS_OFFLINE=1
-
-# 颜色定义
-RED='\033[0;31m'
-GREEN='\033[0;32m'
-YELLOW='\033[1;33m'
-BLUE='\033[0;34m'
-NC='\033[0m' # No Color
-
-# 定义所有场景
-SCENARIOS=(
-    "unitree_g1_pack_camera"
-    "unitree_z1_dual_arm_cleanup_pencils"
-    "unitree_z1_dual_arm_stackbox"
-    "unitree_z1_dual_arm_stackbox_v2"
-    "unitree_z1_stackbox"
-)
-
-# 定义case数量
-CASES=(1 2 3 4)
-
-# 记录开始时间
-START_TIME=$(date +%s)
-LOG_FILE="run_all_cases_$(date +%Y%m%d_%H%M%S).log"
-
-echo -e "${BLUE}========================================${NC}"
-echo -e "${BLUE}开始执行所有场景的case${NC}"
-echo -e "${BLUE}总共: ${#SCENARIOS[@]} 个场景 x ${#CASES[@]} 个case = $((${#SCENARIOS[@]} * ${#CASES[@]})) 个任务${NC}"
-echo -e "${BLUE}日志文件: ${LOG_FILE}${NC}"
-echo -e "${BLUE}========================================${NC}"
-echo ""
-
-# 初始化计数器
-TOTAL_CASES=$((${#SCENARIOS[@]} * ${#CASES[@]}))
-CURRENT_CASE=0
-SUCCESS_COUNT=0
-FAIL_COUNT=0
-
-# 记录失败的case
-declare -a FAILED_CASES
-
-# 遍历所有场景
-for scenario in "${SCENARIOS[@]}"; do
-    echo -e "${YELLOW}>>> 场景: ${scenario}${NC}"
-
-    # 遍历所有case
-    for case_num in "${CASES[@]}"; do
-        CURRENT_CASE=$((CURRENT_CASE + 1))
-        case_dir="${scenario}/case${case_num}"
-        script_path="${case_dir}/run_world_model_interaction.sh"
-
-        echo -e "${BLUE}[${CURRENT_CASE}/${TOTAL_CASES}] 执行: ${case_dir}${NC}"
-
-        # 检查脚本是否存在
-        if [ ! -f "${script_path}" ]; then
-            echo -e "${RED}错误: 脚本不存在 ${script_path}${NC}"
-            FAIL_COUNT=$((FAIL_COUNT + 1))
-            FAILED_CASES+=("${case_dir} (脚本不存在)")
-            continue
-        fi
-
-        # 执行脚本
-        echo "开始时间: $(date '+%Y-%m-%d %H:%M:%S')"
-
-        if bash "${script_path}" >> "${LOG_FILE}" 2>&1; then
-            echo -e "${GREEN}✓ 成功: ${case_dir}${NC}"
-            SUCCESS_COUNT=$((SUCCESS_COUNT + 1))
-        else
-            echo -e "${RED}✗ 失败: ${case_dir}${NC}"
-            FAIL_COUNT=$((FAIL_COUNT + 1))
-            FAILED_CASES+=("${case_dir}")
-        fi
-
-        echo "结束时间: $(date '+%Y-%m-%d %H:%M:%S')"
-        echo ""
-    done
-
-    echo ""
-done
-
-# 计算总耗时
-END_TIME=$(date +%s)
-DURATION=$((END_TIME - START_TIME))
-HOURS=$((DURATION / 3600))
-MINUTES=$(((DURATION % 3600) / 60))
-SECONDS=$((DURATION % 60))
-
-# 输出总结
-echo -e "${BLUE}========================================${NC}"
-echo -e "${BLUE}执行完成！${NC}"
-echo -e "${BLUE}========================================${NC}"
-echo -e "总任务数: ${TOTAL_CASES}"
-echo -e "${GREEN}成功: ${SUCCESS_COUNT}${NC}"
-echo -e "${RED}失败: ${FAIL_COUNT}${NC}"
-echo -e "总耗时: ${HOURS}小时 ${MINUTES}分钟 ${SECONDS}秒"
-echo -e "详细日志: ${LOG_FILE}"
-echo ""
-
-# 如果有失败的case，列出来
-if [ ${FAIL_COUNT} -gt 0 ]; then
-    echo -e "${RED}失败的case列表:${NC}"
-    for failed_case in "${FAILED_CASES[@]}"; do
-        echo -e "${RED}  - ${failed_case}${NC}"
-    done
-    echo ""
-fi
-
-echo -e "${BLUE}========================================${NC}"
--- a/run_all_cases_20260218_190150.log
+++ b/run_all_cases_20260218_190150.log
@@ -1,504 +0,0 @@
-2026-02-18 19:01:56.891895: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
-2026-02-18 19:01:56.940243: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
-2026-02-18 19:01:56.940285: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
-2026-02-18 19:01:56.941395: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
-2026-02-18 19:01:56.948327: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
-To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
-2026-02-18 19:01:57.870809: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
-Global seed set to 123
->>> Loading prepared model from ckpts/unifolm_wma_dual.ckpt.prepared.pt ...
->>> Prepared model loaded.
-INFO:root:***** Configing Data *****
->>> unitree_z1_stackbox: 1 data samples loaded.
->>> unitree_z1_stackbox: data stats loaded.
->>> unitree_z1_stackbox: normalizer initiated.
->>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
->>> unitree_z1_dual_arm_stackbox: data stats loaded.
->>> unitree_z1_dual_arm_stackbox: normalizer initiated.
->>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
->>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
->>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
->>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
->>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
->>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
->>> unitree_g1_pack_camera: 1 data samples loaded.
->>> unitree_g1_pack_camera: data stats loaded.
->>> unitree_g1_pack_camera: normalizer initiated.
->>> Dataset is successfully loaded ...
-    ✓ KV fused: 66 attention layers
-    TRT output 'y': [1, 4, 16, 40, 64] torch.float32
-    TRT output 'hs_a_0': [1, 16, 320, 40, 64] torch.float32
-    TRT output 'hs_a_1': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_2': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_3': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_4': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_5': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_6': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_7': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_8': [1, 16, 320, 40, 64] torch.float32
->>> TRT backbone loaded from /home/qhy/unifolm-world-model-action/scripts/evaluation/../../trt_engines/video_backbone.engine
->>> Generate 16 frames under each generation ...
-DEBUG:h5py._conv:Creating converter from 3 to 5
-DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
-DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
-DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
-
-  0%|          | 0/11 [00:00<?, ?it/s][02/18/2026-19:02:10] [TRT] [W] Using default stream in enqueueV3() may lead to performance issues due to additional calls to cudaStreamSynchronize() by TensorRT to ensure correct synchronization. Please use non-default stream instead.
-
-  9%|▉         | 1/11 [00:17<02:51, 17.15s/it]>>> Step 0: generating actions ...
->>> Step 0: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 1: generating actions ...
-DEBUG:PIL.Image:Importing BlpImagePlugin
-DEBUG:PIL.Image:Importing BlpImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BufrStubImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BufrStubImagePlugin
-DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing DcxImagePlugin
-DEBUG:PIL.Image:Importing DcxImagePlugin
-DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing EpsImagePlugin
-DEBUG:PIL.Image:Importing EpsImagePlugin
-DEBUG:PIL.Image:Importing FitsImagePlugin
-DEBUG:PIL.Image:Importing FitsImagePlugin
-DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FliImagePlugin
-DEBUG:PIL.Image:Importing FliImagePlugin
-DEBUG:PIL.Image:Importing FpxImagePlugin
-DEBUG:PIL.Image:Importing FpxImagePlugin
-DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing FtexImagePlugin
-DEBUG:PIL.Image:Importing FtexImagePlugin
-DEBUG:PIL.Image:Importing GbrImagePlugin
-DEBUG:PIL.Image:Importing GbrImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GribStubImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GribStubImagePlugin
-DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing IcnsImagePlugin
-DEBUG:PIL.Image:Importing IcnsImagePlugin
-DEBUG:PIL.Image:Importing IcoImagePlugin
-DEBUG:PIL.Image:Importing IcoImagePlugin
-DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImtImagePlugin
-DEBUG:PIL.Image:Importing ImtImagePlugin
-DEBUG:PIL.Image:Importing IptcImagePlugin
-DEBUG:PIL.Image:Importing IptcImagePlugin
-DEBUG:PIL.Image:Importing JpegImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing McIdasImagePlugin
-DEBUG:PIL.Image:Importing JpegImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing McIdasImagePlugin
-DEBUG:PIL.Image:Importing MicImagePlugin
-DEBUG:PIL.Image:Importing MicImagePlugin
-DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing MpegImagePlugin
-DEBUG:PIL.Image:Importing MpegImagePlugin
-DEBUG:PIL.Image:Importing MpoImagePlugin
-DEBUG:PIL.Image:Importing MpoImagePlugin
-DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing PalmImagePlugin
-DEBUG:PIL.Image:Importing PalmImagePlugin
-DEBUG:PIL.Image:Importing PcdImagePlugin
-DEBUG:PIL.Image:Importing PcdImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PdfImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PdfImagePlugin
-DEBUG:PIL.Image:Importing PixarImagePlugin
-DEBUG:PIL.Image:Importing PixarImagePlugin
-DEBUG:PIL.Image:Importing PngImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PsdImagePlugin
-DEBUG:PIL.Image:Importing PngImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PsdImagePlugin
-DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing SgiImagePlugin
-DEBUG:PIL.Image:Importing SgiImagePlugin
-DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SunImagePlugin
-DEBUG:PIL.Image:Importing SunImagePlugin
-DEBUG:PIL.Image:Importing TgaImagePlugin
-DEBUG:PIL.Image:Importing TgaImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing WebPImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing WebPImagePlugin
-DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing XbmImagePlugin
-DEBUG:PIL.Image:Importing XbmImagePlugin
-DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XVThumbImagePlugin
-DEBUG:PIL.Image:Importing XVThumbImagePlugin
-
- 18%|█▊        | 2/11 [00:33<02:31, 16.87s/it]
- 27%|██▋       | 3/11 [00:50<02:14, 16.76s/it]
- 36%|███▋      | 4/11 [01:07<01:57, 16.81s/it]
- 45%|████▌     | 5/11 [01:24<01:41, 16.85s/it]
- 55%|█████▍    | 6/11 [01:41<01:24, 16.82s/it]
- 64%|██████▎   | 7/11 [01:57<01:07, 16.82s/it]
- 73%|███████▎  | 8/11 [02:14<00:50, 16.83s/it]
- 82%|████████▏ | 9/11 [02:31<00:33, 16.80s/it]
- 91%|█████████ | 10/11 [02:48<00:16, 16.81s/it]
-100%|██████████| 11/11 [03:05<00:00, 16.81s/it]
-100%|██████████| 11/11 [03:05<00:00, 16.83s/it]
->>> Step 1: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 2: generating actions ...
->>> Step 2: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 3: generating actions ...
->>> Step 3: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 4: generating actions ...
->>> Step 4: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 5: generating actions ...
->>> Step 5: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 6: generating actions ...
->>> Step 6: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 7: generating actions ...
->>> Step 7: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 8: generating actions ...
->>> Step 8: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 9: generating actions ...
->>> Step 9: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 10: generating actions ...
->>> Step 10: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
-
-real	3m49.072s
-user	4m16.055s
-sys	0m44.636s
-2026-02-18 19:05:45.956647: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
-2026-02-18 19:05:46.004149: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
-2026-02-18 19:05:46.004193: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
-2026-02-18 19:05:46.005265: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
-2026-02-18 19:05:46.012074: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
-To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
-2026-02-18 19:05:46.932966: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
-Global seed set to 123
->>> Loading prepared model from ckpts/unifolm_wma_dual.ckpt.prepared.pt ...
->>> Prepared model loaded.
-INFO:root:***** Configing Data *****
->>> unitree_z1_stackbox: 1 data samples loaded.
->>> unitree_z1_stackbox: data stats loaded.
->>> unitree_z1_stackbox: normalizer initiated.
->>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
->>> unitree_z1_dual_arm_stackbox: data stats loaded.
->>> unitree_z1_dual_arm_stackbox: normalizer initiated.
->>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
->>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
->>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
->>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
->>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
->>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
->>> unitree_g1_pack_camera: 1 data samples loaded.
->>> unitree_g1_pack_camera: data stats loaded.
->>> unitree_g1_pack_camera: normalizer initiated.
->>> Dataset is successfully loaded ...
-    ✓ KV fused: 66 attention layers
-    TRT output 'y': [1, 4, 16, 40, 64] torch.float32
-    TRT output 'hs_a_0': [1, 16, 320, 40, 64] torch.float32
-    TRT output 'hs_a_1': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_2': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_3': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_4': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_5': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_6': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_7': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_8': [1, 16, 320, 40, 64] torch.float32
->>> TRT backbone loaded from /home/qhy/unifolm-world-model-action/scripts/evaluation/../../trt_engines/video_backbone.engine
->>> Generate 16 frames under each generation ...
-DEBUG:h5py._conv:Creating converter from 3 to 5
-DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
-DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
-DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
-
-  0%|          | 0/11 [00:00<?, ?it/s][02/18/2026-19:05:59] [TRT] [W] Using default stream in enqueueV3() may lead to performance issues due to additional calls to cudaStreamSynchronize() by TensorRT to ensure correct synchronization. Please use non-default stream instead.
-
-  9%|▉         | 1/11 [00:16<02:47, 16.71s/it]>>> Step 0: generating actions ...
->>> Step 0: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 1: generating actions ...
-DEBUG:PIL.Image:Importing BlpImagePlugin
-DEBUG:PIL.Image:Importing BlpImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BufrStubImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BufrStubImagePlugin
-DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing DcxImagePlugin
-DEBUG:PIL.Image:Importing DcxImagePlugin
-DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing EpsImagePlugin
-DEBUG:PIL.Image:Importing EpsImagePlugin
-DEBUG:PIL.Image:Importing FitsImagePlugin
-DEBUG:PIL.Image:Importing FitsImagePlugin
-DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FliImagePlugin
-DEBUG:PIL.Image:Importing FliImagePlugin
-DEBUG:PIL.Image:Importing FpxImagePlugin
-DEBUG:PIL.Image:Importing FpxImagePlugin
-DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing FtexImagePlugin
-DEBUG:PIL.Image:Importing FtexImagePlugin
-DEBUG:PIL.Image:Importing GbrImagePlugin
-DEBUG:PIL.Image:Importing GbrImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GribStubImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GribStubImagePlugin
-DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing IcnsImagePlugin
-DEBUG:PIL.Image:Importing IcnsImagePlugin
-DEBUG:PIL.Image:Importing IcoImagePlugin
-DEBUG:PIL.Image:Importing IcoImagePlugin
-DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImtImagePlugin
-DEBUG:PIL.Image:Importing ImtImagePlugin
-DEBUG:PIL.Image:Importing IptcImagePlugin
-DEBUG:PIL.Image:Importing IptcImagePlugin
-DEBUG:PIL.Image:Importing JpegImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing McIdasImagePlugin
-DEBUG:PIL.Image:Importing JpegImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing McIdasImagePlugin
-DEBUG:PIL.Image:Importing MicImagePlugin
-DEBUG:PIL.Image:Importing MicImagePlugin
-DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing MpegImagePlugin
-DEBUG:PIL.Image:Importing MpegImagePlugin
-DEBUG:PIL.Image:Importing MpoImagePlugin
-DEBUG:PIL.Image:Importing MpoImagePlugin
-DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing PalmImagePlugin
-DEBUG:PIL.Image:Importing PalmImagePlugin
-DEBUG:PIL.Image:Importing PcdImagePlugin
-DEBUG:PIL.Image:Importing PcdImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PdfImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PdfImagePlugin
-DEBUG:PIL.Image:Importing PixarImagePlugin
-DEBUG:PIL.Image:Importing PixarImagePlugin
-DEBUG:PIL.Image:Importing PngImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PsdImagePlugin
-DEBUG:PIL.Image:Importing PngImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PsdImagePlugin
-DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing SgiImagePlugin
-DEBUG:PIL.Image:Importing SgiImagePlugin
-DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SunImagePlugin
-DEBUG:PIL.Image:Importing SunImagePlugin
-DEBUG:PIL.Image:Importing TgaImagePlugin
-DEBUG:PIL.Image:Importing TgaImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing WebPImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing WebPImagePlugin
-DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing XbmImagePlugin
-DEBUG:PIL.Image:Importing XbmImagePlugin
-DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XVThumbImagePlugin
-DEBUG:PIL.Image:Importing XVThumbImagePlugin
-
- 18%|█▊        | 2/11 [00:33<02:30, 16.75s/it]
- 27%|██▋       | 3/11 [00:50<02:15, 16.91s/it]
- 36%|███▋      | 4/11 [01:07<01:59, 17.02s/it]
- 45%|████▌     | 5/11 [01:24<01:41, 16.98s/it]
- 55%|█████▍    | 6/11 [01:41<01:24, 16.94s/it]
- 64%|██████▎   | 7/11 [01:58<01:07, 16.90s/it]
- 73%|███████▎  | 8/11 [02:15<00:50, 16.83s/it]
- 82%|████████▏ | 9/11 [02:31<00:33, 16.80s/it]
- 91%|█████████ | 10/11 [02:49<00:16, 16.94s/it]
-100%|██████████| 11/11 [03:06<00:00, 16.97s/it]
-100%|██████████| 11/11 [03:06<00:00, 16.91s/it]
->>> Step 1: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 2: generating actions ...
->>> Step 2: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 3: generating actions ...
->>> Step 3: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 4: generating actions ...
->>> Step 4: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 5: generating actions ...
->>> Step 5: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 6: generating actions ...
->>> Step 6: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 7: generating actions ...
->>> Step 7: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 8: generating actions ...
->>> Step 8: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 9: generating actions ...
->>> Step 9: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 10: generating actions ...
->>> Step 10: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
-
-real	3m49.162s
-user	4m12.814s
-sys	0m45.565s
-2026-02-18 19:09:35.113634: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
-2026-02-18 19:09:35.161428: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
-2026-02-18 19:09:35.161474: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
-2026-02-18 19:09:35.162551: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
-2026-02-18 19:09:35.169325: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
-To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
-2026-02-18 19:09:36.089250: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
-Global seed set to 123
->>> Loading prepared model from ckpts/unifolm_wma_dual.ckpt.prepared.pt ...
->>> Prepared model loaded.
-INFO:root:***** Configing Data *****
->>> unitree_z1_stackbox: 1 data samples loaded.
->>> unitree_z1_stackbox: data stats loaded.
->>> unitree_z1_stackbox: normalizer initiated.
->>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
->>> unitree_z1_dual_arm_stackbox: data stats loaded.
->>> unitree_z1_dual_arm_stackbox: normalizer initiated.
->>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
->>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
->>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
->>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
->>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
->>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
->>> unitree_g1_pack_camera: 1 data samples loaded.
->>> unitree_g1_pack_camera: data stats loaded.
->>> unitree_g1_pack_camera: normalizer initiated.
->>> Dataset is successfully loaded ...
-    ✓ KV fused: 66 attention layers
-    TRT output 'y': [1, 4, 16, 40, 64] torch.float32
-    TRT output 'hs_a_0': [1, 16, 320, 40, 64] torch.float32
-    TRT output 'hs_a_1': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_2': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_3': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_4': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_5': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_6': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_7': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_8': [1, 16, 320, 40, 64] torch.float32
->>> TRT backbone loaded from /home/qhy/unifolm-world-model-action/scripts/evaluation/../../trt_engines/video_backbone.engine
->>> Generate 16 frames under each generation ...
-DEBUG:h5py._conv:Creating converter from 3 to 5
-DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
-DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
-DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
-
-  0%|          | 0/11 [00:00<?, ?it/s][02/18/2026-19:09:49] [TRT] [W] Using default stream in enqueueV3() may lead to performance issues due to additional calls to cudaStreamSynchronize() by TensorRT to ensure correct synchronization. Please use non-default stream instead.
-
-  9%|▉         | 1/11 [00:16<02:45, 16.53s/it]>>> Step 0: generating actions ...
->>> Step 0: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 1: generating actions ...
-DEBUG:PIL.Image:Importing BlpImagePlugin
-DEBUG:PIL.Image:Importing BlpImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BufrStubImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BufrStubImagePlugin
-DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing DcxImagePlugin
-DEBUG:PIL.Image:Importing DcxImagePlugin
-DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing EpsImagePlugin
-DEBUG:PIL.Image:Importing EpsImagePlugin
-DEBUG:PIL.Image:Importing FitsImagePlugin
-DEBUG:PIL.Image:Importing FitsImagePlugin
-DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FliImagePlugin
-DEBUG:PIL.Image:Importing FliImagePlugin
-DEBUG:PIL.Image:Importing FpxImagePlugin
-DEBUG:PIL.Image:Importing FpxImagePlugin
-DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing FtexImagePlugin
-DEBUG:PIL.Image:Importing FtexImagePlugin
-DEBUG:PIL.Image:Importing GbrImagePlugin
-DEBUG:PIL.Image:Importing GbrImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GribStubImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GribStubImagePlugin
-DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing IcnsImagePlugin
-DEBUG:PIL.Image:Importing IcnsImagePlugin
-DEBUG:PIL.Image:Importing IcoImagePlugin
-DEBUG:PIL.Image:Importing IcoImagePlugin
-DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImtImagePlugin
-DEBUG:PIL.Image:Importing ImtImagePlugin
-DEBUG:PIL.Image:Importing IptcImagePlugin
-DEBUG:PIL.Image:Importing IptcImagePlugin
-DEBUG:PIL.Image:Importing JpegImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing McIdasImagePlugin
-DEBUG:PIL.Image:Importing JpegImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing McIdasImagePlugin
-DEBUG:PIL.Image:Importing MicImagePlugin
-DEBUG:PIL.Image:Importing MicImagePlugin
-DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing MpegImagePlugin
-DEBUG:PIL.Image:Importing MpegImagePlugin
-DEBUG:PIL.Image:Importing MpoImagePlugin
-DEBUG:PIL.Image:Importing MpoImagePlugin
-DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing PalmImagePlugin
-DEBUG:PIL.Image:Importing PalmImagePlugin
-DEBUG:PIL.Image:Importing PcdImagePlugin
-DEBUG:PIL.Image:Importing PcdImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PdfImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PdfImagePlugin
-DEBUG:PIL.Image:Importing PixarImagePlugin
-DEBUG:PIL.Image:Importing PixarImagePlugin
-DEBUG:PIL.Image:Importing PngImagePlugin
--- a/run_all_psnr.sh
+++ b/run_all_psnr.sh
@@ -1,61 +0,0 @@
-#!/bin/bash
-set -e
-
-SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
-cd "$SCRIPT_DIR"
-
-SCENARIOS=(
-    unitree_g1_pack_camera
-    unitree_z1_dual_arm_cleanup_pencils
-    unitree_z1_dual_arm_stackbox
-    unitree_z1_dual_arm_stackbox_v2
-    unitree_z1_stackbox
-)
-
-CASES=(case1 case2 case3 case4)
-
-total=0
-success=0
-fail=0
-
-for scenario in "${SCENARIOS[@]}"; do
-    for case in "${CASES[@]}"; do
-        case_dir="${scenario}/${case}"
-        gt_video="${case_dir}/${scenario}_${case}.mp4"
-        pred_video=$(ls "${case_dir}"/output/inference/*_full_fs*.mp4 2>/dev/null | head -1)
-        output_file="${case_dir}/psnr_result.json"
-
-        total=$((total + 1))
-        echo "=========================================="
-        echo "[${total}/20] ${case_dir}"
-
-        if [ ! -f "$gt_video" ]; then
-            echo "  SKIP: GT video not found: $gt_video"
-            fail=$((fail + 1))
-            continue
-        fi
-        if [ -z "$pred_video" ]; then
-            echo "  SKIP: pred video not found in ${case_dir}/output/inference/"
-            fail=$((fail + 1))
-            continue
-        fi
-
-        echo "  GT:   $gt_video"
-        echo "  Pred: $pred_video"
-        echo "  Out:  $output_file"
-
-        if python3 psnr_score_for_challenge.py \
-            --gt_video "$gt_video" \
-            --pred_video "$pred_video" \
-            --output_file "$output_file"; then
-            success=$((success + 1))
-            echo "  DONE"
-        else
-            fail=$((fail + 1))
-            echo "  FAILED"
-        fi
-    done
-done
-
-echo "=========================================="
-echo "Finished: ${success} success, ${fail} fail, ${total} total"
--- a/scripts/evaluation/base_model_inference.py
+++ b/scripts/evaluation/base_model_inference.py
@@ -16,9 +16,6 @@ from collections import OrderedDict
 from unifolm_wma.models.samplers.ddim import DDIMSampler
 from unifolm_wma.utils.utils import instantiate_from_config

-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.allow_tf32 = True
-

 def get_filelist(data_dir: str, postfixes: list[str]) -> list[str]:
    """
--- a/scripts/evaluation/real_eval_server.py
+++ b/scripts/evaluation/real_eval_server.py
@@ -19,9 +19,6 @@ from fastapi.responses import JSONResponse
 from typing import Any, Dict, Optional, Tuple, List
 from datetime import datetime

-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.allow_tf32 = True
-
 from unifolm_wma.utils.utils import instantiate_from_config
 from unifolm_wma.models.samplers.ddim import DDIMSampler

--- a/scripts/evaluation/world_model_interaction.py
+++ b/scripts/evaluation/world_model_interaction.py
@@ -1,4 +1,5 @@
 import argparse, os, glob
+from contextlib import nullcontext
 import pandas as pd
 import random
 import torch
@@ -9,8 +10,6 @@ import logging
 import einops
 import warnings
 import imageio
-import atexit
-from concurrent.futures import ThreadPoolExecutor

 from pytorch_lightning import seed_everything
 from omegaconf import OmegaConf
@@ -18,12 +17,8 @@ from tqdm import tqdm
 from einops import rearrange, repeat
 from collections import OrderedDict
 from torch import nn
-from eval_utils import populate_queues
+from eval_utils import populate_queues, log_to_tensorboard
 from collections import deque
-from typing import Optional, List, Any
-
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.allow_tf32 = True
 from torch import Tensor
 from torch.utils.tensorboard import SummaryWriter
 from PIL import Image
@@ -44,6 +39,68 @@ def get_device_from_parameters(module: nn.Module) -> torch.device:
    return next(iter(module.parameters())).device


+def apply_precision_settings(model: nn.Module, args: argparse.Namespace) -> nn.Module:
+    """Apply precision settings to model components based on command-line arguments.
+
+    Args:
+        model (nn.Module): The model to apply precision settings to.
+        args (argparse.Namespace): Parsed command-line arguments containing precision settings.
+
+    Returns:
+        nn.Module: Model with precision settings applied.
+    """
+    print(f">>> Applying precision settings:")
+    print(f"    - Diffusion dtype: {args.diffusion_dtype}")
+    print(f"    - Projector mode: {args.projector_mode}")
+    print(f"    - Encoder mode: {args.encoder_mode}")
+    print(f"    - VAE dtype: {args.vae_dtype}")
+
+    # 1. Set Diffusion backbone precision
+    if args.diffusion_dtype == "bf16":
+        # Convert diffusion model weights to bf16
+        model.model.to(torch.bfloat16)
+        model.diffusion_autocast_dtype = torch.bfloat16
+        print("    ✓ Diffusion model weights converted to bfloat16")
+    else:
+        model.diffusion_autocast_dtype = None
+        print("    ✓ Diffusion model using fp32")
+
+    # 2. Set Projector precision
+    if args.projector_mode == "bf16_full":
+        model.state_projector.to(torch.bfloat16)
+        model.action_projector.to(torch.bfloat16)
+        model.projector_autocast_dtype = None
+        print("    ✓ Projectors converted to bfloat16")
+    elif args.projector_mode == "autocast":
+        model.projector_autocast_dtype = torch.bfloat16
+        print("    ✓ Projectors will use autocast (weights fp32, compute bf16)")
+    else:
+        model.projector_autocast_dtype = None
+        # fp32 mode: do nothing, keep original precision
+
+    # 3. Set Encoder precision
+    if args.encoder_mode == "bf16_full":
+        model.embedder.to(torch.bfloat16)
+        model.image_proj_model.to(torch.bfloat16)
+        model.encoder_autocast_dtype = None
+        print("    ✓ Encoders converted to bfloat16")
+    elif args.encoder_mode == "autocast":
+        model.encoder_autocast_dtype = torch.bfloat16
+        print("    ✓ Encoders will use autocast (weights fp32, compute bf16)")
+    else:
+        model.encoder_autocast_dtype = None
+        # fp32 mode: do nothing, keep original precision
+
+    # 4. Set VAE precision
+    if args.vae_dtype == "bf16":
+        model.first_stage_model.to(torch.bfloat16)
+        print("    ✓ VAE converted to bfloat16")
+    else:
+        print("    ✓ VAE kept in fp32 for best quality")
+
+    return model
+
+
 def write_video(video_path: str, stacked_frames: list, fps: int) -> None:
    """Save a list of frames to a video file.

@@ -156,81 +213,6 @@ def save_results(video: Tensor, filename: str, fps: int = 8) -> None:
                               options={'crf': '10'})


-# ========== Async I/O ==========
-_io_executor: Optional[ThreadPoolExecutor] = None
-_io_futures: List[Any] = []
-
-
-def _get_io_executor() -> ThreadPoolExecutor:
-    global _io_executor
-    if _io_executor is None:
-        _io_executor = ThreadPoolExecutor(max_workers=2)
-    return _io_executor
-
-
-def _flush_io():
-    """Wait for all pending async I/O to finish."""
-    global _io_futures
-    for fut in _io_futures:
-        try:
-            fut.result()
-        except Exception as e:
-            print(f">>> [async I/O] error: {e}")
-    _io_futures.clear()
-
-
-atexit.register(_flush_io)
-
-
-def _save_results_sync(video_cpu: Tensor, filename: str, fps: int) -> None:
-    """Synchronous save on CPU tensor (runs in background thread)."""
-    video = torch.clamp(video_cpu.float(), -1., 1.)
-    n = video.shape[0]
-    video = video.permute(2, 0, 1, 3, 4)
-    frame_grids = [
-        torchvision.utils.make_grid(framesheet, nrow=int(n), padding=0)
-        for framesheet in video
-    ]
-    grid = torch.stack(frame_grids, dim=0)
-    grid = (grid + 1.0) / 2.0
-    grid = (grid * 255).to(torch.uint8).permute(0, 2, 3, 1)
-    torchvision.io.write_video(filename,
-                               grid,
-                               fps=fps,
-                               video_codec='h264',
-                               options={'crf': '10'})
-
-
-def save_results_async(video: Tensor, filename: str, fps: int = 8) -> None:
-    """Submit video saving to background thread pool."""
-    video_cpu = video.detach().cpu()
-    fut = _get_io_executor().submit(_save_results_sync, video_cpu, filename, fps)
-    _io_futures.append(fut)
-
-
-def _log_to_tb_sync(writer, video_cpu: Tensor, tag: str, fps: int) -> None:
-    """Synchronous TensorBoard log on CPU tensor (runs in background thread)."""
-    if video_cpu.dim() == 5:
-        n = video_cpu.shape[0]
-        video = video_cpu.permute(2, 0, 1, 3, 4)
-        frame_grids = [
-            torchvision.utils.make_grid(framesheet, nrow=int(n), padding=0)
-            for framesheet in video
-        ]
-        grid = torch.stack(frame_grids, dim=0)
-        grid = (grid + 1.0) / 2.0
-        grid = grid.unsqueeze(dim=0)
-        writer.add_video(tag, grid, fps=fps)
-
-
-def log_to_tensorboard_async(writer, data: Tensor, tag: str, fps: int = 10) -> None:
-    """Submit TensorBoard logging to background thread pool."""
-    if isinstance(data, torch.Tensor) and data.dim() == 5:
-        data_cpu = data.detach().cpu()
-        fut = _get_io_executor().submit(_log_to_tb_sync, writer, data_cpu, tag, fps)
-        _io_futures.append(fut)
-
-
 def get_init_frame_path(data_dir: str, sample: dict) -> str:
    """Construct the init_frame path from directory and sample metadata.

@@ -343,6 +325,11 @@ def get_latent_z(model, videos: Tensor) -> Tensor:
    """
    b, c, t, h, w = videos.shape
    x = rearrange(videos, 'b c t h w -> (b t) c h w')
+
+    # Auto-detect VAE dtype and convert input
+    vae_dtype = next(model.first_stage_model.parameters()).dtype
+    x = x.to(dtype=vae_dtype)
+
    z = model.encode_first_stage(x)
    z = rearrange(z, '(b t) c h w -> b c t h w', b=b, t=t)
    return z
@@ -408,8 +395,7 @@ def image_guided_synthesis_sim_mode(
        timestep_spacing: str = 'uniform',
        guidance_rescale: float = 0.0,
        sim_mode: bool = True,
-        decode_video: bool = True,
-        **kwargs) -> tuple[torch.Tensor | None, torch.Tensor, torch.Tensor]:
+        **kwargs) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Performs image-guided video generation in a simulation-style mode with optional multimodal guidance (image, state, action, text).

@@ -432,13 +418,10 @@ def image_guided_synthesis_sim_mode(
        timestep_spacing (str): Timestep sampling method in DDIM sampler. Typically "uniform" or "linspace".
        guidance_rescale (float): Guidance rescaling factor to mitigate overexposure from classifier-free guidance.
        sim_mode (bool): Whether to perform world-model interaction or decision-making using the world-model.
-        decode_video (bool): Whether to decode latent samples to pixel-space video.
-            Set to False to skip VAE decode for speed when only actions/states are needed.
        **kwargs: Additional arguments passed to the DDIM sampler.

    Returns:
-        batch_variants (torch.Tensor | None): Predicted pixel-space video frames [B, C, T, H, W],
-            or None when decode_video=False.
+        batch_variants (torch.Tensor): Predicted pixel-space video frames [B, C, T, H, W].
        actions (torch.Tensor): Predicted action sequences [B, T, D] from diffusion decoding.
        states (torch.Tensor): Predicted state sequences [B, T, D] from diffusion decoding.
    """
@@ -448,10 +431,22 @@ def image_guided_synthesis_sim_mode(

    fs = torch.tensor([fs] * batch_size, dtype=torch.long, device=model.device)

+    # Auto-detect model dtype and convert inputs accordingly
+    model_dtype = next(model.embedder.parameters()).dtype
+
    img = observation['observation.images.top'].permute(0, 2, 1, 3, 4)
-    cond_img = rearrange(img, 'b o c h w -> (b o) c h w')[-1:]
-    cond_img_emb = model.embedder(cond_img)
-    cond_img_emb = model.image_proj_model(cond_img_emb)
+    cond_img = rearrange(img, 'b o c h w -> (b o) c h w')[-1:].to(dtype=model_dtype)
+
+    # Encoder autocast: weights stay fp32, compute in bf16
+    enc_ac_dtype = getattr(model, 'encoder_autocast_dtype', None)
+    if enc_ac_dtype is not None and model.device.type == 'cuda':
+        enc_ctx = torch.autocast('cuda', dtype=enc_ac_dtype)
+    else:
+        enc_ctx = nullcontext()
+
+    with enc_ctx:
+        cond_img_emb = model.embedder(cond_img)
+        cond_img_emb = model.image_proj_model(cond_img_emb)

    if model.model.conditioning_key == 'hybrid':
        z = get_latent_z(model, img.permute(0, 2, 1, 3, 4))
@@ -465,11 +460,22 @@ def image_guided_synthesis_sim_mode(
        prompts = [""] * batch_size
    cond_ins_emb = model.get_learned_conditioning(prompts)

-    cond_state_emb = model.state_projector(observation['observation.state'])
-    cond_state_emb = cond_state_emb + model.agent_state_pos_emb
+    # Auto-detect projector dtype and convert inputs
+    projector_dtype = next(model.state_projector.parameters()).dtype

-    cond_action_emb = model.action_projector(observation['action'])
-    cond_action_emb = cond_action_emb + model.agent_action_pos_emb
+    # Projector autocast: weights stay fp32, compute in bf16
+    proj_ac_dtype = getattr(model, 'projector_autocast_dtype', None)
+    if proj_ac_dtype is not None and model.device.type == 'cuda':
+        proj_ctx = torch.autocast('cuda', dtype=proj_ac_dtype)
+    else:
+        proj_ctx = nullcontext()
+
+    with proj_ctx:
+        cond_state_emb = model.state_projector(observation['observation.state'].to(dtype=projector_dtype))
+        cond_state_emb = cond_state_emb + model.agent_state_pos_emb
+
+        cond_action_emb = model.action_projector(observation['action'].to(dtype=projector_dtype))
+        cond_action_emb = cond_action_emb + model.agent_action_pos_emb

    if not sim_mode:
        cond_action_emb = torch.zeros_like(cond_action_emb)
@@ -491,9 +497,17 @@ def image_guided_synthesis_sim_mode(
    kwargs.update({"unconditional_conditioning_img_nonetext": None})
    cond_mask = None
    cond_z0 = None
-    batch_variants = None
+
+    # Setup autocast context for diffusion sampling
+    autocast_dtype = getattr(model, 'diffusion_autocast_dtype', None)
+    if autocast_dtype is not None and model.device.type == 'cuda':
+        autocast_ctx = torch.autocast('cuda', dtype=autocast_dtype)
+    else:
+        autocast_ctx = nullcontext()
+
    if ddim_sampler is not None:
-        samples, actions, states, intermedia = ddim_sampler.sample(
+        with autocast_ctx:
+            samples, actions, states, intermedia = ddim_sampler.sample(
            S=ddim_steps,
            conditioning=cond,
            batch_size=batch_size,
@@ -510,10 +524,9 @@ def image_guided_synthesis_sim_mode(
            guidance_rescale=guidance_rescale,
            **kwargs)

-        if decode_video:
-            # Reconstruct from latent to pixel space
-            batch_images = model.decode_first_stage(samples)
-            batch_variants = batch_images
+        # Reconstruct from latent to pixel space
+        batch_images = model.decode_first_stage(samples)
+        batch_variants = batch_images

    return batch_variants, actions, states

@@ -540,56 +553,37 @@ def run_inference(args: argparse.Namespace, gpu_num: int, gpu_no: int) -> None:
    csv_path = os.path.join(args.prompt_dir, f"{args.dataset}.csv")
    df = pd.read_csv(csv_path)

-    # Load config (always needed for data setup)
+    # Load config
    config = OmegaConf.load(args.config)
+    config['model']['params']['wma_config']['params'][
+        'use_checkpoint'] = False
+    model = instantiate_from_config(config.model)
+    model.perframe_ae = args.perframe_ae
+    assert os.path.exists(args.ckpt_path), "Error: checkpoint Not Found!"
+    model = load_model_checkpoint(model, args.ckpt_path)
+    model.eval()
+    print(f'>>> Load pre-trained model ...')

-    prepared_path = args.ckpt_path + ".prepared.pt"
-    if os.path.exists(prepared_path):
-        # ---- Fast path: load the fully-prepared model ----
-        print(f">>> Loading prepared model from {prepared_path} ...")
-        model = torch.load(prepared_path,
-                           map_location=f"cuda:{gpu_no}",
-                           weights_only=False,
-                           mmap=True)
-        model.eval()
-        print(f">>> Prepared model loaded.")
-    else:
-        # ---- Normal path: construct + load checkpoint ----
-        config['model']['params']['wma_config']['params'][
-            'use_checkpoint'] = False
-        model = instantiate_from_config(config.model)
-        model.perframe_ae = args.perframe_ae
+    # Apply precision settings before moving to GPU
+    model = apply_precision_settings(model, args)

-        assert os.path.exists(args.ckpt_path), "Error: checkpoint Not Found!"
-        model = load_model_checkpoint(model, args.ckpt_path)
-        model.eval()
-        model = model.cuda(gpu_no)
-        print(f'>>> Load pre-trained model ...')
+    # Export precision-converted checkpoint if requested
+    if args.export_precision_ckpt:
+        export_path = args.export_precision_ckpt
+        os.makedirs(os.path.dirname(export_path) or '.', exist_ok=True)
+        torch.save({"state_dict": model.state_dict()}, export_path)
+        print(f">>> Precision-converted checkpoint saved to: {export_path}")
+        return

-        # Save prepared model for fast loading next time
-        print(f">>> Saving prepared model to {prepared_path} ...")
-        torch.save(model, prepared_path)
-        print(f">>> Prepared model saved ({os.path.getsize(prepared_path) / 1024**3:.1f} GB).")
-
-    # Build normalizer (always needed, independent of model loading path)
+    # Build unnomalizer
    logging.info("***** Configing Data *****")
    data = instantiate_from_config(config.data)
    data.setup()
    print(">>> Dataset is successfully loaded ...")

+    model = model.cuda(gpu_no)
    device = get_device_from_parameters(model)

-    # Fuse KV projections in attention layers (to_k + to_v → to_kv)
-    from unifolm_wma.modules.attention import CrossAttention
-    kv_count = sum(1 for m in model.modules()
-                   if isinstance(m, CrossAttention) and m.fuse_kv())
-    print(f"    ✓ KV fused: {kv_count} attention layers")
-
-    # Load TRT backbone if engine exists
-    trt_engine_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', 'trt_engines', 'video_backbone.engine')
-    if os.path.exists(trt_engine_path):
-        model.model.diffusion_model.load_trt_backbone(trt_engine_path)
-
    # Run over data
    assert (args.height % 16 == 0) and (
        args.width % 16
@@ -704,8 +698,7 @@ def run_inference(args: argparse.Namespace, gpu_num: int, gpu_no: int) -> None:
                    fs=model_input_fs,
                    timestep_spacing=args.timestep_spacing,
                    guidance_rescale=args.guidance_rescale,
-                    sim_mode=False,
-                    decode_video=not args.fast_policy_no_decode)
+                    sim_mode=False)

                # Update future actions in the observation queues
                for idx in range(len(pred_actions[0])):
@@ -762,31 +755,29 @@ def run_inference(args: argparse.Namespace, gpu_num: int, gpu_no: int) -> None:
                    cond_obs_queues = populate_queues(cond_obs_queues,
                                                      observation)

-                # Save the imagen videos for decision-making (async)
-                if pred_videos_0 is not None:
-                    sample_tag = f"{args.dataset}-vid{sample['videoid']}-dm-fs-{fs}/itr-{itr}"
-                    log_to_tensorboard_async(writer,
-                                             pred_videos_0,
-                                             sample_tag,
-                                             fps=args.save_fps)
+                # Save the imagen videos for decision-making
+                sample_tag = f"{args.dataset}-vid{sample['videoid']}-dm-fs-{fs}/itr-{itr}"
+                log_to_tensorboard(writer,
+                                   pred_videos_0,
+                                   sample_tag,
+                                   fps=args.save_fps)
                # Save videos environment changes via world-model interaction
                sample_tag = f"{args.dataset}-vid{sample['videoid']}-wd-fs-{fs}/itr-{itr}"
-                log_to_tensorboard_async(writer,
-                                         pred_videos_1,
-                                         sample_tag,
-                                         fps=args.save_fps)
+                log_to_tensorboard(writer,
+                                   pred_videos_1,
+                                   sample_tag,
+                                   fps=args.save_fps)

                # Save the imagen videos for decision-making
-                if pred_videos_0 is not None:
-                    sample_video_file = f'{video_save_dir}/dm/{fs}/itr-{itr}.mp4'
-                    save_results_async(pred_videos_0,
-                                       sample_video_file,
-                                       fps=args.save_fps)
+                sample_video_file = f'{video_save_dir}/dm/{fs}/itr-{itr}.mp4'
+                save_results(pred_videos_0.cpu(),
+                             sample_video_file,
+                             fps=args.save_fps)
                # Save videos environment changes via world-model interaction
                sample_video_file = f'{video_save_dir}/wm/{fs}/itr-{itr}.mp4'
-                save_results_async(pred_videos_1,
-                                   sample_video_file,
-                                   fps=args.save_fps)
+                save_results(pred_videos_1.cpu(),
+                             sample_video_file,
+                             fps=args.save_fps)

                print('>' * 24)
                # Collect the result of world-model interactions
@@ -794,15 +785,12 @@ def run_inference(args: argparse.Namespace, gpu_num: int, gpu_no: int) -> None:

            full_video = torch.cat(wm_video, dim=2)
            sample_tag = f"{args.dataset}-vid{sample['videoid']}-wd-fs-{fs}/full"
-            log_to_tensorboard_async(writer,
-                                     full_video,
-                                     sample_tag,
-                                     fps=args.save_fps)
+            log_to_tensorboard(writer,
+                               full_video,
+                               sample_tag,
+                               fps=args.save_fps)
            sample_full_video_file = f"{video_save_dir}/../{sample['videoid']}_full_fs{fs}.mp4"
-            save_results_async(full_video, sample_full_video_file, fps=args.save_fps)
-
-    # Wait for all async I/O to complete
-    _flush_io()
+            save_results(full_video, sample_full_video_file, fps=args.save_fps)


 def get_parser():
@@ -917,15 +905,39 @@ def get_parser():
                        action='store_true',
                        default=False,
                        help="not using the predicted states as comparison")
-    parser.add_argument(
-        "--fast_policy_no_decode",
-        action='store_true',
-        default=False,
-        help="Speed mode: policy pass only predicts actions, skip policy video decode/log/save.")
    parser.add_argument("--save_fps",
                        type=int,
                        default=8,
                        help="fps for the saving video")
+    parser.add_argument(
+        "--diffusion_dtype",
+        type=str,
+        choices=["fp32", "bf16"],
+        default="bf16",
+        help="Diffusion backbone precision (fp32/bf16)")
+    parser.add_argument(
+        "--projector_mode",
+        type=str,
+        choices=["fp32", "autocast", "bf16_full"],
+        default="bf16_full",
+        help="Projector precision mode (fp32/autocast/bf16_full)")
+    parser.add_argument(
+        "--encoder_mode",
+        type=str,
+        choices=["fp32", "autocast", "bf16_full"],
+        default="bf16_full",
+        help="Encoder precision mode (fp32/autocast/bf16_full)")
+    parser.add_argument(
+        "--vae_dtype",
+        type=str,
+        choices=["fp32", "bf16"],
+        default="fp32",
+        help="VAE precision (fp32/bf16, most affects image quality)")
+    parser.add_argument(
+        "--export_precision_ckpt",
+        type=str,
+        default=None,
+        help="Export precision-converted checkpoint to this path, then exit.")
    return parser


--- a/scripts/export_trt.py
+++ b/scripts/export_trt.py
@@ -1,87 +0,0 @@
-"""Export video UNet backbone to ONNX, then convert to TensorRT engine.
-
-Usage:
-    python scripts/export_trt.py \
-        --ckpt ckpts/unifolm_wma_dual.ckpt.prepared.pt \
-        --config configs/inference/world_model_interaction.yaml \
-        --out_dir trt_engines
-"""
-
-import os
-import sys
-import argparse
-
-import torch
-import tensorrt as trt
-from omegaconf import OmegaConf
-
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
-from unifolm_wma.utils.utils import instantiate_from_config
-from unifolm_wma.trt_utils import export_backbone_onnx
-
-
-def load_model(config_path, ckpt_path):
-    if ckpt_path.endswith('.prepared.pt'):
-        model = torch.load(ckpt_path, map_location='cpu')
-    else:
-        config = OmegaConf.load(config_path)
-        model = instantiate_from_config(config.model)
-        state_dict = torch.load(ckpt_path, map_location='cpu')
-        if 'state_dict' in state_dict:
-            state_dict = state_dict['state_dict']
-        model.load_state_dict(state_dict, strict=False)
-    model.eval().cuda()
-    return model
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--ckpt', required=True)
-    parser.add_argument('--config', default='configs/inference/world_model_interaction.yaml')
-    parser.add_argument('--out_dir', default='trt_engines')
-    parser.add_argument('--context_len', type=int, default=95)
-    parser.add_argument('--fp16', action='store_true', default=True)
-    args = parser.parse_args()
-
-    os.makedirs(args.out_dir, exist_ok=True)
-    onnx_path = os.path.join(args.out_dir, 'video_backbone.onnx')
-    engine_path = os.path.join(args.out_dir, 'video_backbone.engine')
-
-    if os.path.exists(onnx_path):
-        print(f">>> ONNX already exists at {onnx_path}, skipping export.")
-        n_outputs = 10
-    else:
-        print(">>> Loading model ...")
-        model = load_model(args.config, args.ckpt)
-        print(">>> Exporting ONNX ...")
-        with torch.no_grad():
-            n_outputs = export_backbone_onnx(model, onnx_path, context_len=args.context_len)
-        del model
-        torch.cuda.empty_cache()
-
-    print(">>> Converting ONNX -> TensorRT engine ...")
-    logger = trt.Logger(trt.Logger.WARNING)
-    builder = trt.Builder(logger)
-    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
-    parser = trt.OnnxParser(network, logger)
-
-    if not parser.parse_from_file(os.path.abspath(onnx_path)):
-        for i in range(parser.num_errors):
-            print(f"    ONNX parse error: {parser.get_error(i)}")
-        raise RuntimeError("ONNX parsing failed")
-
-    config = builder.create_builder_config()
-    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 16 << 30)
-    if args.fp16:
-        config.set_flag(trt.BuilderFlag.FP16)
-
-    engine_bytes = builder.build_serialized_network(network, config)
-    with open(engine_path, 'wb') as f:
-        f.write(engine_bytes)
-
-    print(f"\n>>> Done! Engine saved to {engine_path}")
-    print(f"    Outputs: 1 y + {n_outputs - 1} hs_a tensors")
-
-
-if __name__ == '__main__':
-    main()
--- a/scripts/trainer.py
+++ b/scripts/trainer.py
@@ -11,9 +11,6 @@ from unifolm_wma.utils.utils import instantiate_from_config
 from unifolm_wma.utils.train import get_trainer_callbacks, get_trainer_logger, get_trainer_strategy
 from unifolm_wma.utils.train import set_logger, init_workspace, load_checkpoints, get_num_parameters

-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.allow_tf32 = True
-

 def get_parser(**parser_kwargs):
    parser = argparse.ArgumentParser(**parser_kwargs)
--- a/src/unifolm_wma/models/ddpms.py
+++ b/src/unifolm_wma/models/ddpms.py
@@ -1105,6 +1105,10 @@ class LatentDiffusion(DDPM):
        else:
            reshape_back = False

+        # Align input dtype with VAE weights (e.g. fp32 samples → bf16 VAE)
+        vae_dtype = next(self.first_stage_model.parameters()).dtype
+        z = z.to(dtype=vae_dtype)
+
        if not self.perframe_ae:
            z = 1. / self.scale_factor * z
            results = self.first_stage_model.decode(z, **kwargs)
@@ -2457,7 +2461,6 @@ class DiffusionWrapper(pl.LightningModule):
        Returns:
            Output from the inner diffusion model (tensor or tuple, depending on the model).
        """
-
        if self.conditioning_key is None:
            out = self.diffusion_model(x, t)
        elif self.conditioning_key == 'concat':
--- a/src/unifolm_wma/models/diffusion_head/conditional_unet1d.py
+++ b/src/unifolm_wma/models/diffusion_head/conditional_unet1d.py
@@ -501,10 +501,6 @@ class ConditionalUnet1D(nn.Module):
        self.last_frame_only = last_frame_only
        self.horizon = horizon

-        # Context precomputation cache
-        self._global_cond_cache_enabled = False
-        self._global_cond_cache = {}
-
    def forward(self,
                sample: torch.Tensor,
                timestep: Union[torch.Tensor, float, int],
@@ -534,20 +530,14 @@ class ConditionalUnet1D(nn.Module):
        B, T, D = sample.shape
        if self.use_linear_act_proj:
            sample = self.proj_in_action(sample.unsqueeze(-1))
-            _gc_key = (cond['image'].data_ptr(), cond['agent_pos'].data_ptr())
-            if self._global_cond_cache_enabled and _gc_key in self._global_cond_cache:
-                global_cond = self._global_cond_cache[_gc_key]
-            else:
-                global_cond = self.obs_encoder(cond)
-                global_cond = rearrange(global_cond,
-                                        '(b t) d -> b 1 (t d)',
-                                        b=B,
-                                        t=self.n_obs_steps)
-                global_cond = repeat(global_cond,
-                                     'b c d -> b (repeat c) d',
-                                     repeat=T)
-                if self._global_cond_cache_enabled:
-                    self._global_cond_cache[_gc_key] = global_cond
+            global_cond = self.obs_encoder(cond)
+            global_cond = rearrange(global_cond,
+                                    '(b t) d -> b 1 (t d)',
+                                    b=B,
+                                    t=self.n_obs_steps)
+            global_cond = repeat(global_cond,
+                                 'b c d -> b (repeat c) d',
+                                 repeat=T)
        else:
            sample = einops.rearrange(sample, 'b h t -> b t h')
            sample = self.proj_in_horizon(sample)
--- a/src/unifolm_wma/models/samplers/ddim.py
+++ b/src/unifolm_wma/models/samplers/ddim.py
@@ -6,8 +6,6 @@ from unifolm_wma.utils.diffusion import make_ddim_sampling_parameters, make_ddim
 from unifolm_wma.utils.common import noise_like
 from unifolm_wma.utils.common import extract_into_tensor
 from tqdm import tqdm
-from unifolm_wma.modules.attention import enable_cross_attn_kv_cache, disable_cross_attn_kv_cache
-from unifolm_wma.modules.networks.wma_model import enable_ctx_cache, disable_ctx_cache


 class DDIMSampler(object):
@@ -69,12 +67,11 @@ class DDIMSampler(object):
            ddim_timesteps=self.ddim_timesteps,
            eta=ddim_eta,
            verbose=verbose)
-        # Ensure tensors are on correct device for efficient indexing
-        self.register_buffer('ddim_sigmas', to_torch(torch.as_tensor(ddim_sigmas)))
-        self.register_buffer('ddim_alphas', to_torch(torch.as_tensor(ddim_alphas)))
-        self.register_buffer('ddim_alphas_prev', to_torch(torch.as_tensor(ddim_alphas_prev)))
+        self.register_buffer('ddim_sigmas', ddim_sigmas)
+        self.register_buffer('ddim_alphas', ddim_alphas)
+        self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
        self.register_buffer('ddim_sqrt_one_minus_alphas',
-                             to_torch(torch.as_tensor(np.sqrt(1. - ddim_alphas))))
+                             np.sqrt(1. - ddim_alphas))
        sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
            (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) *
            (1 - self.alphas_cumprod / self.alphas_cumprod_prev))
@@ -244,70 +241,63 @@ class DDIMSampler(object):

        dp_ddim_scheduler_action.set_timesteps(len(timesteps))
        dp_ddim_scheduler_state.set_timesteps(len(timesteps))
-        ts = torch.empty((b, ), device=device, dtype=torch.long)
-        enable_cross_attn_kv_cache(self.model)
-        enable_ctx_cache(self.model)
-        try:
-            for i, step in enumerate(iterator):
-                index = total_steps - i - 1
-                ts.fill_(step)
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((b, ), step, device=device, dtype=torch.long)

-                # Use mask to blend noised original latent (img_orig) & new sampled latent (img)
-                if mask is not None:
-                    assert x0 is not None
-                    if clean_cond:
-                        img_orig = x0
-                    else:
-                        img_orig = self.model.q_sample(x0, ts)
-                    img = img_orig * mask + (1. - mask) * img
+            # Use mask to blend noised original latent (img_orig) & new sampled latent (img)
+            if mask is not None:
+                assert x0 is not None
+                if clean_cond:
+                    img_orig = x0
+                else:
+                    img_orig = self.model.q_sample(x0, ts)
+                img = img_orig * mask + (1. - mask) * img

-                outs = self.p_sample_ddim(
-                    img,
-                    action,
-                    state,
-                    cond,
-                    ts,
-                    index=index,
-                    use_original_steps=ddim_use_original_steps,
-                    quantize_denoised=quantize_denoised,
-                    temperature=temperature,
-                    noise_dropout=noise_dropout,
-                    score_corrector=score_corrector,
-                    corrector_kwargs=corrector_kwargs,
-                    unconditional_guidance_scale=unconditional_guidance_scale,
-                    unconditional_conditioning=unconditional_conditioning,
-                    mask=mask,
-                    x0=x0,
-                    fs=fs,
-                    guidance_rescale=guidance_rescale,
-                    **kwargs)
+            outs = self.p_sample_ddim(
+                img,
+                action,
+                state,
+                cond,
+                ts,
+                index=index,
+                use_original_steps=ddim_use_original_steps,
+                quantize_denoised=quantize_denoised,
+                temperature=temperature,
+                noise_dropout=noise_dropout,
+                score_corrector=score_corrector,
+                corrector_kwargs=corrector_kwargs,
+                unconditional_guidance_scale=unconditional_guidance_scale,
+                unconditional_conditioning=unconditional_conditioning,
+                mask=mask,
+                x0=x0,
+                fs=fs,
+                guidance_rescale=guidance_rescale,
+                **kwargs)

-                img, pred_x0, model_output_action, model_output_state = outs
+            img, pred_x0, model_output_action, model_output_state = outs

-                action = dp_ddim_scheduler_action.step(
-                    model_output_action,
-                    step,
-                    action,
-                    generator=None,
-                ).prev_sample
-                state = dp_ddim_scheduler_state.step(
-                    model_output_state,
-                    step,
-                    state,
-                    generator=None,
-                ).prev_sample
+            action = dp_ddim_scheduler_action.step(
+                model_output_action,
+                step,
+                action,
+                generator=None,
+            ).prev_sample
+            state = dp_ddim_scheduler_state.step(
+                model_output_state,
+                step,
+                state,
+                generator=None,
+            ).prev_sample

-                if callback: callback(i)
-                if img_callback: img_callback(pred_x0, i)
+            if callback: callback(i)
+            if img_callback: img_callback(pred_x0, i)

-                if index % log_every_t == 0 or index == total_steps - 1:
-                    intermediates['x_inter'].append(img)
-                    intermediates['pred_x0'].append(pred_x0)
-                    intermediates['x_inter_action'].append(action)
-                    intermediates['x_inter_state'].append(state)
-        finally:
-            disable_cross_attn_kv_cache(self.model)
-            disable_ctx_cache(self.model)
+            if index % log_every_t == 0 or index == total_steps - 1:
+                intermediates['x_inter'].append(img)
+                intermediates['pred_x0'].append(pred_x0)
+                intermediates['x_inter_action'].append(action)
+                intermediates['x_inter_state'].append(state)

        return img, action, state, intermediates

@@ -335,6 +325,10 @@ class DDIMSampler(object):
                      guidance_rescale=0.0,
                      **kwargs):
        b, *_, device = *x.shape, x.device
+        if x.dim() == 5:
+            is_video = True
+        else:
+            is_video = False

        if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
            model_output, model_output_action, model_output_state = self.model.apply_model(
@@ -383,11 +377,17 @@ class DDIMSampler(object):
        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
        sigmas = self.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas

-        # Use 0-d tensors directly (already on device); broadcasting handles shape
-        a_t = alphas[index]
-        a_prev = alphas_prev[index]
-        sigma_t = sigmas[index]
-        sqrt_one_minus_at = sqrt_one_minus_alphas[index]
+        if is_video:
+            size = (b, 1, 1, 1, 1)
+        else:
+            size = (b, 1, 1, 1)
+
+        a_t = torch.full(size, alphas[index], device=device)
+        a_prev = torch.full(size, alphas_prev[index], device=device)
+        sigma_t = torch.full(size, sigmas[index], device=device)
+        sqrt_one_minus_at = torch.full(size,
+                                       sqrt_one_minus_alphas[index],
+                                       device=device)

        if self.model.parameterization != "v":
            pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
@@ -395,8 +395,12 @@ class DDIMSampler(object):
            pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output)

        if self.model.use_dynamic_rescale:
-            scale_t = self.ddim_scale_arr[index]
-            prev_scale_t = self.ddim_scale_arr_prev[index]
+            scale_t = torch.full(size,
+                                 self.ddim_scale_arr[index],
+                                 device=device)
+            prev_scale_t = torch.full(size,
+                                      self.ddim_scale_arr_prev[index],
+                                      device=device)
            rescale = (prev_scale_t / scale_t)
            pred_x0 *= rescale

--- a/src/unifolm_wma/modules/attention.py
+++ b/src/unifolm_wma/modules/attention.py
@@ -98,10 +98,6 @@ class CrossAttention(nn.Module):
        self.text_context_len = text_context_len
        self.agent_state_context_len = agent_state_context_len
        self.agent_action_context_len = agent_action_context_len
-        self._kv_cache = {}
-        self._kv_cache_enabled = False
-        self._kv_fused = False
-
        self.cross_attention_scale_learnable = cross_attention_scale_learnable
        if self.image_cross_attention:
            self.to_k_ip = nn.Linear(context_dim, inner_dim, bias=False)
@@ -118,27 +114,6 @@ class CrossAttention(nn.Module):
                self.register_parameter('alpha_caa',
                                        nn.Parameter(torch.tensor(0.)))

-    def fuse_kv(self):
-        """Fuse to_k/to_v into to_kv (2 Linear → 1). Works for all layers."""
-        k_w = self.to_k.weight  # (inner_dim, context_dim)
-        v_w = self.to_v.weight
-        self.to_kv = nn.Linear(k_w.shape[1], k_w.shape[0] * 2, bias=False)
-        self.to_kv.weight = nn.Parameter(torch.cat([k_w, v_w], dim=0))
-        del self.to_k, self.to_v
-        if self.image_cross_attention:
-            for suffix in ('_ip', '_as', '_aa'):
-                k_attr = f'to_k{suffix}'
-                v_attr = f'to_v{suffix}'
-                kw = getattr(self, k_attr).weight
-                vw = getattr(self, v_attr).weight
-                fused = nn.Linear(kw.shape[1], kw.shape[0] * 2, bias=False)
-                fused.weight = nn.Parameter(torch.cat([kw, vw], dim=0))
-                setattr(self, f'to_kv{suffix}', fused)
-                delattr(self, k_attr)
-                delattr(self, v_attr)
-        self._kv_fused = True
-        return True
-
    def forward(self, x, context=None, mask=None):
        spatial_self_attn = (context is None)
        k_ip, v_ip, out_ip = None, None, None
@@ -150,7 +125,7 @@ class CrossAttention(nn.Module):
        context = default(context, x)

        if self.image_cross_attention and not spatial_self_attn:
-            assert 1 > 2, ">>> ERROR: should setup xformers and use efficient_forward ..."
+            # assert 1 > 2, ">>> ERROR: should setup xformers and use efficient_forward ..."
            context_agent_state = context[:, :self.agent_state_context_len, :]
            context_agent_action = context[:,
                                           self.agent_state_context_len:self.
@@ -165,28 +140,19 @@ class CrossAttention(nn.Module):
                                    self.agent_action_context_len +
                                    self.text_context_len:, :]

-            if self._kv_fused:
-                k, v = self.to_kv(context_ins).chunk(2, dim=-1)
-                k_ip, v_ip = self.to_kv_ip(context_image).chunk(2, dim=-1)
-                k_as, v_as = self.to_kv_as(context_agent_state).chunk(2, dim=-1)
-                k_aa, v_aa = self.to_kv_aa(context_agent_action).chunk(2, dim=-1)
-            else:
-                k = self.to_k(context_ins)
-                v = self.to_v(context_ins)
-                k_ip = self.to_k_ip(context_image)
-                v_ip = self.to_v_ip(context_image)
-                k_as = self.to_k_as(context_agent_state)
-                v_as = self.to_v_as(context_agent_state)
-                k_aa = self.to_k_aa(context_agent_action)
-                v_aa = self.to_v_aa(context_agent_action)
+            k = self.to_k(context_ins)
+            v = self.to_v(context_ins)
+            k_ip = self.to_k_ip(context_image)
+            v_ip = self.to_v_ip(context_image)
+            k_as = self.to_k_as(context_agent_state)
+            v_as = self.to_v_as(context_agent_state)
+            k_aa = self.to_k_aa(context_agent_action)
+            v_aa = self.to_v_aa(context_agent_action)
        else:
            if not spatial_self_attn:
                context = context[:, :self.text_context_len, :]
-            if self._kv_fused:
-                k, v = self.to_kv(context).chunk(2, dim=-1)
-            else:
-                k = self.to_k(context)
-                v = self.to_v(context)
+            k = self.to_k(context)
+            v = self.to_v(context)

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h),
                      (q, k, v))
@@ -270,162 +236,134 @@ class CrossAttention(nn.Module):
        k_ip, v_ip, out_ip = None, None, None
        k_as, v_as, out_as = None, None, None
        k_aa, v_aa, out_aa = None, None, None
-        attn_mask_aa = None

-        h = self.heads
        q = self.to_q(x)
        context = default(context, x)

-        b, _, _ = q.shape
-        q = q.unsqueeze(3).reshape(b, q.shape[1], h, self.dim_head).permute(0, 2, 1, 3).reshape(b * h, q.shape[1], self.dim_head).contiguous()
-
-        def _reshape_kv(t):
-            return t.unsqueeze(3).reshape(b, t.shape[1], h, self.dim_head).permute(0, 2, 1, 3).reshape(b * h, t.shape[1], self.dim_head).contiguous()
-
-        use_cache = self._kv_cache_enabled and not spatial_self_attn
-        cache_hit = use_cache and len(self._kv_cache) > 0
-
-        if cache_hit:
-            k = self._kv_cache['k']
-            v = self._kv_cache['v']
-            k_ip = self._kv_cache.get('k_ip')
-            v_ip = self._kv_cache.get('v_ip')
-            k_as = self._kv_cache.get('k_as')
-            v_as = self._kv_cache.get('v_as')
-            k_aa = self._kv_cache.get('k_aa')
-            v_aa = self._kv_cache.get('v_aa')
-            attn_mask_aa = self._kv_cache.get('attn_mask_aa')
-        elif self.image_cross_attention and not spatial_self_attn:
+        if self.image_cross_attention and not spatial_self_attn:
            if context.shape[1] == self.text_context_len + self.video_length:
                context_ins, context_image = context[:, :self.text_context_len, :], context[:,self.text_context_len:, :]
-                if self._kv_fused:
-                    k, v = self.to_kv(context).chunk(2, dim=-1)
-                    k_ip, v_ip = self.to_kv_ip(context_image).chunk(2, dim=-1)
-                else:
-                    k = self.to_k(context)
-                    v = self.to_v(context)
-                    k_ip = self.to_k_ip(context_image)
-                    v_ip = self.to_v_ip(context_image)
-                k, v = map(_reshape_kv, (k, v))
-                k_ip, v_ip = map(_reshape_kv, (k_ip, v_ip))
-                if use_cache:
-                    self._kv_cache = {'k': k, 'v': v, 'k_ip': k_ip, 'v_ip': v_ip}
+                k = self.to_k(context)
+                v = self.to_v(context)
+                k_ip = self.to_k_ip(context_image)
+                v_ip = self.to_v_ip(context_image)
            elif context.shape[1] == self.agent_state_context_len + self.text_context_len + self.video_length:
                context_agent_state = context[:, :self.agent_state_context_len, :]
                context_ins = context[:, self.agent_state_context_len:self.agent_state_context_len+self.text_context_len, :]
                context_image = context[:, self.agent_state_context_len+self.text_context_len:, :]
-                if self._kv_fused:
-                    k, v = self.to_kv(context_ins).chunk(2, dim=-1)
-                    k_ip, v_ip = self.to_kv_ip(context_image).chunk(2, dim=-1)
-                    k_as, v_as = self.to_kv_as(context_agent_state).chunk(2, dim=-1)
-                else:
-                    k = self.to_k(context_ins)
-                    v = self.to_v(context_ins)
-                    k_ip = self.to_k_ip(context_image)
-                    v_ip = self.to_v_ip(context_image)
-                    k_as = self.to_k_as(context_agent_state)
-                    v_as = self.to_v_as(context_agent_state)
-                k, v = map(_reshape_kv, (k, v))
-                k_ip, v_ip = map(_reshape_kv, (k_ip, v_ip))
-                k_as, v_as = map(_reshape_kv, (k_as, v_as))
-                if use_cache:
-                    self._kv_cache = {'k': k, 'v': v, 'k_ip': k_ip, 'v_ip': v_ip, 'k_as': k_as, 'v_as': v_as}
+                k = self.to_k(context_ins)
+                v = self.to_v(context_ins)
+                k_ip = self.to_k_ip(context_image)
+                v_ip = self.to_v_ip(context_image)
+                k_as = self.to_k_as(context_agent_state)
+                v_as = self.to_v_as(context_agent_state)
            else:
                context_agent_state = context[:, :self.agent_state_context_len, :]
                context_agent_action = context[:, self.agent_state_context_len:self.agent_state_context_len+self.agent_action_context_len, :]
                context_ins = context[:, self.agent_state_context_len+self.agent_action_context_len:self.agent_state_context_len+self.agent_action_context_len+self.text_context_len, :]
                context_image = context[:, self.agent_state_context_len+self.agent_action_context_len+self.text_context_len:, :]

-                if self._kv_fused:
-                    k, v = self.to_kv(context_ins).chunk(2, dim=-1)
-                    k_ip, v_ip = self.to_kv_ip(context_image).chunk(2, dim=-1)
-                    k_as, v_as = self.to_kv_as(context_agent_state).chunk(2, dim=-1)
-                    k_aa, v_aa = self.to_kv_aa(context_agent_action).chunk(2, dim=-1)
-                else:
-                    k = self.to_k(context_ins)
-                    v = self.to_v(context_ins)
-                    k_ip = self.to_k_ip(context_image)
-                    v_ip = self.to_v_ip(context_image)
-                    k_as = self.to_k_as(context_agent_state)
-                    v_as = self.to_v_as(context_agent_state)
-                    k_aa = self.to_k_aa(context_agent_action)
-                    v_aa = self.to_v_aa(context_agent_action)
+                k = self.to_k(context_ins)
+                v = self.to_v(context_ins)
+                k_ip = self.to_k_ip(context_image)
+                v_ip = self.to_v_ip(context_image)
+                k_as = self.to_k_as(context_agent_state)
+                v_as = self.to_v_as(context_agent_state)
+                k_aa = self.to_k_aa(context_agent_action)
+                v_aa = self.to_v_aa(context_agent_action)

-                k, v = map(_reshape_kv, (k, v))
-                k_ip, v_ip = map(_reshape_kv, (k_ip, v_ip))
-                k_as, v_as = map(_reshape_kv, (k_as, v_as))
-                k_aa, v_aa = map(_reshape_kv, (k_aa, v_aa))
-
-                attn_mask_aa_raw = self._get_attn_mask_aa(x.shape[0],
-                                                          q.shape[1],
-                                                          k_aa.shape[1],
-                                                          block_size=16,
-                                                          device=k_aa.device)
-                attn_mask_aa = attn_mask_aa_raw.unsqueeze(1).repeat(1, h, 1, 1).reshape(
-                    b * h, attn_mask_aa_raw.shape[1], attn_mask_aa_raw.shape[2]).to(q.dtype)
-
-                if use_cache:
-                    self._kv_cache = {
-                        'k': k, 'v': v, 'k_ip': k_ip, 'v_ip': v_ip,
-                        'k_as': k_as, 'v_as': v_as, 'k_aa': k_aa, 'v_aa': v_aa,
-                        'attn_mask_aa': attn_mask_aa,
-                    }
+                attn_mask_aa = self._get_attn_mask_aa(x.shape[0],
+                                                      q.shape[1],
+                                                      k_aa.shape[1],
+                                                      block_size=16).to(k_aa.device)
        else:
            if not spatial_self_attn:
                assert 1 > 2, ">>> ERROR: you should never go into here ..."
                context = context[:, :self.text_context_len, :]
-            if self._kv_fused:
-                k, v = self.to_kv(context).chunk(2, dim=-1)
-            else:
-                k = self.to_k(context)
-                v = self.to_v(context)
-            k, v = map(_reshape_kv, (k, v))
-            if use_cache:
-                self._kv_cache = {'k': k, 'v': v}
+            k = self.to_k(context)
+            v = self.to_v(context)
+
+        b, _, _ = q.shape
+        q = q.unsqueeze(3).reshape(b, q.shape[1], self.heads, self.dim_head).permute(0, 2, 1, 3).reshape(b * self.heads, q.shape[1], self.dim_head).contiguous()
        if k is not None:
+            k, v = map(
+                lambda t: t.unsqueeze(3).reshape(b, t.shape[
+                    1], self.heads, self.dim_head).permute(0, 2, 1, 3).reshape(
+                        b * self.heads, t.shape[1], self.dim_head).contiguous(),
+                (k, v),
+            )
            out = xformers.ops.memory_efficient_attention(q,
                                                          k,
                                                          v,
                                                          attn_bias=None,
                                                          op=None)
            out = (out.unsqueeze(0).reshape(
-                b, h, out.shape[1],
+                b, self.heads, out.shape[1],
                self.dim_head).permute(0, 2, 1,
                                       3).reshape(b, out.shape[1],
-                                                  h * self.dim_head))
+                                                  self.heads * self.dim_head))

        if k_ip is not None:
+            # For image cross-attention
+            k_ip, v_ip = map(
+                lambda t: t.unsqueeze(3).reshape(b, t.shape[
+                    1], self.heads, self.dim_head).permute(0, 2, 1, 3).reshape(
+                        b * self.heads, t.shape[1], self.dim_head).contiguous(
+                        ),
+                (k_ip, v_ip),
+            )
            out_ip = xformers.ops.memory_efficient_attention(q,
                                                             k_ip,
                                                             v_ip,
                                                             attn_bias=None,
                                                             op=None)
            out_ip = (out_ip.unsqueeze(0).reshape(
-                b, h, out_ip.shape[1],
+                b, self.heads, out_ip.shape[1],
                self.dim_head).permute(0, 2, 1,
                                       3).reshape(b, out_ip.shape[1],
-                                                  h * self.dim_head))
+                                                  self.heads * self.dim_head))

        if k_as is not None:
+            # For agent state cross-attention
+            k_as, v_as = map(
+                lambda t: t.unsqueeze(3).reshape(b, t.shape[
+                    1], self.heads, self.dim_head).permute(0, 2, 1, 3).reshape(
+                        b * self.heads, t.shape[1], self.dim_head).contiguous(
+                        ),
+                (k_as, v_as),
+            )
            out_as = xformers.ops.memory_efficient_attention(q,
                                                             k_as,
                                                             v_as,
                                                             attn_bias=None,
                                                             op=None)
            out_as = (out_as.unsqueeze(0).reshape(
-                b, h, out_as.shape[1],
+                b, self.heads, out_as.shape[1],
                self.dim_head).permute(0, 2, 1,
                                       3).reshape(b, out_as.shape[1],
-                                                  h * self.dim_head))
-
+                                                  self.heads * self.dim_head))
        if k_aa is not None:
+            # For agent action cross-attention
+            k_aa, v_aa = map(
+                lambda t: t.unsqueeze(3).reshape(b, t.shape[
+                    1], self.heads, self.dim_head).permute(0, 2, 1, 3).reshape(
+                        b * self.heads, t.shape[1], self.dim_head).contiguous(
+                        ),
+                (k_aa, v_aa),
+            )
+
+            attn_mask_aa = attn_mask_aa.unsqueeze(1).repeat(1,self.heads,1,1).reshape(
+                    b * self.heads, attn_mask_aa.shape[1], attn_mask_aa.shape[2])
+            attn_mask_aa = attn_mask_aa.to(q.dtype)
+
            out_aa = xformers.ops.memory_efficient_attention(
                q, k_aa, v_aa, attn_bias=attn_mask_aa, op=None)
+
            out_aa = (out_aa.unsqueeze(0).reshape(
-                b, h, out_aa.shape[1],
+                b, self.heads, out_aa.shape[1],
                self.dim_head).permute(0, 2, 1,
                                       3).reshape(b, out_aa.shape[1],
-                                                  h * self.dim_head))
+                                                  self.heads * self.dim_head))
        if exists(mask):
            raise NotImplementedError

@@ -448,43 +386,17 @@ class CrossAttention(nn.Module):

        return self.to_out(out)

-    def _get_attn_mask_aa(self, b, l1, l2, block_size=16, device=None):
-        cache_key = (b, l1, l2, block_size)
-        if hasattr(self, '_attn_mask_aa_cache_key') and self._attn_mask_aa_cache_key == cache_key:
-            cached = self._attn_mask_aa_cache
-            if device is not None and cached.device != torch.device(device):
-                cached = cached.to(device)
-                self._attn_mask_aa_cache = cached
-            return cached
-
-        target_device = device if device is not None else 'cpu'
+    def _get_attn_mask_aa(self, b, l1, l2, block_size=16):
        num_token = l2 // block_size
-        start_positions = ((torch.arange(b, device=target_device) % block_size) + 1) * num_token
-        col_indices = torch.arange(l2, device=target_device)
+        start_positions = ((torch.arange(b) % block_size) + 1) * num_token
+        col_indices = torch.arange(l2)
        mask_2d = col_indices.unsqueeze(0) >= start_positions.unsqueeze(1)
        mask = mask_2d.unsqueeze(1).expand(b, l1, l2)
-        attn_mask = torch.zeros(b, l1, l2, dtype=torch.float, device=target_device)
+        attn_mask = torch.zeros_like(mask, dtype=torch.float)
        attn_mask[mask] = float('-inf')
-
-        self._attn_mask_aa_cache_key = cache_key
-        self._attn_mask_aa_cache = attn_mask
        return attn_mask


-def enable_cross_attn_kv_cache(module):
-    for m in module.modules():
-        if isinstance(m, CrossAttention):
-            m._kv_cache_enabled = True
-            m._kv_cache = {}
-
-
-def disable_cross_attn_kv_cache(module):
-    for m in module.modules():
-        if isinstance(m, CrossAttention):
-            m._kv_cache_enabled = False
-            m._kv_cache = {}
-
-
 class BasicTransformerBlock(nn.Module):

    def __init__(self,
--- a/src/unifolm_wma/modules/networks/wma_model.py
+++ b/src/unifolm_wma/modules/networks/wma_model.py
@@ -685,28 +685,6 @@ class WMAModel(nn.Module):
        self.action_token_projector = instantiate_from_config(
            stem_process_config)

-        # Context precomputation cache
-        self._ctx_cache_enabled = False
-        self._ctx_cache = {}
-        self._trt_backbone = None  # TRT engine for video UNet backbone
-        # Reusable CUDA stream for parallel state_unet / action_unet
-        self._state_stream = torch.cuda.Stream()
-
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        state.pop('_state_stream', None)
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        self._state_stream = torch.cuda.Stream()
-
-    def load_trt_backbone(self, engine_path, n_hs_a=9):
-        """Load a TensorRT engine for the video UNet backbone."""
-        from unifolm_wma.trt_utils import TRTBackbone
-        self._trt_backbone = TRTBackbone(engine_path, n_hs_a=n_hs_a)
-        print(f">>> TRT backbone loaded from {engine_path}")
-
    def forward(self,
                x: Tensor,
                x_action: Tensor,
@@ -742,64 +720,58 @@ class WMAModel(nn.Module):
                                   repeat_only=False).type(x.dtype)
        emb = self.time_embed(t_emb)

-        _ctx_key = context.data_ptr()
-        if self._ctx_cache_enabled and _ctx_key in self._ctx_cache:
-            context = self._ctx_cache[_ctx_key]
+        bt, l_context, _ = context.shape
+        if self.base_model_gen_only:
+            assert l_context == 77 + self.n_obs_steps * 16, ">>> ERROR Context dim 1 ..."  ## NOTE HANDCODE
        else:
-            bt, l_context, _ = context.shape
-            if self.base_model_gen_only:
-                assert l_context == 77 + self.n_obs_steps * 16, ">>> ERROR Context dim 1 ..."  ## NOTE HANDCODE
-            else:
-                if l_context == self.n_obs_steps + 77 + t * 16:
-                    context_agent_state = context[:, :self.n_obs_steps]
-                    context_text = context[:, self.n_obs_steps:self.n_obs_steps +
-                                           77, :]
-                    context_img = context[:, self.n_obs_steps + 77:, :]
-                    context_agent_state = context_agent_state.repeat_interleave(
-                        repeats=t, dim=0)
-                    context_text = context_text.repeat_interleave(repeats=t, dim=0)
-                    context_img = rearrange(context_img,
-                                            'b (t l) c -> (b t) l c',
-                                            t=t)
-                    context = torch.cat(
-                        [context_agent_state, context_text, context_img], dim=1)
-                elif l_context == self.n_obs_steps + 16 + 77 + t * 16:
-                    context_agent_state = context[:, :self.n_obs_steps]
-                    context_agent_action = context[:, self.
-                                                   n_obs_steps:self.n_obs_steps +
-                                                   16, :]
-                    context_agent_action = rearrange(
-                        context_agent_action.unsqueeze(2), 'b t l d -> (b t) l d')
-                    context_agent_action = self.action_token_projector(
-                        context_agent_action)
-                    context_agent_action = rearrange(context_agent_action,
-                                                     '(b o) l d -> b o l d',
-                                                     o=t)
-                    context_agent_action = rearrange(context_agent_action,
-                                                     'b o (t l) d -> b o t l d',
-                                                     t=t)
-                    context_agent_action = context_agent_action.permute(
-                        0, 2, 1, 3, 4)
-                    context_agent_action = rearrange(context_agent_action,
-                                                     'b t o l d -> (b t) (o l) d')
+            if l_context == self.n_obs_steps + 77 + t * 16:
+                context_agent_state = context[:, :self.n_obs_steps]
+                context_text = context[:, self.n_obs_steps:self.n_obs_steps +
+                                       77, :]
+                context_img = context[:, self.n_obs_steps + 77:, :]
+                context_agent_state = context_agent_state.repeat_interleave(
+                    repeats=t, dim=0)
+                context_text = context_text.repeat_interleave(repeats=t, dim=0)
+                context_img = rearrange(context_img,
+                                        'b (t l) c -> (b t) l c',
+                                        t=t)
+                context = torch.cat(
+                    [context_agent_state, context_text, context_img], dim=1)
+            elif l_context == self.n_obs_steps + 16 + 77 + t * 16:
+                context_agent_state = context[:, :self.n_obs_steps]
+                context_agent_action = context[:, self.
+                                               n_obs_steps:self.n_obs_steps +
+                                               16, :]
+                context_agent_action = rearrange(
+                    context_agent_action.unsqueeze(2), 'b t l d -> (b t) l d')
+                context_agent_action = self.action_token_projector(
+                    context_agent_action)
+                context_agent_action = rearrange(context_agent_action,
+                                                 '(b o) l d -> b o l d',
+                                                 o=t)
+                context_agent_action = rearrange(context_agent_action,
+                                                 'b o (t l) d -> b o t l d',
+                                                 t=t)
+                context_agent_action = context_agent_action.permute(
+                    0, 2, 1, 3, 4)
+                context_agent_action = rearrange(context_agent_action,
+                                                 'b t o l d -> (b t) (o l) d')

-                    context_text = context[:, self.n_obs_steps +
-                                           16:self.n_obs_steps + 16 + 77, :]
-                    context_text = context_text.repeat_interleave(repeats=t, dim=0)
+                context_text = context[:, self.n_obs_steps +
+                                       16:self.n_obs_steps + 16 + 77, :]
+                context_text = context_text.repeat_interleave(repeats=t, dim=0)

-                    context_img = context[:, self.n_obs_steps + 16 + 77:, :]
-                    context_img = rearrange(context_img,
-                                            'b (t l) c -> (b t) l c',
-                                            t=t)
-                    context_agent_state = context_agent_state.repeat_interleave(
-                        repeats=t, dim=0)
-                    context = torch.cat([
-                        context_agent_state, context_agent_action, context_text,
-                        context_img
-                    ],
-                                        dim=1)
-            if self._ctx_cache_enabled:
-                self._ctx_cache[_ctx_key] = context
+                context_img = context[:, self.n_obs_steps + 16 + 77:, :]
+                context_img = rearrange(context_img,
+                                        'b (t l) c -> (b t) l c',
+                                        t=t)
+                context_agent_state = context_agent_state.repeat_interleave(
+                    repeats=t, dim=0)
+                context = torch.cat([
+                    context_agent_state, context_agent_action, context_text,
+                    context_img
+                ],
+                                    dim=1)

        emb = emb.repeat_interleave(repeats=t, dim=0)

@@ -819,92 +791,58 @@ class WMAModel(nn.Module):
            fs_embed = fs_embed.repeat_interleave(repeats=t, dim=0)
            emb = emb + fs_embed

-        if self._trt_backbone is not None:
-            # TRT path: run backbone via TensorRT engine
-            h_in = x.type(self.dtype).contiguous()
-            y, hs_a = self._trt_backbone(h_in, emb.contiguous(), context.contiguous())
-        else:
-            # PyTorch path: original backbone
-            h = x.type(self.dtype)
-            adapter_idx = 0
-            hs = []
-            hs_a = []
-            for id, module in enumerate(self.input_blocks):
-                h = module(h, emb, context=context, batch_size=b)
-                if id == 0 and self.addition_attention:
-                    h = self.init_attn(h, emb, context=context, batch_size=b)
-                # plug-in adapter features
-                if ((id + 1) % 3 == 0) and features_adapter is not None:
-                    h = h + features_adapter[adapter_idx]
-                    adapter_idx += 1
-                if id != 0:
-                    if isinstance(module[0], Downsample):
-                        hs_a.append(
-                            rearrange(hs[-1], '(b t) c h w -> b t c h w', t=t))
-                hs.append(h)
-            hs_a.append(rearrange(h, '(b t) c h w -> b t c h w', t=t))
-
-            if features_adapter is not None:
-                assert len(
-                    features_adapter) == adapter_idx, 'Wrong features_adapter'
-            h = self.middle_block(h, emb, context=context, batch_size=b)
-            hs_a.append(rearrange(h, '(b t) c h w -> b t c h w', t=t))
-
-            hs_out = []
-            for module in self.output_blocks:
-                h = torch.cat([h, hs.pop()], dim=1)
-                h = module(h, emb, context=context, batch_size=b)
-                if isinstance(module[-1], Upsample):
+        h = x.type(self.dtype)
+        adapter_idx = 0
+        hs = []
+        hs_a = []
+        for id, module in enumerate(self.input_blocks):
+            h = module(h, emb, context=context, batch_size=b)
+            if id == 0 and self.addition_attention:
+                h = self.init_attn(h, emb, context=context, batch_size=b)
+            # plug-in adapter features
+            if ((id + 1) % 3 == 0) and features_adapter is not None:
+                h = h + features_adapter[adapter_idx]
+                adapter_idx += 1
+            if id != 0:
+                if isinstance(module[0], Downsample):
                    hs_a.append(
-                        rearrange(hs_out[-1], '(b t) c h w -> b t c h w', t=t))
-                hs_out.append(h)
-            h = h.type(x.dtype)
-            hs_a.append(rearrange(hs_out[-1], '(b t) c h w -> b t c h w', t=t))
+                        rearrange(hs[-1], '(b t) c h w -> b t c h w', t=t))
+            hs.append(h)
+        hs_a.append(rearrange(h, '(b t) c h w -> b t c h w', t=t))

-            y = self.out(h)
-            y = rearrange(y, '(b t) c h w -> b c t h w', b=b)
+        if features_adapter is not None:
+            assert len(
+                features_adapter) == adapter_idx, 'Wrong features_adapter'
+        h = self.middle_block(h, emb, context=context, batch_size=b)
+        hs_a.append(rearrange(h, '(b t) c h w -> b t c h w', t=t))
+
+        hs_out = []
+        for module in self.output_blocks:
+            h = torch.cat([h, hs.pop()], dim=1)
+            h = module(h, emb, context=context, batch_size=b)
+            if isinstance(module[-1], Upsample):
+                hs_a.append(
+                    rearrange(hs_out[-1], '(b t) c h w -> b t c h w', t=t))
+            hs_out.append(h)
+        h = h.type(x.dtype)
+        hs_a.append(rearrange(hs_out[-1], '(b t) c h w -> b t c h w', t=t))
+
+        y = self.out(h)
+        y = rearrange(y, '(b t) c h w -> b c t h w', b=b)

        if not self.base_model_gen_only:
            ba, _, _ = x_action.shape
-            ts_state = timesteps[:ba] if b > 1 else timesteps
-            # Run action_unet and state_unet in parallel via CUDA streams
-            s_stream = self._state_stream
-            s_stream.wait_stream(torch.cuda.current_stream())
-            with torch.cuda.stream(s_stream):
-                s_y = self.state_unet(x_state, ts_state, hs_a,
-                                      context_action[:2], **kwargs)
            a_y = self.action_unet(x_action, timesteps[:ba], hs_a,
                                   context_action[:2], **kwargs)
-            torch.cuda.current_stream().wait_stream(s_stream)
+            # Predict state
+            if b > 1:
+                s_y = self.state_unet(x_state, timesteps[:ba], hs_a,
+                                      context_action[:2], **kwargs)
+            else:
+                s_y = self.state_unet(x_state, timesteps, hs_a,
+                                      context_action[:2], **kwargs)
        else:
            a_y = torch.zeros_like(x_action)
            s_y = torch.zeros_like(x_state)

        return y, a_y, s_y
-
-
-def enable_ctx_cache(model):
-    """Enable context precomputation cache on WMAModel and its action/state UNets."""
-    for m in model.modules():
-        if isinstance(m, WMAModel):
-            m._ctx_cache_enabled = True
-            m._ctx_cache = {}
-    # conditional_unet1d cache
-    from unifolm_wma.models.diffusion_head.conditional_unet1d import ConditionalUnet1D
-    for m in model.modules():
-        if isinstance(m, ConditionalUnet1D):
-            m._global_cond_cache_enabled = True
-            m._global_cond_cache = {}
-
-
-def disable_ctx_cache(model):
-    """Disable and clear context precomputation cache."""
-    for m in model.modules():
-        if isinstance(m, WMAModel):
-            m._ctx_cache_enabled = False
-            m._ctx_cache = {}
-    from unifolm_wma.models.diffusion_head.conditional_unet1d import ConditionalUnet1D
-    for m in model.modules():
-        if isinstance(m, ConditionalUnet1D):
-            m._global_cond_cache_enabled = False
-            m._global_cond_cache = {}
--- a/src/unifolm_wma/trt_utils.py
+++ b/src/unifolm_wma/trt_utils.py
@@ -1,151 +0,0 @@
-"""TensorRT acceleration utilities for the video UNet backbone."""
-
-import torch
-import torch.nn as nn
-from einops import rearrange
-from unifolm_wma.modules.networks.wma_model import Downsample, Upsample
-
-
-class VideoBackboneForExport(nn.Module):
-    """Wrapper that isolates the video UNet backbone for ONNX export.
-
-    Takes already-preprocessed inputs (after context/time embedding prep)
-    and returns y + hs_a as a flat tuple.
-    """
-
-    def __init__(self, wma_model):
-        super().__init__()
-        self.input_blocks = wma_model.input_blocks
-        self.middle_block = wma_model.middle_block
-        self.output_blocks = wma_model.output_blocks
-        self.out = wma_model.out
-        self.addition_attention = wma_model.addition_attention
-        if self.addition_attention:
-            self.init_attn = wma_model.init_attn
-        self.dtype = wma_model.dtype
-
-    def forward(self, h, emb, context):
-        t = 16
-        b = 1
-
-        hs = []
-        hs_a = []
-        h = h.type(self.dtype)
-        for id, module in enumerate(self.input_blocks):
-            h = module(h, emb, context=context, batch_size=b)
-            if id == 0 and self.addition_attention:
-                h = self.init_attn(h, emb, context=context, batch_size=b)
-            if id != 0:
-                if isinstance(module[0], Downsample):
-                    hs_a.append(rearrange(hs[-1], '(b t) c h w -> b t c h w', t=t))
-            hs.append(h)
-        hs_a.append(rearrange(h, '(b t) c h w -> b t c h w', t=t))
-
-        h = self.middle_block(h, emb, context=context, batch_size=b)
-        hs_a.append(rearrange(h, '(b t) c h w -> b t c h w', t=t))
-
-        hs_out = []
-        for module in self.output_blocks:
-            h = torch.cat([h, hs.pop()], dim=1)
-            h = module(h, emb, context=context, batch_size=b)
-            if isinstance(module[-1], Upsample):
-                hs_a.append(rearrange(hs_out[-1], '(b t) c h w -> b t c h w', t=t))
-            hs_out.append(h)
-        hs_a.append(rearrange(hs_out[-1], '(b t) c h w -> b t c h w', t=t))
-
-        y = self.out(h.type(h.dtype))
-        y = rearrange(y, '(b t) c h w -> b c t h w', b=b)
-        return (y, *hs_a)
-
-
-def export_backbone_onnx(model, save_path, context_len=95):
-    wma = model.model.diffusion_model
-    wrapper = VideoBackboneForExport(wma)
-    wrapper.eval().cuda()
-
-    for m in wrapper.modules():
-        if hasattr(m, 'checkpoint'):
-            m.checkpoint = False
-        if hasattr(m, 'use_checkpoint'):
-            m.use_checkpoint = False
-
-    import xformers.ops
-    _orig_mea = xformers.ops.memory_efficient_attention
-    def _sdpa_replacement(q, k, v, attn_bias=None, op=None, **kw):
-        return torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias)
-    xformers.ops.memory_efficient_attention = _sdpa_replacement
-
-    BT = 16
-    emb_dim = wma.model_channels * 4
-    ctx_dim = 1024
-    in_ch = wma.in_channels
-
-    dummy_h = torch.randn(BT, in_ch, 40, 64, device='cuda', dtype=torch.float32)
-    dummy_emb = torch.randn(BT, emb_dim, device='cuda', dtype=torch.float32)
-    dummy_ctx = torch.randn(BT, context_len, ctx_dim, device='cuda', dtype=torch.float32)
-
-    with torch.no_grad():
-        outputs = wrapper(dummy_h, dummy_emb, dummy_ctx)
-        n_outputs = len(outputs)
-        print(f">>> Backbone has {n_outputs} outputs (1 y + {n_outputs-1} hs_a)")
-        for i, o in enumerate(outputs):
-            print(f"    output[{i}]: {o.shape} {o.dtype}")
-
-    output_names = ['y'] + [f'hs_a_{i}' for i in range(n_outputs - 1)]
-
-    torch.onnx.export(
-        wrapper,
-        (dummy_h, dummy_emb, dummy_ctx),
-        save_path,
-        input_names=['h', 'emb', 'context'],
-        output_names=output_names,
-        opset_version=17,
-        do_constant_folding=True,
-    )
-    print(f">>> ONNX exported to {save_path}")
-    xformers.ops.memory_efficient_attention = _orig_mea
-    return n_outputs
-
-
-class TRTBackbone:
-    """TensorRT runtime wrapper for the video UNet backbone."""
-
-    def __init__(self, engine_path, n_hs_a=9):
-        import tensorrt as trt
-
-        self.logger = trt.Logger(trt.Logger.WARNING)
-        with open(engine_path, 'rb') as f:
-            runtime = trt.Runtime(self.logger)
-            self.engine = runtime.deserialize_cuda_engine(f.read())
-        self.context = self.engine.create_execution_context()
-        self.n_hs_a = n_hs_a
-
-        import numpy as np
-        self.output_buffers = {}
-        for i in range(self.engine.num_io_tensors):
-            name = self.engine.get_tensor_name(i)
-            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT:
-                shape = self.engine.get_tensor_shape(name)
-                np_dtype = trt.nptype(self.engine.get_tensor_dtype(name))
-                buf = torch.empty(list(shape), dtype=torch.from_numpy(np.empty(0, dtype=np_dtype)).dtype, device='cuda')
-                self.output_buffers[name] = buf
-                print(f"    TRT output '{name}': {list(shape)} {buf.dtype}")
-
-    def __call__(self, h, emb, context):
-        import tensorrt as trt
-        for name, tensor in [('h', h), ('emb', emb), ('context', context)]:
-            expected_dtype = trt.nptype(self.engine.get_tensor_dtype(name))
-            torch_expected = torch.from_numpy(__import__('numpy').empty(0, dtype=expected_dtype)).dtype
-            if tensor.dtype != torch_expected:
-                tensor = tensor.to(torch_expected)
-            self.context.set_tensor_address(name, tensor.contiguous().data_ptr())
-
-        for name, buf in self.output_buffers.items():
-            self.context.set_tensor_address(name, buf.data_ptr())
-
-        self.context.execute_async_v3(torch.cuda.current_stream().cuda_stream)
-        torch.cuda.synchronize()
-
-        y = self.output_buffers['y']
-        hs_a = [self.output_buffers[f'hs_a_{i}'] for i in range(self.n_hs_a)]
-        return y, hs_a
--- a/unitree_g1_pack_camera/case1/output.log
+++ b/unitree_g1_pack_camera/case1/output.log
@@ -1,13 +1,32 @@
-2026-02-18 19:01:56.891895: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
-2026-02-18 19:01:56.940243: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
-2026-02-18 19:01:56.940285: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
-2026-02-18 19:01:56.941395: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
-2026-02-18 19:01:56.948327: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
-To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
-2026-02-18 19:01:57.870809: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+2026-02-08 05:20:49.828675: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 05:20:49.831563: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 05:20:49.861366: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 05:20:49.861402: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 05:20:49.862974: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 05:20:49.870402: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 05:20:49.870647: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 05:20:50.486843: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
 Global seed set to 123
->>> Loading prepared model from ckpts/unifolm_wma_dual.ckpt.prepared.pt ...
->>> Prepared model loaded.
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
 INFO:root:***** Configing Data *****
 >>> unitree_z1_stackbox: 1 data samples loaded.
 >>> unitree_z1_stackbox: data stats loaded.
@@ -25,125 +44,71 @@ INFO:root:***** Configing Data *****
 >>> unitree_g1_pack_camera: data stats loaded.
 >>> unitree_g1_pack_camera: normalizer initiated.
 >>> Dataset is successfully loaded ...
-    ✓ KV fused: 66 attention layers
-    TRT output 'y': [1, 4, 16, 40, 64] torch.float32
-    TRT output 'hs_a_0': [1, 16, 320, 40, 64] torch.float32
-    TRT output 'hs_a_1': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_2': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_3': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_4': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_5': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_6': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_7': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_8': [1, 16, 320, 40, 64] torch.float32
->>> TRT backbone loaded from /home/qhy/unifolm-world-model-action/scripts/evaluation/../../trt_engines/video_backbone.engine
 >>> Generate 16 frames under each generation ...
 DEBUG:h5py._conv:Creating converter from 3 to 5
 DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
 DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
 DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
-
-  0%|          | 0/11 [00:00<?, ?it/s][02/18/2026-19:02:10] [TRT] [W] Using default stream in enqueueV3() may lead to performance issues due to additional calls to cudaStreamSynchronize() by TensorRT to ensure correct synchronization. Please use non-default stream instead.
+
+  0%|          | 0/11 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
 >>> Step 0: generating actions ...
-  9%|▉         | 1/11 [00:17<02:51, 17.15s/it]>>> Step 0: generating actions ...
->>> Step 0: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
 >>> Step 0: interacting with world model ...
 DEBUG:PIL.Image:Importing BlpImagePlugin
 DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BufrStubImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
 DEBUG:PIL.Image:Importing BufrStubImagePlugin
 DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing DcxImagePlugin
 DEBUG:PIL.Image:Importing DcxImagePlugin
 DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing EpsImagePlugin
 DEBUG:PIL.Image:Importing EpsImagePlugin
 DEBUG:PIL.Image:Importing FitsImagePlugin
-DEBUG:PIL.Image:Importing FitsImagePlugin
 DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FliImagePlugin
 DEBUG:PIL.Image:Importing FliImagePlugin
 DEBUG:PIL.Image:Importing FpxImagePlugin
 DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing FtexImagePlugin
 DEBUG:PIL.Image:Importing FtexImagePlugin
 DEBUG:PIL.Image:Importing GbrImagePlugin
 DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GribStubImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
 DEBUG:PIL.Image:Importing GribStubImagePlugin
 DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing IcnsImagePlugin
 DEBUG:PIL.Image:Importing IcnsImagePlugin
 DEBUG:PIL.Image:Importing IcoImagePlugin
-DEBUG:PIL.Image:Importing IcoImagePlugin
 DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImtImagePlugin
 DEBUG:PIL.Image:Importing ImtImagePlugin
 DEBUG:PIL.Image:Importing IptcImagePlugin
 DEBUG:PIL.Image:Importing JpegImagePlugin
 DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing McIdasImagePlugin
-DEBUG:PIL.Image:Importing JpegImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
 DEBUG:PIL.Image:Importing McIdasImagePlugin
 DEBUG:PIL.Image:Importing MicImagePlugin
 DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing MpegImagePlugin
 DEBUG:PIL.Image:Importing MpegImagePlugin
 DEBUG:PIL.Image:Importing MpoImagePlugin
-DEBUG:PIL.Image:Importing MpoImagePlugin
 DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing PalmImagePlugin
 DEBUG:PIL.Image:Importing PalmImagePlugin
 DEBUG:PIL.Image:Importing PcdImagePlugin
 DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PdfImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
 DEBUG:PIL.Image:Importing PdfImagePlugin
 DEBUG:PIL.Image:Importing PixarImagePlugin
 DEBUG:PIL.Image:Importing PngImagePlugin
 DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PsdImagePlugin
-DEBUG:PIL.Image:Importing PngImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
 DEBUG:PIL.Image:Importing PsdImagePlugin
 DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing SgiImagePlugin
 DEBUG:PIL.Image:Importing SgiImagePlugin
 DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SunImagePlugin
 DEBUG:PIL.Image:Importing SunImagePlugin
 DEBUG:PIL.Image:Importing TgaImagePlugin
 DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing WebPImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
 DEBUG:PIL.Image:Importing WebPImagePlugin
 DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing XbmImagePlugin
 DEBUG:PIL.Image:Importing XbmImagePlugin
 DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XVThumbImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  9%|▉         | 1/11 [01:38<16:25, 98.56s/it]
 18%|█▊        | 2/11 [03:16<14:44, 98.31s/it]
 27%|██▋       | 3/11 [04:55<13:06, 98.33s/it]
 36%|███▋      | 4/11 [06:36<11:37, 99.66s/it]
@@ -174,6 +139,6 @@ DEBUG:PIL.Image:Importing XVThumbImagePlugin
 >>> Step 6: generating actions ...
 >>> Step 6: interacting with world model ...
 >>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 7: generating actions ...
->>> Step 7: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_g1_pack_camera/case1/psnr_result.json
+++ b/unitree_g1_pack_camera/case1/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_g1_pack_camera/case1/unitree_g1_pack_camera_case1.mp4",
-    "pred_video": "unitree_g1_pack_camera/case1/output/inference/0_full_fs6.mp4",
-    "psnr": 35.615362167470806
+    "pred_video": "unitree_g1_pack_camera/case1/output/inference/unitree_g1_pack_camera_case1_amd.mp4",
+    "psnr": 16.415668383379177
 }
--- a/unitree_g1_pack_camera/case2/output.log
+++ b/unitree_g1_pack_camera/case2/output.log
@@ -1,13 +1,32 @@
-2026-02-18 19:05:45.956647: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
-2026-02-18 19:05:46.004149: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
-2026-02-18 19:05:46.004193: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
-2026-02-18 19:05:46.005265: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
-2026-02-18 19:05:46.012074: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
-To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
-2026-02-18 19:05:46.932966: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
-Global seed set to 123
->>> Loading prepared model from ckpts/unifolm_wma_dual.ckpt.prepared.pt ...
->>> Prepared model loaded.
+2026-02-08 05:06:45.806187: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 05:06:45.809295: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 05:06:45.840950: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 05:06:45.840981: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 05:06:45.842814: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 05:06:45.851049: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 05:06:45.851316: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 05:06:47.225477: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
 INFO:root:***** Configing Data *****
 >>> unitree_z1_stackbox: 1 data samples loaded.
 >>> unitree_z1_stackbox: data stats loaded.
@@ -25,125 +44,71 @@ INFO:root:***** Configing Data *****
 >>> unitree_g1_pack_camera: data stats loaded.
 >>> unitree_g1_pack_camera: normalizer initiated.
 >>> Dataset is successfully loaded ...
-    ✓ KV fused: 66 attention layers
-    TRT output 'y': [1, 4, 16, 40, 64] torch.float32
-    TRT output 'hs_a_0': [1, 16, 320, 40, 64] torch.float32
-    TRT output 'hs_a_1': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_2': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_3': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_4': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_5': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_6': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_7': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_8': [1, 16, 320, 40, 64] torch.float32
->>> TRT backbone loaded from /home/qhy/unifolm-world-model-action/scripts/evaluation/../../trt_engines/video_backbone.engine
 >>> Generate 16 frames under each generation ...
 DEBUG:h5py._conv:Creating converter from 3 to 5
 DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
 DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
 DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
-
-  0%|          | 0/11 [00:00<?, ?it/s][02/18/2026-19:05:59] [TRT] [W] Using default stream in enqueueV3() may lead to performance issues due to additional calls to cudaStreamSynchronize() by TensorRT to ensure correct synchronization. Please use non-default stream instead.
+
+  0%|          | 0/11 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
 >>> Step 0: generating actions ...
-  9%|▉         | 1/11 [00:16<02:47, 16.71s/it]>>> Step 0: generating actions ...
->>> Step 0: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
 >>> Step 0: interacting with world model ...
 DEBUG:PIL.Image:Importing BlpImagePlugin
 DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BufrStubImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
 DEBUG:PIL.Image:Importing BufrStubImagePlugin
 DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing DcxImagePlugin
 DEBUG:PIL.Image:Importing DcxImagePlugin
 DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing EpsImagePlugin
 DEBUG:PIL.Image:Importing EpsImagePlugin
 DEBUG:PIL.Image:Importing FitsImagePlugin
-DEBUG:PIL.Image:Importing FitsImagePlugin
 DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FliImagePlugin
 DEBUG:PIL.Image:Importing FliImagePlugin
 DEBUG:PIL.Image:Importing FpxImagePlugin
 DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing FtexImagePlugin
 DEBUG:PIL.Image:Importing FtexImagePlugin
 DEBUG:PIL.Image:Importing GbrImagePlugin
 DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GribStubImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
 DEBUG:PIL.Image:Importing GribStubImagePlugin
 DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing IcnsImagePlugin
 DEBUG:PIL.Image:Importing IcnsImagePlugin
 DEBUG:PIL.Image:Importing IcoImagePlugin
-DEBUG:PIL.Image:Importing IcoImagePlugin
 DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImtImagePlugin
 DEBUG:PIL.Image:Importing ImtImagePlugin
 DEBUG:PIL.Image:Importing IptcImagePlugin
 DEBUG:PIL.Image:Importing JpegImagePlugin
 DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing McIdasImagePlugin
-DEBUG:PIL.Image:Importing JpegImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
 DEBUG:PIL.Image:Importing McIdasImagePlugin
 DEBUG:PIL.Image:Importing MicImagePlugin
 DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing MpegImagePlugin
 DEBUG:PIL.Image:Importing MpegImagePlugin
 DEBUG:PIL.Image:Importing MpoImagePlugin
-DEBUG:PIL.Image:Importing MpoImagePlugin
 DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing PalmImagePlugin
 DEBUG:PIL.Image:Importing PalmImagePlugin
 DEBUG:PIL.Image:Importing PcdImagePlugin
 DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PdfImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
 DEBUG:PIL.Image:Importing PdfImagePlugin
 DEBUG:PIL.Image:Importing PixarImagePlugin
 DEBUG:PIL.Image:Importing PngImagePlugin
 DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PsdImagePlugin
-DEBUG:PIL.Image:Importing PngImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
 DEBUG:PIL.Image:Importing PsdImagePlugin
 DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing SgiImagePlugin
 DEBUG:PIL.Image:Importing SgiImagePlugin
 DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SunImagePlugin
 DEBUG:PIL.Image:Importing SunImagePlugin
 DEBUG:PIL.Image:Importing TgaImagePlugin
 DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing WebPImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
 DEBUG:PIL.Image:Importing WebPImagePlugin
 DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing XbmImagePlugin
 DEBUG:PIL.Image:Importing XbmImagePlugin
 DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XVThumbImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  9%|▉         | 1/11 [01:37<16:14, 97.41s/it]
 18%|█▊        | 2/11 [03:14<14:35, 97.22s/it]
 27%|██▋       | 3/11 [04:51<12:58, 97.33s/it]
 36%|███▋      | 4/11 [06:29<11:22, 97.47s/it]
@@ -174,6 +139,6 @@ DEBUG:PIL.Image:Importing XVThumbImagePlugin
 >>> Step 6: generating actions ...
 >>> Step 6: interacting with world model ...
 >>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 7: generating actions ...
->>> Step 7: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_g1_pack_camera/case2/psnr_result.json
+++ b/unitree_g1_pack_camera/case2/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_g1_pack_camera/case2/unitree_g1_pack_camera_case2.mp4",
-    "pred_video": "unitree_g1_pack_camera/case2/output/inference/50_full_fs6.mp4",
-    "psnr": 34.61979248212279
+    "pred_video": "unitree_g1_pack_camera/case2/output/inference/unitree_g1_pack_camera_case2_amd.mp4",
+    "psnr": 19.515250190529375
 }
--- a/unitree_g1_pack_camera/case3/output.log
+++ b/unitree_g1_pack_camera/case3/output.log
@@ -1,13 +1,32 @@
-2026-02-18 19:09:35.113634: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
-2026-02-18 19:09:35.161428: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
-2026-02-18 19:09:35.161474: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
-2026-02-18 19:09:35.162551: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
-2026-02-18 19:09:35.169325: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
-To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
-2026-02-18 19:09:36.089250: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
-Global seed set to 123
->>> Loading prepared model from ckpts/unifolm_wma_dual.ckpt.prepared.pt ...
->>> Prepared model loaded.
+2026-02-08 05:08:32.803904: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 05:08:32.807010: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 05:08:32.837936: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 05:08:32.837978: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 05:08:32.839785: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 05:08:32.847835: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 05:08:32.848223: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 05:08:34.120114: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
 INFO:root:***** Configing Data *****
 >>> unitree_z1_stackbox: 1 data samples loaded.
 >>> unitree_z1_stackbox: data stats loaded.
@@ -25,122 +44,101 @@ INFO:root:***** Configing Data *****
 >>> unitree_g1_pack_camera: data stats loaded.
 >>> unitree_g1_pack_camera: normalizer initiated.
 >>> Dataset is successfully loaded ...
-    ✓ KV fused: 66 attention layers
-    TRT output 'y': [1, 4, 16, 40, 64] torch.float32
-    TRT output 'hs_a_0': [1, 16, 320, 40, 64] torch.float32
-    TRT output 'hs_a_1': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_2': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_3': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_4': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_5': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_6': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_7': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_8': [1, 16, 320, 40, 64] torch.float32
->>> TRT backbone loaded from /home/qhy/unifolm-world-model-action/scripts/evaluation/../../trt_engines/video_backbone.engine
 >>> Generate 16 frames under each generation ...
 DEBUG:h5py._conv:Creating converter from 3 to 5
 DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
 DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
 DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
-
-  0%|          | 0/11 [00:00<?, ?it/s][02/18/2026-19:09:49] [TRT] [W] Using default stream in enqueueV3() may lead to performance issues due to additional calls to cudaStreamSynchronize() by TensorRT to ensure correct synchronization. Please use non-default stream instead.
+
+  0%|          | 0/11 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
 >>> Step 0: generating actions ...
-  9%|▉         | 1/11 [00:16<02:45, 16.53s/it]>>> Step 0: generating actions ...
->>> Step 0: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
 >>> Step 0: interacting with world model ...
 DEBUG:PIL.Image:Importing BlpImagePlugin
 DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BufrStubImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
 DEBUG:PIL.Image:Importing BufrStubImagePlugin
 DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing DcxImagePlugin
 DEBUG:PIL.Image:Importing DcxImagePlugin
 DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing EpsImagePlugin
 DEBUG:PIL.Image:Importing EpsImagePlugin
 DEBUG:PIL.Image:Importing FitsImagePlugin
-DEBUG:PIL.Image:Importing FitsImagePlugin
 DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FliImagePlugin
 DEBUG:PIL.Image:Importing FliImagePlugin
 DEBUG:PIL.Image:Importing FpxImagePlugin
 DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing FtexImagePlugin
 DEBUG:PIL.Image:Importing FtexImagePlugin
 DEBUG:PIL.Image:Importing GbrImagePlugin
 DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GribStubImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
 DEBUG:PIL.Image:Importing GribStubImagePlugin
 DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing IcnsImagePlugin
 DEBUG:PIL.Image:Importing IcnsImagePlugin
 DEBUG:PIL.Image:Importing IcoImagePlugin
-DEBUG:PIL.Image:Importing IcoImagePlugin
 DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImtImagePlugin
 DEBUG:PIL.Image:Importing ImtImagePlugin
 DEBUG:PIL.Image:Importing IptcImagePlugin
 DEBUG:PIL.Image:Importing JpegImagePlugin
 DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing McIdasImagePlugin
-DEBUG:PIL.Image:Importing JpegImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
 DEBUG:PIL.Image:Importing McIdasImagePlugin
 DEBUG:PIL.Image:Importing MicImagePlugin
 DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing MpegImagePlugin
 DEBUG:PIL.Image:Importing MpegImagePlugin
 DEBUG:PIL.Image:Importing MpoImagePlugin
-DEBUG:PIL.Image:Importing MpoImagePlugin
 DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing PalmImagePlugin
 DEBUG:PIL.Image:Importing PalmImagePlugin
 DEBUG:PIL.Image:Importing PcdImagePlugin
 DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PdfImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
 DEBUG:PIL.Image:Importing PdfImagePlugin
 DEBUG:PIL.Image:Importing PixarImagePlugin
 DEBUG:PIL.Image:Importing PngImagePlugin
 DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PsdImagePlugin
-DEBUG:PIL.Image:Importing PngImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
 DEBUG:PIL.Image:Importing PsdImagePlugin
 DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing SgiImagePlugin
 DEBUG:PIL.Image:Importing SgiImagePlugin
 DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SunImagePlugin
 DEBUG:PIL.Image:Importing SunImagePlugin
 DEBUG:PIL.Image:Importing TgaImagePlugin
 DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing WebPImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
 DEBUG:PIL.Image:Importing WebPImagePlugin
 DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing XbmImagePlugin
 DEBUG:PIL.Image:Importing XbmImagePlugin
 DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XVThumbImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  9%|▉         | 1/11 [01:39<16:34, 99.46s/it]
+ 18%|█▊        | 2/11 [03:18<14:55, 99.48s/it]
+ 27%|██▋       | 3/11 [04:58<13:16, 99.60s/it]
+ 36%|███▋      | 4/11 [06:38<11:37, 99.69s/it]
+ 45%|████▌     | 5/11 [08:18<09:58, 99.68s/it]
+ 55%|█████▍    | 6/11 [09:57<08:18, 99.66s/it]
+ 64%|██████▎   | 7/11 [11:37<06:38, 99.62s/it]
+ 73%|███████▎  | 8/11 [13:16<04:58, 99.55s/it]
+ 82%|████████▏ | 9/11 [14:56<03:19, 99.50s/it]
+ 91%|█████████ | 10/11 [16:35<01:39, 99.43s/it]
+100%|██████████| 11/11 [18:14<00:00, 99.36s/it]
+100%|██████████| 11/11 [18:14<00:00, 99.51s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 6: generating actions ...
+>>> Step 6: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_g1_pack_camera/case3/psnr_result.json
+++ b/unitree_g1_pack_camera/case3/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_g1_pack_camera/case3/unitree_g1_pack_camera_case3.mp4",
-    "pred_video": "unitree_g1_pack_camera/case3/output/inference/100_full_fs6.mp4",
-    "psnr": 37.034952654534486
+    "pred_video": "unitree_g1_pack_camera/case3/output/inference/unitree_g1_pack_camera_case3_amd.mp4",
+    "psnr": 19.429578160315536
 }
--- a/unitree_g1_pack_camera/case4/output.log
+++ b/unitree_g1_pack_camera/case4/output.log
@@ -0,0 +1,144 @@
+2026-02-08 05:29:19.728303: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 05:29:19.731620: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 05:29:19.761276: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 05:29:19.761301: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 05:29:19.762880: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 05:29:19.770578: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 05:29:19.771072: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 05:29:21.043661: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/11 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  9%|▉         | 1/11 [01:37<16:18, 97.81s/it]
+ 18%|█▊        | 2/11 [03:15<14:38, 97.56s/it]
+ 27%|██▋       | 3/11 [04:52<12:59, 97.48s/it]
+ 36%|███▋      | 4/11 [06:29<11:21, 97.38s/it]
+ 45%|████▌     | 5/11 [08:06<09:43, 97.28s/it]
+ 55%|█████▍    | 6/11 [09:44<08:06, 97.35s/it]
+ 64%|██████▎   | 7/11 [11:21<06:29, 97.36s/it]
+ 73%|███████▎  | 8/11 [12:59<04:52, 97.38s/it]
+ 82%|████████▏ | 9/11 [14:36<03:14, 97.39s/it]
+ 91%|█████████ | 10/11 [16:14<01:37, 97.42s/it]
+100%|██████████| 11/11 [17:51<00:00, 97.42s/it]
+100%|██████████| 11/11 [17:51<00:00, 97.41s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 6: generating actions ...
+>>> Step 6: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_g1_pack_camera/case4/psnr_result.json
+++ b/unitree_g1_pack_camera/case4/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_g1_pack_camera/case4/unitree_g1_pack_camera_case4.mp4",
-    "pred_video": "unitree_g1_pack_camera/case4/output/inference/200_full_fs6.mp4",
-    "psnr": 31.43390896360405
+    "pred_video": "unitree_g1_pack_camera/case4/output/inference/unitree_g1_pack_camera_case4_amd.mp4",
+    "psnr": 17.80386833747375
 }
--- a/unitree_z1_dual_arm_cleanup_pencils/case1/output.log
+++ b/unitree_z1_dual_arm_cleanup_pencils/case1/output.log
@@ -1,11 +1,15 @@
-2026-02-10 15:38:28.973314: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
-2026-02-10 15:38:29.023024: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
-2026-02-10 15:38:29.023070: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
-2026-02-10 15:38:29.024393: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
-2026-02-10 15:38:29.031901: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
-To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
-2026-02-10 15:38:29.955454: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
-Global seed set to 123
+2026-02-08 12:22:55.885867: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 12:22:55.890510: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 12:22:55.938683: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 12:22:55.938759: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 12:22:55.941091: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 12:22:55.952450: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 12:22:55.952933: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 12:22:56.593653: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
 INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
 INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
 INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
@@ -14,11 +18,24 @@ INFO:root:Loaded ViT-H-14 model config.
 DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
 DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
 INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
 INFO:root:Loaded ViT-H-14 model config.
 DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
 INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:149: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
 >>> model checkpoint loaded.
 >>> Load pre-trained model ...
+>>> Applying precision settings:
+    - Diffusion dtype: bf16
+    - Projector mode: bf16_full
+    - Encoder mode: bf16_full
+    - VAE dtype: bf16
+    ✓ Diffusion model weights converted to bfloat16
+    ✓ Projectors converted to bfloat16
+    ✓ Encoders converted to bfloat16
+    ✓ VAE converted to bfloat16
 INFO:root:***** Configing Data *****
 >>> unitree_z1_stackbox: 1 data samples loaded.
 >>> unitree_z1_stackbox: data stats loaded.
@@ -41,7 +58,13 @@ DEBUG:h5py._conv:Creating converter from 3 to 5
 DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
 DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
 DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
-
+
+  0%|          | 0/8 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
 >>> Step 0: generating actions ...
 >>> Step 0: interacting with world model ...
 DEBUG:PIL.Image:Importing BlpImagePlugin
@@ -92,7 +115,7 @@ DEBUG:PIL.Image:Importing WmfImagePlugin
 DEBUG:PIL.Image:Importing WmfImagePlugin
 DEBUG:PIL.Image:Importing XbmImagePlugin
 DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XVThumbImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin

 12%|█▎        | 1/8 [01:24<09:53, 84.82s/it]
 25%|██▌       | 2/8 [02:49<08:26, 84.48s/it]
@@ -116,6 +139,6 @@ DEBUG:PIL.Image:Importing XVThumbImagePlugin
 >>> Step 4: generating actions ...
 >>> Step 4: interacting with world model ...
 >>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 5: generating actions ...
->>> Step 5: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_z1_dual_arm_cleanup_pencils/case1/psnr_result.json
+++ b/unitree_z1_dual_arm_cleanup_pencils/case1/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_cleanup_pencils/case1/unitree_z1_dual_arm_cleanup_pencils_case1.mp4",
-    "pred_video": "unitree_z1_dual_arm_cleanup_pencils/case1/output/inference/0_full_fs4.mp4",
-    "psnr": 47.911564449209735
+    "pred_video": "unitree_z1_dual_arm_cleanup_pencils/case1/output/inference/unitree_z1_dual_arm_cleanup_pencils_case1_amd.mp4",
+    "psnr": 19.586376345676264
 }
--- a/unitree_z1_dual_arm_cleanup_pencils/case1/psnr_result1.json
+++ b/unitree_z1_dual_arm_cleanup_pencils/case1/psnr_result1.json
@@ -0,0 +1,5 @@
+{
+    "gt_video": "/mnt/ASC1637/unifolm-world-model-action/unitree_z1_dual_arm_cleanup_pencils/case1/output/inference/unitree_z1_dual_arm_cleanup_pencils_case1_amd.mp4",
+    "pred_video": "/mnt/ASC1637/unifolm-world-model-action/unitree_z1_dual_arm_cleanup_pencils/case1/output/inference/0_full_fs4.mp4",
+    "psnr": 30.44844270035179
+}
--- a/unitree_z1_dual_arm_cleanup_pencils/case1/run_world_model_interaction.sh
+++ b/unitree_z1_dual_arm_cleanup_pencils/case1/run_world_model_interaction.sh
@@ -4,7 +4,7 @@ dataset="unitree_z1_dual_arm_cleanup_pencils"
 {
    time CUDA_VISIBLE_DEVICES=0 python3 scripts/evaluation/world_model_interaction.py \
        --seed 123 \
-        --ckpt_path ckpts/unifolm_wma_dual.ckpt \
+        --ckpt_path ckpts/unifolm_wma_dual_mix_bf16.ckpt \
        --config configs/inference/world_model_interaction.yaml \
        --savedir "${res_dir}/output" \
        --bs 1 --height 320 --width 512 \
@@ -20,5 +20,6 @@ dataset="unitree_z1_dual_arm_cleanup_pencils"
        --n_iter 8 \
        --timestep_spacing 'uniform_trailing' \
        --guidance_rescale 0.7 \
-        --perframe_ae
+        --perframe_ae \
+        --vae_dtype bf16
 } 2>&1 | tee "${res_dir}/output.log"
--- a/unitree_z1_dual_arm_cleanup_pencils/case2/output.log
+++ b/unitree_z1_dual_arm_cleanup_pencils/case2/output.log
@@ -0,0 +1,137 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 06:59:34.465946: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 06:59:34.469367: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 06:59:34.500805: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 06:59:34.500837: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 06:59:34.502917: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 06:59:34.511434: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 06:59:34.511678: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 06:59:35.478194: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/8 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+ 12%|█▎        | 1/8 [01:37<11:23, 97.57s/it]
+ 25%|██▌       | 2/8 [03:14<09:44, 97.48s/it]
+ 38%|███▊      | 3/8 [04:52<08:07, 97.47s/it]
+ 50%|█████     | 4/8 [06:29<06:29, 97.49s/it]
+ 62%|██████▎   | 5/8 [08:07<04:52, 97.42s/it]
+ 75%|███████▌  | 6/8 [09:44<03:14, 97.32s/it]
+ 88%|████████▊ | 7/8 [11:21<01:37, 97.34s/it]
+100%|██████████| 8/8 [12:59<00:00, 97.36s/it]
+100%|██████████| 8/8 [12:59<00:00, 97.40s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_z1_dual_arm_cleanup_pencils/case2/psnr_result.json
+++ b/unitree_z1_dual_arm_cleanup_pencils/case2/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_cleanup_pencils/case2/unitree_z1_dual_arm_cleanup_pencils_case2.mp4",
-    "pred_video": "unitree_z1_dual_arm_cleanup_pencils/case2/output/inference/50_full_fs4.mp4",
-    "psnr": 48.344571927558974
+    "pred_video": "unitree_z1_dual_arm_cleanup_pencils/case2/output/inference/unitree_z1_dual_arm_cleanup_pencils_case2_amd.mp4",
+    "psnr": 20.484298972158296
 }
--- a/unitree_z1_dual_arm_cleanup_pencils/case3/output.log
+++ b/unitree_z1_dual_arm_cleanup_pencils/case3/output.log
@@ -0,0 +1,137 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 07:18:52.629976: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 07:18:52.633025: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:18:52.663985: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 07:18:52.664018: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 07:18:52.665837: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 07:18:52.673889: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:18:52.674218: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 07:18:53.298338: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/8 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+ 12%|█▎        | 1/8 [01:40<11:43, 100.54s/it]
+ 25%|██▌       | 2/8 [03:20<10:02, 100.36s/it]
+ 38%|███▊      | 3/8 [05:01<08:21, 100.32s/it]
+ 50%|█████     | 4/8 [06:41<06:41, 100.36s/it]
+ 62%|██████▎   | 5/8 [08:21<05:00, 100.30s/it]
+ 75%|███████▌  | 6/8 [10:01<03:20, 100.28s/it]
+ 88%|████████▊ | 7/8 [11:42<01:40, 100.34s/it]
+100%|██████████| 8/8 [13:22<00:00, 100.36s/it]
+100%|██████████| 8/8 [13:22<00:00, 100.34s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_z1_dual_arm_cleanup_pencils/case3/psnr_result.json
+++ b/unitree_z1_dual_arm_cleanup_pencils/case3/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_cleanup_pencils/case3/unitree_z1_dual_arm_cleanup_pencils_case3.mp4",
-    "pred_video": "unitree_z1_dual_arm_cleanup_pencils/case3/output/inference/100_full_fs4.mp4",
-    "psnr": 41.152374490134825
+    "pred_video": "unitree_z1_dual_arm_cleanup_pencils/case3/output/inference/unitree_z1_dual_arm_cleanup_pencils_case3_amd.mp4",
+    "psnr": 21.20205061239349
 }
--- a/unitree_z1_dual_arm_cleanup_pencils/case4/output.log
+++ b/unitree_z1_dual_arm_cleanup_pencils/case4/output.log
@@ -0,0 +1,137 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 07:22:15.333099: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 07:22:15.336215: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:22:15.366489: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 07:22:15.366522: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 07:22:15.368294: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 07:22:15.376202: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:22:15.376444: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 07:22:15.995383: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/8 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+ 12%|█▎        | 1/8 [01:37<11:23, 97.68s/it]
+ 25%|██▌       | 2/8 [03:15<09:47, 97.83s/it]
+ 38%|███▊      | 3/8 [04:53<08:09, 97.91s/it]
+ 50%|█████     | 4/8 [06:31<06:32, 98.03s/it]
+ 62%|██████▎   | 5/8 [08:10<04:54, 98.11s/it]
+ 75%|███████▌  | 6/8 [09:48<03:16, 98.18s/it]
+ 88%|████████▊ | 7/8 [11:26<01:38, 98.24s/it]
+100%|██████████| 8/8 [13:04<00:00, 98.16s/it]
+100%|██████████| 8/8 [13:04<00:00, 98.09s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_z1_dual_arm_cleanup_pencils/case4/psnr_result.json
+++ b/unitree_z1_dual_arm_cleanup_pencils/case4/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_cleanup_pencils/case4/unitree_z1_dual_arm_cleanup_pencils_case4.mp4",
-    "pred_video": "unitree_z1_dual_arm_cleanup_pencils/case4/output/inference/200_full_fs4.mp4",
-    "psnr": 46.025723557253855
+    "pred_video": "unitree_z1_dual_arm_cleanup_pencils/case4/output/inference/unitree_z1_dual_arm_cleanup_pencils_case4_amd.mp4",
+    "psnr": 21.130122583788612
 }
--- a/unitree_z1_dual_arm_stackbox/case1/output.log
+++ b/unitree_z1_dual_arm_stackbox/case1/output.log
@@ -0,0 +1,134 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 07:24:40.357099: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 07:24:40.360365: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:24:40.391744: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 07:24:40.391772: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 07:24:40.393608: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 07:24:40.401837: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:24:40.402077: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 07:24:41.022382: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/7 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+ 14%|█▍        | 1/7 [01:41<10:09, 101.63s/it]
+ 29%|██▊       | 2/7 [03:20<08:18, 99.78s/it] 
+ 43%|████▎     | 3/7 [04:58<06:36, 99.24s/it]
+ 57%|█████▋    | 4/7 [06:37<04:57, 99.05s/it]
+ 71%|███████▏  | 5/7 [08:16<03:17, 98.90s/it]
+ 86%|████████▌ | 6/7 [09:54<01:38, 98.80s/it]
+100%|██████████| 7/7 [11:33<00:00, 98.70s/it]
+100%|██████████| 7/7 [11:33<00:00, 99.03s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
--- a/unitree_z1_dual_arm_stackbox/case1/psnr_result.json
+++ b/unitree_z1_dual_arm_stackbox/case1/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_stackbox/case1/unitree_z1_dual_arm_stackbox_case1.mp4",
-    "pred_video": "unitree_z1_dual_arm_stackbox/case1/output/inference/5_full_fs4.mp4",
-    "psnr": 44.3480149502738
+    "pred_video": "unitree_z1_dual_arm_stackbox/case1/output/inference/unitree_z1_dual_arm_stackbox_case1_amd.mp4",
+    "psnr": 21.258130518117493
 }
--- a/unitree_z1_dual_arm_stackbox/case1/run_world_model_interaction.sh
+++ b/unitree_z1_dual_arm_stackbox/case1/run_world_model_interaction.sh
@@ -2,7 +2,7 @@ res_dir="unitree_z1_dual_arm_stackbox/case1"
 dataset="unitree_z1_dual_arm_stackbox"

 {
-    time CUDA_VISIBLE_DEVICES=0 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=7 python3 scripts/evaluation/world_model_interaction.py \
        --seed 123 \
        --ckpt_path ckpts/unifolm_wma_dual.ckpt \
        --config configs/inference/world_model_interaction.yaml \
--- a/unitree_z1_dual_arm_stackbox/case2/output.log
+++ b/unitree_z1_dual_arm_stackbox/case2/output.log
@@ -0,0 +1,134 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 07:25:18.653033: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 07:25:18.656060: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:25:18.687077: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 07:25:18.687119: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 07:25:18.688915: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 07:25:18.697008: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:25:18.697255: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 07:25:19.338303: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/7 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+ 14%|█▍        | 1/7 [01:39<09:56, 99.35s/it]
+ 29%|██▊       | 2/7 [03:18<08:17, 99.50s/it]
+ 43%|████▎     | 3/7 [04:58<06:38, 99.54s/it]
+ 57%|█████▋    | 4/7 [06:38<04:58, 99.52s/it]
+ 71%|███████▏  | 5/7 [08:17<03:19, 99.55s/it]
+ 86%|████████▌ | 6/7 [09:57<01:39, 99.53s/it]
+100%|██████████| 7/7 [11:36<00:00, 99.50s/it]
+100%|██████████| 7/7 [11:36<00:00, 99.51s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
--- a/unitree_z1_dual_arm_stackbox/case2/psnr_result.json
+++ b/unitree_z1_dual_arm_stackbox/case2/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_stackbox/case2/unitree_z1_dual_arm_stackbox_case2.mp4",
-    "pred_video": "unitree_z1_dual_arm_stackbox/case2/output/inference/15_full_fs4.mp4",
-    "psnr": 39.867728254007716
+    "pred_video": "unitree_z1_dual_arm_stackbox/case2/output/inference/unitree_z1_dual_arm_stackbox_case2_amd.mp4",
+    "psnr": 23.878153424077645
 }
--- a/unitree_z1_dual_arm_stackbox/case2/run_world_model_interaction.sh
+++ b/unitree_z1_dual_arm_stackbox/case2/run_world_model_interaction.sh
@@ -2,7 +2,7 @@ res_dir="unitree_z1_dual_arm_stackbox/case2"
 dataset="unitree_z1_dual_arm_stackbox"

 {
-    time CUDA_VISIBLE_DEVICES=0 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=6 python3 scripts/evaluation/world_model_interaction.py \
        --seed 123 \
        --ckpt_path ckpts/unifolm_wma_dual.ckpt \
        --config configs/inference/world_model_interaction.yaml \
--- a/unitree_z1_dual_arm_stackbox/case3/output.log
+++ b/unitree_z1_dual_arm_stackbox/case3/output.log
@@ -0,0 +1,134 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 07:35:33.682231: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 07:35:33.685275: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:35:33.716682: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 07:35:33.716728: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 07:35:33.718523: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 07:35:33.726756: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:35:33.727105: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 07:35:34.356722: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/7 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+ 14%|█▍        | 1/7 [01:41<10:06, 101.02s/it]
+ 29%|██▊       | 2/7 [03:23<08:29, 101.84s/it]
+ 43%|████▎     | 3/7 [05:04<06:45, 101.43s/it]
+ 57%|█████▋    | 4/7 [06:45<05:04, 101.42s/it]
+ 71%|███████▏  | 5/7 [08:27<03:22, 101.40s/it]
+ 86%|████████▌ | 6/7 [10:08<01:41, 101.39s/it]
+100%|██████████| 7/7 [11:49<00:00, 101.33s/it]
+100%|██████████| 7/7 [11:49<00:00, 101.39s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
--- a/unitree_z1_dual_arm_stackbox/case3/psnr_result.json
+++ b/unitree_z1_dual_arm_stackbox/case3/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_stackbox/case3/unitree_z1_dual_arm_stackbox_case3.mp4",
-    "pred_video": "unitree_z1_dual_arm_stackbox/case3/output/inference/25_full_fs4.mp4",
-    "psnr": 39.19101039445159
+    "pred_video": "unitree_z1_dual_arm_stackbox/case3/output/inference/unitree_z1_dual_arm_stackbox_case3_amd.mp4",
+    "psnr": 25.400458754751128
 }
--- a/unitree_z1_dual_arm_stackbox/case4/output.log
+++ b/unitree_z1_dual_arm_stackbox/case4/output.log
@@ -0,0 +1,134 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 07:38:45.572744: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 07:38:45.576864: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:38:45.624825: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 07:38:45.624883: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 07:38:45.627150: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 07:38:45.638316: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:38:45.638803: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 07:38:46.426363: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/7 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+ 14%|█▍        | 1/7 [01:38<09:52, 98.73s/it]
+ 29%|██▊       | 2/7 [03:17<08:14, 98.85s/it]
+ 43%|████▎     | 3/7 [04:56<06:35, 98.80s/it]
+ 57%|█████▋    | 4/7 [06:35<04:56, 98.94s/it]
+ 71%|███████▏  | 5/7 [08:14<03:17, 98.93s/it]
+ 86%|████████▌ | 6/7 [09:53<01:38, 98.89s/it]
+100%|██████████| 7/7 [11:31<00:00, 98.81s/it]
+100%|██████████| 7/7 [11:31<00:00, 98.85s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
--- a/unitree_z1_dual_arm_stackbox/case4/psnr_result.json
+++ b/unitree_z1_dual_arm_stackbox/case4/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_stackbox/case4/unitree_z1_dual_arm_stackbox_case4.mp4",
-    "pred_video": "unitree_z1_dual_arm_stackbox/case4/output/inference/35_full_fs4.mp4",
-    "psnr": 40.29563315341769
+    "pred_video": "unitree_z1_dual_arm_stackbox/case4/output/inference/unitree_z1_dual_arm_stackbox_case4_amd.mp4",
+    "psnr": 24.098958457373858
 }
--- a/unitree_z1_dual_arm_stackbox_v2/case1/output.log
+++ b/unitree_z1_dual_arm_stackbox_v2/case1/output.log
@@ -1,13 +1,34 @@
-2026-02-18 18:49:49.117856: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
-2026-02-18 18:49:49.165270: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
-2026-02-18 18:49:49.165322: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
-2026-02-18 18:49:49.166382: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
-2026-02-18 18:49:49.173299: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
-To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
-2026-02-18 18:49:50.090214: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 07:51:23.961486: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 07:51:24.200063: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:51:24.522299: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 07:51:24.522350: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 07:51:24.528237: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 07:51:24.579400: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:51:24.579644: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 07:51:25.781311: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
 Global seed set to 123
->>> Loading prepared model from ckpts/unifolm_wma_dual.ckpt.prepared.pt ...
->>> Prepared model loaded.
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
 INFO:root:***** Configing Data *****
 >>> unitree_z1_stackbox: 1 data samples loaded.
 >>> unitree_z1_stackbox: data stats loaded.
@@ -25,28 +46,19 @@ INFO:root:***** Configing Data *****
 >>> unitree_g1_pack_camera: data stats loaded.
 >>> unitree_g1_pack_camera: normalizer initiated.
 >>> Dataset is successfully loaded ...
-    ✓ KV fused: 66 attention layers
-    TRT output 'y': [1, 4, 16, 40, 64] torch.float32
-    TRT output 'hs_a_0': [1, 16, 320, 40, 64] torch.float32
-    TRT output 'hs_a_1': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_2': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_3': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_4': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_5': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_6': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_7': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_8': [1, 16, 320, 40, 64] torch.float32
->>> TRT backbone loaded from /home/qhy/unifolm-world-model-action/scripts/evaluation/../../trt_engines/video_backbone.engine
 >>> Generate 16 frames under each generation ...
 DEBUG:h5py._conv:Creating converter from 3 to 5
 DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
 DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
 DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
-
-  0%|          | 0/11 [00:00<?, ?it/s][02/18/2026-18:50:03] [TRT] [W] Using default stream in enqueueV3() may lead to performance issues due to additional calls to cudaStreamSynchronize() by TensorRT to ensure correct synchronization. Please use non-default stream instead.
+
+  0%|          | 0/11 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
 >>> Step 0: generating actions ...
-  9%|▉         | 1/11 [00:15<02:38, 15.88s/it]>>> Step 0: generating actions ...
->>> Step 0: interacting with world model ...
 >>> Step 0: interacting with world model ...
 DEBUG:PIL.Image:Importing BlpImagePlugin
 DEBUG:PIL.Image:Importing BmpImagePlugin
@@ -96,7 +108,9 @@ DEBUG:PIL.Image:Importing WmfImagePlugin
 DEBUG:PIL.Image:Importing WmfImagePlugin
 DEBUG:PIL.Image:Importing XbmImagePlugin
 DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  9%|▉         | 1/11 [01:38<16:20, 98.04s/it]
 18%|█▊        | 2/11 [03:15<14:40, 97.81s/it]
 27%|██▋       | 3/11 [04:53<13:01, 97.72s/it]
 36%|███▋      | 4/11 [06:31<11:24, 97.71s/it]
@@ -127,6 +141,6 @@ DEBUG:PIL.Image:Importing XVThumbImagePlugin
 >>> Step 6: generating actions ...
 >>> Step 6: interacting with world model ...
 >>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 7: generating actions ...
->>> Step 7: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_z1_dual_arm_stackbox_v2/case1/psnr_result.json
+++ b/unitree_z1_dual_arm_stackbox_v2/case1/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_stackbox_v2/case1/unitree_z1_dual_arm_stackbox_v2_case1.mp4",
-    "pred_video": "unitree_z1_dual_arm_stackbox_v2/case1/output/inference/5_full_fs4.mp4",
-    "psnr": 27.62636266067224
+    "pred_video": "unitree_z1_dual_arm_stackbox_v2/case1/output/inference/unitree_z1_dual_arm_stackbox_v2_case1_amd.mp4",
+    "psnr": 18.126776535969576
 }
--- a/unitree_z1_dual_arm_stackbox_v2/case1/run_world_model_interaction.sh
+++ b/unitree_z1_dual_arm_stackbox_v2/case1/run_world_model_interaction.sh
@@ -2,7 +2,7 @@ res_dir="unitree_z1_dual_arm_stackbox_v2/case1"
 dataset="unitree_z1_dual_arm_stackbox_v2"

 {
-    time CUDA_VISIBLE_DEVICES=0 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=7 python3 scripts/evaluation/world_model_interaction.py \
        --seed 123 \
        --ckpt_path ckpts/unifolm_wma_dual.ckpt \
        --config configs/inference/world_model_interaction.yaml \
@@ -20,6 +20,5 @@ dataset="unitree_z1_dual_arm_stackbox_v2"
        --n_iter 11 \
        --timestep_spacing 'uniform_trailing' \
        --guidance_rescale 0.7 \
-        --perframe_ae \
-        --fast_policy_no_decode
+        --perframe_ae
 } 2>&1 | tee "${res_dir}/output.log"
--- a/unitree_z1_dual_arm_stackbox_v2/case2/output.log
+++ b/unitree_z1_dual_arm_stackbox_v2/case2/output.log
@@ -0,0 +1,146 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 07:56:31.144789: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 07:56:31.148256: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:56:31.178870: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 07:56:31.178898: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 07:56:31.180683: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 07:56:31.188800: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:56:31.189142: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 07:56:31.810098: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/11 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  9%|▉         | 1/11 [01:40<16:41, 100.16s/it]
+ 18%|█▊        | 2/11 [03:20<15:04, 100.47s/it]
+ 27%|██▋       | 3/11 [05:01<13:24, 100.62s/it]
+ 36%|███▋      | 4/11 [06:42<11:44, 100.69s/it]
+ 45%|████▌     | 5/11 [08:22<10:02, 100.48s/it]
+ 55%|█████▍    | 6/11 [10:02<08:21, 100.33s/it]
+ 64%|██████▎   | 7/11 [11:42<06:40, 100.23s/it]
+ 73%|███████▎  | 8/11 [13:22<05:00, 100.23s/it]
+ 82%|████████▏ | 9/11 [15:03<03:20, 100.23s/it]
+ 91%|█████████ | 10/11 [16:43<01:40, 100.33s/it]
+100%|██████████| 11/11 [18:24<00:00, 100.41s/it]
+100%|██████████| 11/11 [18:24<00:00, 100.39s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 6: generating actions ...
+>>> Step 6: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_z1_dual_arm_stackbox_v2/case2/psnr_result.json
+++ b/unitree_z1_dual_arm_stackbox_v2/case2/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_stackbox_v2/case2/unitree_z1_dual_arm_stackbox_v2_case2.mp4",
-    "pred_video": "unitree_z1_dual_arm_stackbox_v2/case2/output/inference/15_full_fs4.mp4",
-    "psnr": 33.90444714332389
+    "pred_video": "unitree_z1_dual_arm_stackbox_v2/case2/output/inference/unitree_z1_dual_arm_stackbox_v2_case2_amd.mp4",
+    "psnr": 19.38130614773096
 }
--- a/unitree_z1_dual_arm_stackbox_v2/case3/output.log
+++ b/unitree_z1_dual_arm_stackbox_v2/case3/output.log
@@ -0,0 +1,146 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 07:56:04.467082: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 07:56:04.470145: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:56:04.502248: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 07:56:04.502277: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 07:56:04.504088: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 07:56:04.512557: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 07:56:04.512830: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 07:56:05.259641: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/11 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  9%|▉         | 1/11 [01:38<16:20, 98.03s/it]
+ 18%|█▊        | 2/11 [03:16<14:43, 98.19s/it]
+ 27%|██▋       | 3/11 [04:55<13:08, 98.54s/it]
+ 36%|███▋      | 4/11 [06:33<11:29, 98.52s/it]
+ 45%|████▌     | 5/11 [08:11<09:50, 98.38s/it]
+ 55%|█████▍    | 6/11 [09:49<08:10, 98.11s/it]
+ 64%|██████▎   | 7/11 [11:27<06:31, 97.97s/it]
+ 73%|███████▎  | 8/11 [13:04<04:53, 97.83s/it]
+ 82%|████████▏ | 9/11 [14:42<03:15, 97.72s/it]
+ 91%|█████████ | 10/11 [16:19<01:37, 97.71s/it]
+100%|██████████| 11/11 [17:57<00:00, 97.74s/it]
+100%|██████████| 11/11 [17:57<00:00, 97.97s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 6: generating actions ...
+>>> Step 6: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_z1_dual_arm_stackbox_v2/case3/psnr_result.json
+++ b/unitree_z1_dual_arm_stackbox_v2/case3/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_stackbox_v2/case3/unitree_z1_dual_arm_stackbox_v2_case3.mp4",
-    "pred_video": "unitree_z1_dual_arm_stackbox_v2/case3/output/inference/25_full_fs4.mp4",
-    "psnr": 34.50192428908007
+    "pred_video": "unitree_z1_dual_arm_stackbox_v2/case3/output/inference/unitree_z1_dual_arm_stackbox_v2_case3_amd.mp4",
+    "psnr": 18.74462122425683
 }
--- a/unitree_z1_dual_arm_stackbox_v2/case4/output.log
+++ b/unitree_z1_dual_arm_stackbox_v2/case4/output.log
@@ -1,13 +1,34 @@
-2026-02-18 18:54:56.403136: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
-2026-02-18 18:54:56.451144: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
-2026-02-18 18:54:56.451189: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
-2026-02-18 18:54:56.452312: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
-2026-02-18 18:54:56.459281: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
-To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
-2026-02-18 18:54:57.381032: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 08:04:16.104516: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 08:04:16.109112: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:04:16.138703: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 08:04:16.138737: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 08:04:16.140302: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 08:04:16.147672: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:04:16.147903: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 08:04:17.363218: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
 Global seed set to 123
->>> Loading prepared model from ckpts/unifolm_wma_dual.ckpt.prepared.pt ...
->>> Prepared model loaded.
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
 INFO:root:***** Configing Data *****
 >>> unitree_z1_stackbox: 1 data samples loaded.
 >>> unitree_z1_stackbox: data stats loaded.
@@ -25,125 +46,71 @@ INFO:root:***** Configing Data *****
 >>> unitree_g1_pack_camera: data stats loaded.
 >>> unitree_g1_pack_camera: normalizer initiated.
 >>> Dataset is successfully loaded ...
-    ✓ KV fused: 66 attention layers
-    TRT output 'y': [1, 4, 16, 40, 64] torch.float32
-    TRT output 'hs_a_0': [1, 16, 320, 40, 64] torch.float32
-    TRT output 'hs_a_1': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_2': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_3': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_4': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_5': [1, 16, 1280, 5, 8] torch.float32
-    TRT output 'hs_a_6': [1, 16, 1280, 10, 16] torch.float32
-    TRT output 'hs_a_7': [1, 16, 640, 20, 32] torch.float32
-    TRT output 'hs_a_8': [1, 16, 320, 40, 64] torch.float32
->>> TRT backbone loaded from /home/qhy/unifolm-world-model-action/scripts/evaluation/../../trt_engines/video_backbone.engine
 >>> Generate 16 frames under each generation ...
 DEBUG:h5py._conv:Creating converter from 3 to 5
 DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
 DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
 DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
-
-  0%|          | 0/11 [00:00<?, ?it/s][02/18/2026-18:55:10] [TRT] [W] Using default stream in enqueueV3() may lead to performance issues due to additional calls to cudaStreamSynchronize() by TensorRT to ensure correct synchronization. Please use non-default stream instead.
+
+  0%|          | 0/11 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
 >>> Step 0: generating actions ...
-  9%|▉         | 1/11 [00:16<02:45, 16.53s/it]>>> Step 0: generating actions ...
->>> Step 0: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
 >>> Step 0: interacting with world model ...
 DEBUG:PIL.Image:Importing BlpImagePlugin
 DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
-DEBUG:PIL.Image:Importing BufrStubImagePlugin
-DEBUG:PIL.Image:Importing BmpImagePlugin
 DEBUG:PIL.Image:Importing BufrStubImagePlugin
 DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing CurImagePlugin
-DEBUG:PIL.Image:Importing DcxImagePlugin
 DEBUG:PIL.Image:Importing DcxImagePlugin
 DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing DdsImagePlugin
-DEBUG:PIL.Image:Importing EpsImagePlugin
 DEBUG:PIL.Image:Importing EpsImagePlugin
 DEBUG:PIL.Image:Importing FitsImagePlugin
-DEBUG:PIL.Image:Importing FitsImagePlugin
 DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FitsStubImagePlugin
-DEBUG:PIL.Image:Importing FliImagePlugin
 DEBUG:PIL.Image:Importing FliImagePlugin
 DEBUG:PIL.Image:Importing FpxImagePlugin
 DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing FtexImagePlugin
 DEBUG:PIL.Image:Importing FtexImagePlugin
 DEBUG:PIL.Image:Importing GbrImagePlugin
 DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
-DEBUG:PIL.Image:Importing GribStubImagePlugin
-DEBUG:PIL.Image:Importing GifImagePlugin
 DEBUG:PIL.Image:Importing GribStubImagePlugin
 DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
-DEBUG:PIL.Image:Importing IcnsImagePlugin
 DEBUG:PIL.Image:Importing IcnsImagePlugin
 DEBUG:PIL.Image:Importing IcoImagePlugin
-DEBUG:PIL.Image:Importing IcoImagePlugin
 DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImImagePlugin
-DEBUG:PIL.Image:Importing ImtImagePlugin
 DEBUG:PIL.Image:Importing ImtImagePlugin
 DEBUG:PIL.Image:Importing IptcImagePlugin
 DEBUG:PIL.Image:Importing JpegImagePlugin
 DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
-DEBUG:PIL.Image:Importing McIdasImagePlugin
-DEBUG:PIL.Image:Importing JpegImagePlugin
-DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
 DEBUG:PIL.Image:Importing McIdasImagePlugin
 DEBUG:PIL.Image:Importing MicImagePlugin
 DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
-DEBUG:PIL.Image:Importing MpegImagePlugin
 DEBUG:PIL.Image:Importing MpegImagePlugin
 DEBUG:PIL.Image:Importing MpoImagePlugin
-DEBUG:PIL.Image:Importing MpoImagePlugin
 DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing MspImagePlugin
-DEBUG:PIL.Image:Importing PalmImagePlugin
 DEBUG:PIL.Image:Importing PalmImagePlugin
 DEBUG:PIL.Image:Importing PcdImagePlugin
 DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
-DEBUG:PIL.Image:Importing PdfImagePlugin
-DEBUG:PIL.Image:Importing PcxImagePlugin
 DEBUG:PIL.Image:Importing PdfImagePlugin
 DEBUG:PIL.Image:Importing PixarImagePlugin
 DEBUG:PIL.Image:Importing PngImagePlugin
 DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
-DEBUG:PIL.Image:Importing PsdImagePlugin
-DEBUG:PIL.Image:Importing PngImagePlugin
-DEBUG:PIL.Image:Importing PpmImagePlugin
 DEBUG:PIL.Image:Importing PsdImagePlugin
 DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing QoiImagePlugin
-DEBUG:PIL.Image:Importing SgiImagePlugin
 DEBUG:PIL.Image:Importing SgiImagePlugin
 DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SpiderImagePlugin
-DEBUG:PIL.Image:Importing SunImagePlugin
 DEBUG:PIL.Image:Importing SunImagePlugin
 DEBUG:PIL.Image:Importing TgaImagePlugin
 DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
-DEBUG:PIL.Image:Importing WebPImagePlugin
-DEBUG:PIL.Image:Importing TiffImagePlugin
 DEBUG:PIL.Image:Importing WebPImagePlugin
 DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing WmfImagePlugin
-DEBUG:PIL.Image:Importing XbmImagePlugin
 DEBUG:PIL.Image:Importing XbmImagePlugin
 DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XVThumbImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  9%|▉         | 1/11 [01:39<16:32, 99.26s/it]
 18%|█▊        | 2/11 [03:17<14:49, 98.81s/it]
 27%|██▋       | 3/11 [04:56<13:10, 98.76s/it]
 36%|███▋      | 4/11 [06:35<11:31, 98.80s/it]
@@ -174,6 +141,6 @@ DEBUG:PIL.Image:Importing XVThumbImagePlugin
 >>> Step 6: generating actions ...
 >>> Step 6: interacting with world model ...
 >>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 7: generating actions ...
->>> Step 7: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
--- a/unitree_z1_dual_arm_stackbox_v2/case4/psnr_result.json
+++ b/unitree_z1_dual_arm_stackbox_v2/case4/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_dual_arm_stackbox_v2/case4/unitree_z1_dual_arm_stackbox_v2_case4.mp4",
-    "pred_video": "unitree_z1_dual_arm_stackbox_v2/case4/output/inference/35_full_fs4.mp4",
-    "psnr": 25.49270910031428
+    "pred_video": "unitree_z1_dual_arm_stackbox_v2/case4/output/inference/unitree_z1_dual_arm_stackbox_v2_case4_amd.mp4",
+    "psnr": 19.526448380726254
 }
--- a/unitree_z1_dual_arm_stackbox_v2/case4/run_world_model_interaction.sh
+++ b/unitree_z1_dual_arm_stackbox_v2/case4/run_world_model_interaction.sh
@@ -2,7 +2,7 @@ res_dir="unitree_z1_dual_arm_stackbox_v2/case4"
 dataset="unitree_z1_dual_arm_stackbox_v2"

 {
-    time CUDA_VISIBLE_DEVICES=0 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=6 python3 scripts/evaluation/world_model_interaction.py \
        --seed 123 \
        --ckpt_path ckpts/unifolm_wma_dual.ckpt \
        --config configs/inference/world_model_interaction.yaml \
--- a/unitree_z1_stackbox/case1/output.log
+++ b/unitree_z1_stackbox/case1/output.log
@@ -0,0 +1,149 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 08:12:47.424053: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 08:12:47.427280: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:12:47.458253: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 08:12:47.458288: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 08:12:47.462758: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 08:12:47.518283: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:12:47.518566: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 08:12:48.593011: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/12 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  8%|▊         | 1/12 [01:38<18:08, 98.94s/it]
+ 17%|█▋        | 2/12 [03:18<16:30, 99.01s/it]
+ 25%|██▌       | 3/12 [04:57<14:51, 99.07s/it]
+ 33%|███▎      | 4/12 [06:36<13:12, 99.04s/it]
+ 42%|████▏     | 5/12 [08:15<11:33, 99.00s/it]
+ 50%|█████     | 6/12 [09:54<09:54, 99.10s/it]
+ 58%|█████▊    | 7/12 [11:33<08:14, 99.00s/it]
+ 67%|██████▋   | 8/12 [13:13<06:38, 99.58s/it]
+ 75%|███████▌  | 9/12 [14:54<04:59, 99.88s/it]
+ 83%|████████▎ | 10/12 [16:33<03:19, 99.58s/it]
+ 92%|█████████▏| 11/12 [18:12<01:39, 99.39s/it]
+100%|██████████| 12/12 [19:51<00:00, 99.25s/it]
+100%|██████████| 12/12 [19:51<00:00, 99.28s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 6: generating actions ...
+>>> Step 6: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 8: generating actions ...
+>>> Step 8: interacting with world model ...
--- a/unitree_z1_stackbox/case1/psnr_result.json
+++ b/unitree_z1_stackbox/case1/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_stackbox/case1/unitree_z1_stackbox_case1.mp4",
-    "pred_video": "unitree_z1_stackbox/case1/output/inference/5_full_fs4.mp4",
-    "psnr": 42.83913947323794
+    "pred_video": "unitree_z1_stackbox/case1/output/inference/unitree_z1_stackbox_case1_amd.mp4",
+    "psnr": 19.81391789862606
 }
--- a/unitree_z1_stackbox/case1/run_world_model_interaction.sh
+++ b/unitree_z1_stackbox/case1/run_world_model_interaction.sh
@@ -2,7 +2,7 @@ res_dir="unitree_z1_stackbox/case1"
 dataset="unitree_z1_stackbox"

 {
-    time CUDA_VISIBLE_DEVICES=0 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=5 python3 scripts/evaluation/world_model_interaction.py \
        --seed 123 \
        --ckpt_path ckpts/unifolm_wma_dual.ckpt \
        --config configs/inference/world_model_interaction.yaml \
--- a/unitree_z1_stackbox/case2/output.log
+++ b/unitree_z1_stackbox/case2/output.log
@@ -0,0 +1,149 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 08:15:49.934949: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 08:15:49.937974: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:15:49.969069: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 08:15:49.969100: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 08:15:49.970909: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 08:15:49.979005: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:15:49.979255: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 08:15:50.597743: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/12 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  8%|▊         | 1/12 [01:37<17:51, 97.37s/it]
+ 17%|█▋        | 2/12 [03:14<16:13, 97.31s/it]
+ 25%|██▌       | 3/12 [04:51<14:35, 97.26s/it]
+ 33%|███▎      | 4/12 [06:29<12:58, 97.25s/it]
+ 42%|████▏     | 5/12 [08:06<11:20, 97.24s/it]
+ 50%|█████     | 6/12 [09:43<09:43, 97.24s/it]
+ 58%|█████▊    | 7/12 [11:20<08:06, 97.27s/it]
+ 67%|██████▋   | 8/12 [12:58<06:29, 97.36s/it]
+ 75%|███████▌  | 9/12 [14:36<04:52, 97.49s/it]
+ 83%|████████▎ | 10/12 [16:13<03:15, 97.52s/it]
+ 92%|█████████▏| 11/12 [17:51<01:37, 97.47s/it]
+100%|██████████| 12/12 [19:28<00:00, 97.35s/it]
+100%|██████████| 12/12 [19:28<00:00, 97.35s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 6: generating actions ...
+>>> Step 6: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 8: generating actions ...
+>>> Step 8: interacting with world model ...
--- a/unitree_z1_stackbox/case2/psnr_result.json
+++ b/unitree_z1_stackbox/case2/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_stackbox/case2/unitree_z1_stackbox_case2.mp4",
-    "pred_video": "unitree_z1_stackbox/case2/output/inference/15_full_fs4.mp4",
-    "psnr": 48.64571989587276
+    "pred_video": "unitree_z1_stackbox/case2/output/inference/unitree_z1_stackbox_case2_amd.mp4",
+    "psnr": 21.083821459054743
 }
--- a/unitree_z1_stackbox/case3/output.log
+++ b/unitree_z1_stackbox/case3/output.log
@@ -0,0 +1,149 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 08:16:22.299521: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 08:16:22.302545: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:16:22.335354: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 08:16:22.335389: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 08:16:22.337179: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 08:16:22.345296: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:16:22.345548: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 08:16:23.008743: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+[rank: 0] Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/12 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  8%|▊         | 1/12 [01:39<18:16, 99.64s/it]
+ 17%|█▋        | 2/12 [03:19<16:35, 99.56s/it]
+ 25%|██▌       | 3/12 [04:58<14:55, 99.53s/it]
+ 33%|███▎      | 4/12 [06:38<13:16, 99.53s/it]
+ 42%|████▏     | 5/12 [08:17<11:36, 99.54s/it]
+ 50%|█████     | 6/12 [09:57<09:57, 99.57s/it]
+ 58%|█████▊    | 7/12 [11:37<08:18, 99.66s/it]
+ 67%|██████▋   | 8/12 [13:17<06:39, 99.83s/it]
+ 75%|███████▌  | 9/12 [14:57<04:59, 99.93s/it]
+ 83%|████████▎ | 10/12 [16:37<03:19, 99.97s/it]
+ 92%|█████████▏| 11/12 [18:17<01:39, 99.85s/it]
+100%|██████████| 12/12 [19:56<00:00, 99.71s/it]
+100%|██████████| 12/12 [19:56<00:00, 99.71s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 6: generating actions ...
+>>> Step 6: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 8: generating actions ...
+>>> Step 8: interacting with world model ...
--- a/unitree_z1_stackbox/case3/psnr_result.json
+++ b/unitree_z1_stackbox/case3/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_stackbox/case3/unitree_z1_stackbox_case3.mp4",
-    "pred_video": "unitree_z1_stackbox/case3/output/inference/25_full_fs4.mp4",
-    "psnr": 45.127553229898034
+    "pred_video": "unitree_z1_stackbox/case3/output/inference/unitree_z1_stackbox_case3_amd.mp4",
+    "psnr": 21.322784880212172
 }
--- a/unitree_z1_stackbox/case4/output.log
+++ b/unitree_z1_stackbox/case4/output.log
@@ -0,0 +1,149 @@
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
+  __import__("pkg_resources").declare_namespace(__name__)
+2026-02-08 08:25:54.657305: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-08 08:25:54.660628: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:25:54.691237: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-08 08:25:54.691275: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-08 08:25:54.693046: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-08 08:25:54.701142: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-08 08:25:54.701413: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+2026-02-08 08:25:55.801367: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+Global seed set to 123
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
+  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+INFO:mainlogger:LatentVisualDiffusion: Running in v-prediction mode
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+INFO:unifolm_wma.models.diffusion_head.conditional_unet1d:number of parameters: 5.010531e+08
+AE working on z of shape (1, 4, 32, 32) = 4096 dimensions.
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): hf-mirror.com:443
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  checkpoint = torch.load(checkpoint_path, map_location=map_location)
+INFO:root:Loaded ViT-H-14 model config.
+DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0
+INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
+/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:86: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
+  state_dict = torch.load(ckpt, map_location="cpu")
+>>> model checkpoint loaded.
+>>> Load pre-trained model ...
+INFO:root:***** Configing Data *****
+>>> unitree_z1_stackbox: 1 data samples loaded.
+>>> unitree_z1_stackbox: data stats loaded.
+>>> unitree_z1_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox: normalizer initiated.
+>>> unitree_z1_dual_arm_stackbox_v2: 1 data samples loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: data stats loaded.
+>>> unitree_z1_dual_arm_stackbox_v2: normalizer initiated.
+>>> unitree_z1_dual_arm_cleanup_pencils: 1 data samples loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: data stats loaded.
+>>> unitree_z1_dual_arm_cleanup_pencils: normalizer initiated.
+>>> unitree_g1_pack_camera: 1 data samples loaded.
+>>> unitree_g1_pack_camera: data stats loaded.
+>>> unitree_g1_pack_camera: normalizer initiated.
+>>> Dataset is successfully loaded ...
+>>> Generate 16 frames under each generation ...
+DEBUG:h5py._conv:Creating converter from 3 to 5
+DEBUG:PIL.PngImagePlugin:STREAM b'IHDR' 16 13
+DEBUG:PIL.PngImagePlugin:STREAM b'pHYs' 41 9
+DEBUG:PIL.PngImagePlugin:STREAM b'IDAT' 62 4096
+
+  0%|          | 0/12 [00:00<?, ?it/s]/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:5501: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at ../aten/src/ATen/Context.cpp:296.)
+  proj = linear(q, w, b)
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Flash attention support on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:225.)
+  attn_output = scaled_dot_product_attention(
+/mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/torch/nn/functional.py:6278: UserWarning: Memory Efficient attention on Navi31 GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:269.)
+  attn_output = scaled_dot_product_attention(
+>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
+DEBUG:PIL.Image:Importing BlpImagePlugin
+DEBUG:PIL.Image:Importing BmpImagePlugin
+DEBUG:PIL.Image:Importing BufrStubImagePlugin
+DEBUG:PIL.Image:Importing CurImagePlugin
+DEBUG:PIL.Image:Importing DcxImagePlugin
+DEBUG:PIL.Image:Importing DdsImagePlugin
+DEBUG:PIL.Image:Importing EpsImagePlugin
+DEBUG:PIL.Image:Importing FitsImagePlugin
+DEBUG:PIL.Image:Importing FitsStubImagePlugin
+DEBUG:PIL.Image:Importing FliImagePlugin
+DEBUG:PIL.Image:Importing FpxImagePlugin
+DEBUG:PIL.Image:Image: failed to import FpxImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing FtexImagePlugin
+DEBUG:PIL.Image:Importing GbrImagePlugin
+DEBUG:PIL.Image:Importing GifImagePlugin
+DEBUG:PIL.Image:Importing GribStubImagePlugin
+DEBUG:PIL.Image:Importing Hdf5StubImagePlugin
+DEBUG:PIL.Image:Importing IcnsImagePlugin
+DEBUG:PIL.Image:Importing IcoImagePlugin
+DEBUG:PIL.Image:Importing ImImagePlugin
+DEBUG:PIL.Image:Importing ImtImagePlugin
+DEBUG:PIL.Image:Importing IptcImagePlugin
+DEBUG:PIL.Image:Importing JpegImagePlugin
+DEBUG:PIL.Image:Importing Jpeg2KImagePlugin
+DEBUG:PIL.Image:Importing McIdasImagePlugin
+DEBUG:PIL.Image:Importing MicImagePlugin
+DEBUG:PIL.Image:Image: failed to import MicImagePlugin: No module named 'olefile'
+DEBUG:PIL.Image:Importing MpegImagePlugin
+DEBUG:PIL.Image:Importing MpoImagePlugin
+DEBUG:PIL.Image:Importing MspImagePlugin
+DEBUG:PIL.Image:Importing PalmImagePlugin
+DEBUG:PIL.Image:Importing PcdImagePlugin
+DEBUG:PIL.Image:Importing PcxImagePlugin
+DEBUG:PIL.Image:Importing PdfImagePlugin
+DEBUG:PIL.Image:Importing PixarImagePlugin
+DEBUG:PIL.Image:Importing PngImagePlugin
+DEBUG:PIL.Image:Importing PpmImagePlugin
+DEBUG:PIL.Image:Importing PsdImagePlugin
+DEBUG:PIL.Image:Importing QoiImagePlugin
+DEBUG:PIL.Image:Importing SgiImagePlugin
+DEBUG:PIL.Image:Importing SpiderImagePlugin
+DEBUG:PIL.Image:Importing SunImagePlugin
+DEBUG:PIL.Image:Importing TgaImagePlugin
+DEBUG:PIL.Image:Importing TiffImagePlugin
+DEBUG:PIL.Image:Importing WebPImagePlugin
+DEBUG:PIL.Image:Importing WmfImagePlugin
+DEBUG:PIL.Image:Importing XbmImagePlugin
+DEBUG:PIL.Image:Importing XpmImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
+
+  8%|▊         | 1/12 [01:37<17:51, 97.38s/it]
+ 17%|█▋        | 2/12 [03:14<16:12, 97.24s/it]
+ 25%|██▌       | 3/12 [04:51<14:35, 97.28s/it]
+ 33%|███▎      | 4/12 [06:29<12:59, 97.40s/it]
+ 42%|████▏     | 5/12 [08:06<11:21, 97.30s/it]
+ 50%|█████     | 6/12 [09:43<09:43, 97.17s/it]
+ 58%|█████▊    | 7/12 [11:20<08:05, 97.07s/it]
+ 67%|██████▋   | 8/12 [12:57<06:28, 97.02s/it]
+ 75%|███████▌  | 9/12 [14:34<04:50, 96.98s/it]
+ 83%|████████▎ | 10/12 [16:11<03:14, 97.00s/it]
+ 92%|█████████▏| 11/12 [17:48<01:37, 97.06s/it]
+100%|██████████| 12/12 [19:25<00:00, 97.13s/it]
+100%|██████████| 12/12 [19:25<00:00, 97.14s/it]
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 1: generating actions ...
+>>> Step 1: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 2: generating actions ...
+>>> Step 2: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 3: generating actions ...
+>>> Step 3: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 4: generating actions ...
+>>> Step 4: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 5: generating actions ...
+>>> Step 5: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 6: generating actions ...
+>>> Step 6: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 7: generating actions ...
+>>> Step 7: interacting with world model ...
+>>>>>>>>>>>>>>>>>>>>>>>>
+>>> Step 8: generating actions ...
+>>> Step 8: interacting with world model ...
--- a/unitree_z1_stackbox/case4/psnr_result.json
+++ b/unitree_z1_stackbox/case4/psnr_result.json
@@ -1,5 +1,5 @@
 {
    "gt_video": "unitree_z1_stackbox/case4/unitree_z1_stackbox_case4.mp4",
-    "pred_video": "unitree_z1_stackbox/case4/output/inference/35_full_fs4.mp4",
-    "psnr": 50.642542240144444
+    "pred_video": "unitree_z1_stackbox/case4/output/inference/unitree_z1_stackbox_case4_amd.mp4",
+    "psnr": 25.32928948331741
 }
--- a/unitree_z1_stackbox/case4/run_world_model_interaction.sh
+++ b/unitree_z1_stackbox/case4/run_world_model_interaction.sh
@@ -2,7 +2,7 @@ res_dir="unitree_z1_stackbox/case4"
 dataset="unitree_z1_stackbox"

 {
-    time CUDA_VISIBLE_DEVICES=0 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=7 python3 scripts/evaluation/world_model_interaction.py \
        --seed 123 \
        --ckpt_path ckpts/unifolm_wma_dual.ckpt \
        --config configs/inference/world_model_interaction.yaml \
Author	SHA1	Message	Date
olivame	e588182642	修复混合精度vae相关的配置错误，确保在推理阶段正确使用了混合精度模型，并且导出了正确精度的检查点文件。	2026-02-08 12:35:59 +00:00
olivame	e6c55a648c	所有case的baseline，amd版本的ground truth都上传了	2026-02-08 09:42:14 +00:00
				`@@ -0,0 +1 @@`
				`{"framework": "pytorch", "task": "robotics", "allow_remote": true}`