KV 融合实现完成。改动总结：速度微弱提升psnr略微上升

attention.py — 3处改动： 1. __init__ 添加 _kv_fused = False 标志 2.新增 fuse_kv() 方法：将 to_k + to_v → to_kv，同时处理 _ip/_as/_aa 辅助 KV 对 2. bmm_forward 两个分支加_kv_fused 判断，用to_kv().chunk(2, dim=-1) 替代分别调用
fix bugs
2026-02-11 12:36:38 +08:00 · 2026-02-10 22:35:45 +08:00 · 2026-02-10 21:41:48 +08:00 · 2026-02-10 20:01:25 +08:00 · 2026-02-10 19:54:53 +08:00 · 2026-02-10 19:36:17 +08:00
65 changed files with 616 additions and 3473 deletions
--- a/.claude/settings.local.json
+++ b/.claude/settings.local.json
@@ -0,0 +1,15 @@
+{
+  "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
@@ -0,0 +1,2 @@
+eval "$(conda shell.bash hook 2>/dev/null)"
+conda activate unifolm-wma
--- a/.gitignore
+++ b/.gitignore
@@ -55,7 +55,6 @@ coverage.xml
 *.pot

 # Django stuff:
-
 local_settings.py
 db.sqlite3

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

 *.ckpt
-*.0
+
+*.0
+ckpts/unifolm_wma_dual.ckpt.prepared.pt
--- a/case4_run.log
+++ b/case4_run.log
@@ -1,135 +0,0 @@
-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
@@ -1 +0,0 @@
-{"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: '/mnt/ASC1637/unifolm-world-model-action/examples/world_model_interaction_prompts'
+        data_dir: '/home/qhy/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
@@ -1,21 +0,0 @@
-# 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/run.log
+++ b/run.log
@@ -1,150 +0,0 @@
-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/scripts/evaluation/base_model_inference.py
+++ b/scripts/evaluation/base_model_inference.py
@@ -16,6 +16,9 @@ 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,6 +19,9 @@ 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,5 +1,4 @@
 import argparse, os, glob
-from contextlib import nullcontext
 import pandas as pd
 import random
 import torch
@@ -10,6 +9,8 @@ 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
@@ -17,8 +18,12 @@ from tqdm import tqdm
 from einops import rearrange, repeat
 from collections import OrderedDict
 from torch import nn
-from eval_utils import populate_queues, log_to_tensorboard
+from eval_utils import populate_queues
 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
@@ -39,68 +44,6 @@ 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.

@@ -213,6 +156,81 @@ 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.

@@ -325,11 +343,6 @@ 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
@@ -395,7 +408,8 @@ def image_guided_synthesis_sim_mode(
        timestep_spacing: str = 'uniform',
        guidance_rescale: float = 0.0,
        sim_mode: bool = True,
-        **kwargs) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        decode_video: bool = True,
+        **kwargs) -> tuple[torch.Tensor | None, torch.Tensor, torch.Tensor]:
    """
    Performs image-guided video generation in a simulation-style mode with optional multimodal guidance (image, state, action, text).

@@ -418,10 +432,13 @@ 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): Predicted pixel-space video frames [B, C, T, H, W].
+        batch_variants (torch.Tensor | None): Predicted pixel-space video frames [B, C, T, H, W],
+            or None when decode_video=False.
        actions (torch.Tensor): Predicted action sequences [B, T, D] from diffusion decoding.
        states (torch.Tensor): Predicted state sequences [B, T, D] from diffusion decoding.
    """
@@ -431,22 +448,10 @@ 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:].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)
+    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)

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

-    # Auto-detect projector dtype and convert inputs
-    projector_dtype = next(model.state_projector.parameters()).dtype
+    cond_state_emb = model.state_projector(observation['observation.state'])
+    cond_state_emb = cond_state_emb + model.agent_state_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
+    cond_action_emb = model.action_projector(observation['action'])
+    cond_action_emb = cond_action_emb + model.agent_action_pos_emb

    if not sim_mode:
        cond_action_emb = torch.zeros_like(cond_action_emb)
@@ -497,17 +491,9 @@ def image_guided_synthesis_sim_mode(
    kwargs.update({"unconditional_conditioning_img_nonetext": None})
    cond_mask = None
    cond_z0 = 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()
-
+    batch_variants = None
    if ddim_sampler is not None:
-        with autocast_ctx:
-            samples, actions, states, intermedia = ddim_sampler.sample(
+        samples, actions, states, intermedia = ddim_sampler.sample(
            S=ddim_steps,
            conditioning=cond,
            batch_size=batch_size,
@@ -524,9 +510,10 @@ def image_guided_synthesis_sim_mode(
            guidance_rescale=guidance_rescale,
            **kwargs)

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

    return batch_variants, actions, states

@@ -553,37 +540,51 @@ 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
+    # Load config (always needed for data setup)
    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 ...')

-    # Apply precision settings before moving to GPU
-    model = apply_precision_settings(model, args)
+    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

-    # 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
+        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 ...')

-    # Build unnomalizer
+        # 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)
    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")
+
    # Run over data
    assert (args.height % 16 == 0) and (
        args.width % 16
@@ -698,7 +699,8 @@ 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)
+                    sim_mode=False,
+                    decode_video=not args.fast_policy_no_decode)

                # Update future actions in the observation queues
                for idx in range(len(pred_actions[0])):
@@ -755,29 +757,31 @@ 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
-                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 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 videos environment changes via world-model interaction
                sample_tag = f"{args.dataset}-vid{sample['videoid']}-wd-fs-{fs}/itr-{itr}"
-                log_to_tensorboard(writer,
-                                   pred_videos_1,
-                                   sample_tag,
-                                   fps=args.save_fps)
+                log_to_tensorboard_async(writer,
+                                         pred_videos_1,
+                                         sample_tag,
+                                         fps=args.save_fps)

                # Save the imagen videos for decision-making
-                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)
+                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)
                # Save videos environment changes via world-model interaction
                sample_video_file = f'{video_save_dir}/wm/{fs}/itr-{itr}.mp4'
-                save_results(pred_videos_1.cpu(),
-                             sample_video_file,
-                             fps=args.save_fps)
+                save_results_async(pred_videos_1,
+                                   sample_video_file,
+                                   fps=args.save_fps)

                print('>' * 24)
                # Collect the result of world-model interactions
@@ -785,12 +789,15 @@ 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(writer,
-                               full_video,
-                               sample_tag,
-                               fps=args.save_fps)
+            log_to_tensorboard_async(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(full_video, sample_full_video_file, fps=args.save_fps)
+            save_results_async(full_video, sample_full_video_file, fps=args.save_fps)
+
+    # Wait for all async I/O to complete
+    _flush_io()


 def get_parser():
@@ -905,39 +912,15 @@ 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/trainer.py
+++ b/scripts/trainer.py
@@ -11,6 +11,9 @@ 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,10 +1105,6 @@ 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)
@@ -2461,6 +2457,7 @@ 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,6 +501,10 @@ 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],
@@ -530,14 +534,20 @@ class ConditionalUnet1D(nn.Module):
        B, T, D = sample.shape
        if self.use_linear_act_proj:
            sample = self.proj_in_action(sample.unsqueeze(-1))
-            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)
+            _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
        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,6 +6,8 @@ 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):
@@ -67,11 +69,12 @@ class DDIMSampler(object):
            ddim_timesteps=self.ddim_timesteps,
            eta=ddim_eta,
            verbose=verbose)
-        self.register_buffer('ddim_sigmas', ddim_sigmas)
-        self.register_buffer('ddim_alphas', ddim_alphas)
-        self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
+        # 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_sqrt_one_minus_alphas',
-                             np.sqrt(1. - ddim_alphas))
+                             to_torch(torch.as_tensor(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))
@@ -241,63 +244,70 @@ class DDIMSampler(object):

        dp_ddim_scheduler_action.set_timesteps(len(timesteps))
        dp_ddim_scheduler_state.set_timesteps(len(timesteps))
-        for i, step in enumerate(iterator):
-            index = total_steps - i - 1
-            ts = torch.full((b, ), step, device=device, dtype=torch.long)
+        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)

-            # 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)
+                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)

        return img, action, state, intermediates

@@ -325,10 +335,6 @@ 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(
@@ -377,17 +383,11 @@ 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

-        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)
+        # 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 self.model.parameterization != "v":
            pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
@@ -395,12 +395,8 @@ 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 = torch.full(size,
-                                 self.ddim_scale_arr[index],
-                                 device=device)
-            prev_scale_t = torch.full(size,
-                                      self.ddim_scale_arr_prev[index],
-                                      device=device)
+            scale_t = self.ddim_scale_arr[index]
+            prev_scale_t = self.ddim_scale_arr_prev[index]
            rescale = (prev_scale_t / scale_t)
            pred_x0 *= rescale

--- a/src/unifolm_wma/modules/attention.py
+++ b/src/unifolm_wma/modules/attention.py
@@ -98,6 +98,10 @@ 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)
@@ -114,6 +118,27 @@ 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
@@ -125,7 +150,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.
@@ -140,19 +165,28 @@ class CrossAttention(nn.Module):
                                    self.agent_action_context_len +
                                    self.text_context_len:, :]

-            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)
+            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)
        else:
            if not spatial_self_attn:
                context = context[:, :self.text_context_len, :]
-            k = self.to_k(context)
-            v = self.to_v(context)
+            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)

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h),
                      (q, k, v))
@@ -236,134 +270,162 @@ 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)

-        if self.image_cross_attention and not spatial_self_attn:
+        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 context.shape[1] == self.text_context_len + self.video_length:
                context_ins, context_image = context[:, :self.text_context_len, :], context[:,self.text_context_len:, :]
-                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)
+                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}
            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:, :]
-                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)
+                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}
            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:, :]

-                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)
+                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)

-                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)
+                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,
+                    }
        else:
            if not spatial_self_attn:
                assert 1 > 2, ">>> ERROR: you should never go into here ..."
                context = context[:, :self.text_context_len, :]
-            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 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}
        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, self.heads, out.shape[1],
+                b, h, out.shape[1],
                self.dim_head).permute(0, 2, 1,
                                       3).reshape(b, out.shape[1],
-                                                  self.heads * self.dim_head))
+                                                  h * 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, self.heads, out_ip.shape[1],
+                b, h, out_ip.shape[1],
                self.dim_head).permute(0, 2, 1,
                                       3).reshape(b, out_ip.shape[1],
-                                                  self.heads * self.dim_head))
+                                                  h * 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, self.heads, out_as.shape[1],
+                b, h, out_as.shape[1],
                self.dim_head).permute(0, 2, 1,
                                       3).reshape(b, out_as.shape[1],
-                                                  self.heads * self.dim_head))
+                                                  h * 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, self.heads, out_aa.shape[1],
+                b, h, out_aa.shape[1],
                self.dim_head).permute(0, 2, 1,
                                       3).reshape(b, out_aa.shape[1],
-                                                  self.heads * self.dim_head))
+                                                  h * self.dim_head))
        if exists(mask):
            raise NotImplementedError

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

        return self.to_out(out)

-    def _get_attn_mask_aa(self, b, l1, l2, block_size=16):
+    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'
        num_token = l2 // block_size
-        start_positions = ((torch.arange(b) % block_size) + 1) * num_token
-        col_indices = torch.arange(l2)
+        start_positions = ((torch.arange(b, device=target_device) % block_size) + 1) * num_token
+        col_indices = torch.arange(l2, device=target_device)
        mask_2d = col_indices.unsqueeze(0) >= start_positions.unsqueeze(1)
        mask = mask_2d.unsqueeze(1).expand(b, l1, l2)
-        attn_mask = torch.zeros_like(mask, dtype=torch.float)
+        attn_mask = torch.zeros(b, l1, l2, dtype=torch.float, device=target_device)
        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,6 +685,21 @@ 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 = {}
+        # 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 forward(self,
                x: Tensor,
                x_action: Tensor,
@@ -720,58 +735,64 @@ class WMAModel(nn.Module):
                                   repeat_only=False).type(x.dtype)
        emb = self.time_embed(t_emb)

-        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
+        _ctx_key = context.data_ptr()
+        if self._ctx_cache_enabled and _ctx_key in self._ctx_cache:
+            context = self._ctx_cache[_ctx_key]
        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')
+            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')

-                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)
+                    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

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

@@ -832,17 +853,45 @@ class WMAModel(nn.Module):

        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)
-            # 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)
+            torch.cuda.current_stream().wait_stream(s_stream)
        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/unitree_g1_pack_camera/case1/output.log
+++ b/unitree_g1_pack_camera/case1/output.log
@@ -1,144 +0,0 @@
-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
-/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: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]
- 45%|████▌     | 5/11 [08:31<10:29, 104.96s/it]
- 55%|█████▍    | 6/11 [10:10<08:35, 103.07s/it]
- 64%|██████▎   | 7/11 [11:48<06:46, 101.50s/it]
- 73%|███████▎  | 8/11 [13:27<05:01, 100.52s/it]
- 82%|████████▏ | 9/11 [15:05<03:19, 99.79s/it] 
- 91%|█████████ | 10/11 [16:43<01:39, 99.30s/it]
-100%|██████████| 11/11 [18:21<00:00, 98.97s/it]
-100%|██████████| 11/11 [18:21<00:00, 100.16s/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/case1/psnr_result.json
+++ b/unitree_g1_pack_camera/case1/psnr_result.json
@@ -1,5 +0,0 @@
-{
-    "gt_video": "unitree_g1_pack_camera/case1/unitree_g1_pack_camera_case1.mp4",
-    "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,144 +0,0 @@
-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.
->>> 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: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]
- 45%|████▌     | 5/11 [08:07<09:45, 97.57s/it]
- 55%|█████▍    | 6/11 [09:45<08:07, 97.59s/it]
- 64%|██████▎   | 7/11 [11:22<06:30, 97.57s/it]
- 73%|███████▎  | 8/11 [13:00<04:52, 97.54s/it]
- 82%|████████▏ | 9/11 [14:37<03:14, 97.50s/it]
- 91%|█████████ | 10/11 [16:14<01:37, 97.32s/it]
-100%|██████████| 11/11 [17:51<00:00, 97.19s/it]
-100%|██████████| 11/11 [17:51<00:00, 97.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_g1_pack_camera/case2/psnr_result.json
+++ b/unitree_g1_pack_camera/case2/psnr_result.json
@@ -1,5 +0,0 @@
-{
-    "gt_video": "unitree_g1_pack_camera/case2/unitree_g1_pack_camera_case2.mp4",
-    "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,144 +0,0 @@
-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.
->>> 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: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 +0,0 @@
-{
-    "gt_video": "unitree_g1_pack_camera/case3/unitree_g1_pack_camera_case3.mp4",
-    "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
@@ -1,144 +0,0 @@
-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 +0,0 @@
-{
-    "gt_video": "unitree_g1_pack_camera/case4/unitree_g1_pack_camera_case4.mp4",
-    "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,15 +1,11 @@
-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)
+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
 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
@@ -18,24 +14,11 @@ 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.
@@ -58,13 +41,7 @@ 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(
+
  0%|          | 0/8 [00:00<?, ?it/s]>>> Step 0: generating actions ...
 >>> Step 0: interacting with world model ...
 DEBUG:PIL.Image:Importing BlpImagePlugin
@@ -115,7 +92,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:14<08:41, 74.51s/it]
 25%|██▌       | 2/8 [02:29<07:28, 74.79s/it]
@@ -139,6 +116,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/unitree_z1_dual_arm_cleanup_pencils_case1_amd.mp4",
-    "psnr": 19.586376345676264
+    "pred_video": "unitree_z1_dual_arm_cleanup_pencils/case1/output/inference/0_full_fs4.mp4",
+    "psnr": 47.911564449209735
 }
--- a/unitree_z1_dual_arm_cleanup_pencils/case1/psnr_result1.json
+++ b/unitree_z1_dual_arm_cleanup_pencils/case1/psnr_result1.json
@@ -1,5 +0,0 @@
-{
-    "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_mix_bf16.ckpt \
+        --ckpt_path ckpts/unifolm_wma_dual.ckpt \
        --config configs/inference/world_model_interaction.yaml \
        --savedir "${res_dir}/output" \
        --bs 1 --height 320 --width 512 \
@@ -20,6 +20,5 @@ dataset="unitree_z1_dual_arm_cleanup_pencils"
        --n_iter 8 \
        --timestep_spacing 'uniform_trailing' \
        --guidance_rescale 0.7 \
-        --perframe_ae \
-        --vae_dtype bf16
+        --perframe_ae
 } 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
@@ -1,137 +0,0 @@
-/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 +0,0 @@
-{
-    "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/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
@@ -1,137 +0,0 @@
-/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 +0,0 @@
-{
-    "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/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
@@ -1,137 +0,0 @@
-/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 +0,0 @@
-{
-    "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/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
@@ -1,134 +0,0 @@
-/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 +0,0 @@
-{
-    "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/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=7 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=0 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
@@ -1,134 +0,0 @@
-/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 +0,0 @@
-{
-    "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/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=6 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=0 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
@@ -1,134 +0,0 @@
-/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 +0,0 @@
-{
-    "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/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
@@ -1,134 +0,0 @@
-/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 +0,0 @@
-{
-    "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/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,34 +1,13 @@
-/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
+2026-02-11 11:59:27.241485: 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-11 11:59:27.291755: 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-11 11:59:27.291807: 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-11 11:59:27.293169: 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-11 11:59:27.300838: 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-11 11:59:28.228009: 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 ...
+>>> 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.
@@ -46,19 +25,16 @@ 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
 >>> 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(
+
  0%|          | 0/11 [00:00<?, ?it/s]
+  9%|▉         | 1/11 [00:34<05:40, 34.05s/it]>>> Step 0: generating actions ...
+>>> Step 0: interacting with world model ...
 >>>>>>>>>>>>>>>>>>>>>>>>
 >>> Step 1: generating actions ...
 DEBUG:PIL.Image:Importing BlpImagePlugin
@@ -108,9 +84,7 @@ DEBUG:PIL.Image:Importing WmfImagePlugin
 DEBUG:PIL.Image:Importing WebPImagePlugin
 DEBUG:PIL.Image:Importing WmfImagePlugin
 DEBUG:PIL.Image:Importing XbmImagePlugin
-DEBUG:PIL.Image:Importing XVThumbImagePlugin
-
-  9%|▉         | 1/11 [01:38<16:20, 98.04s/it]
+DEBUG:PIL.Image:Importing XpmImagePlugin
 DEBUG:PIL.Image:Importing XVThumbImagePlugin

 18%|█▊        | 2/11 [01:08<05:07, 34.17s/it]
@@ -141,6 +115,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/unitree_z1_dual_arm_stackbox_v2_case1_amd.mp4",
-    "psnr": 18.126776535969576
+    "gt_video": "/home/qhy/unifolm-world-model-action/unitree_z1_dual_arm_stackbox_v2/case1/unitree_z1_dual_arm_stackbox_v2_case1.mp4",
+    "pred_video": "/home/qhy/unifolm-world-model-action/unitree_z1_dual_arm_stackbox_v2/case1/output/inference/5_full_fs4.mp4",
+    "psnr": 28.167025381705358
 }
--- 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=7 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=0 python3 scripts/evaluation/world_model_interaction.py \
        --seed 123 \
        --ckpt_path ckpts/unifolm_wma_dual.ckpt \
        --config configs/inference/world_model_interaction.yaml \
@@ -20,5 +20,6 @@ dataset="unitree_z1_dual_arm_stackbox_v2"
        --n_iter 11 \
        --timestep_spacing 'uniform_trailing' \
        --guidance_rescale 0.7 \
-        --perframe_ae
+        --perframe_ae \
+        --fast_policy_no_decode
 } 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
@@ -1,146 +0,0 @@
-/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 +0,0 @@
-{
-    "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/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
@@ -1,146 +0,0 @@
-/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 +0,0 @@
-{
-    "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/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,146 +0,0 @@
-/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
-/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: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]
- 45%|████▌     | 5/11 [08:14<09:53, 98.85s/it]
- 55%|█████▍    | 6/11 [09:53<08:14, 98.87s/it]
- 64%|██████▎   | 7/11 [11:31<06:34, 98.68s/it]
- 73%|███████▎  | 8/11 [13:09<04:55, 98.49s/it]
- 82%|████████▏ | 9/11 [14:47<03:16, 98.38s/it]
- 91%|█████████ | 10/11 [16:25<01:38, 98.29s/it]
-100%|██████████| 11/11 [18:03<00:00, 98.26s/it]
-100%|██████████| 11/11 [18:03<00:00, 98.54s/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/case4/psnr_result.json
+++ b/unitree_z1_dual_arm_stackbox_v2/case4/psnr_result.json
@@ -1,5 +0,0 @@
-{
-    "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/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=6 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=0 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
@@ -1,149 +0,0 @@
-/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 +0,0 @@
-{
-    "gt_video": "unitree_z1_stackbox/case1/unitree_z1_stackbox_case1.mp4",
-    "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=5 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=0 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
@@ -1,149 +0,0 @@
-/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 +0,0 @@
-{
-    "gt_video": "unitree_z1_stackbox/case2/unitree_z1_stackbox_case2.mp4",
-    "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
@@ -1,149 +0,0 @@
-/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 +0,0 @@
-{
-    "gt_video": "unitree_z1_stackbox/case3/unitree_z1_stackbox_case3.mp4",
-    "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
@@ -1,149 +0,0 @@
-/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 +0,0 @@
-{
-    "gt_video": "unitree_z1_stackbox/case4/unitree_z1_stackbox_case4.mp4",
-    "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=7 python3 scripts/evaluation/world_model_interaction.py \
+    time CUDA_VISIBLE_DEVICES=0 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
qhy	9a08e27a19	KV 融合实现完成。改动总结：速度微弱提升psnr略微上升 attention.py — 3处改动： 1. __init__ 添加 _kv_fused = False 标志 2.新增 fuse_kv() 方法：将 to_k + to_v → to_kv，同时处理 _ip/_as/_aa 辅助 KV 对 2. bmm_forward 两个分支加_kv_fused 判断，用to_kv().chunk(2, dim=-1) 替代分别调用	2026-02-11 12:36:38 +08:00
qhy	b558856e1e	fix bugs	2026-02-10 22:35:45 +08:00
qhy	dcbcb2c377	- state_unet 放到一个独立的 CUDA stream 上执行 - action_unet 在默认 stream 上同时执行 - 用 wait_stream 确保两者都完成后再返回两个 1D UNet 输入完全独立，共享的 hs_a 和 context_action 都是只读的。GPU 利用率只有 ~31%，小张量 kernel 不会打满 GPU，两个 stream 可以真正并行。	2026-02-10 21:41:48 +08:00
qhy	ff43432ef9	结果	2026-02-10 20:01:25 +08:00
qhy	afa12ba031	每步迭代保存异步	2026-02-10 19:54:53 +08:00
qhy	bf4d66c874	跳过模型加载	2026-02-10 19:36:17 +08:00
qhy	9347a4ebe5	实现了Context 预计算和缓存功能，提升了采样效率。 psnr不下降	2026-02-10 17:47:46 +08:00
qhy	223a50f9e0	添加CrossAttention kv缓存，减少重复计算，提升性能，psnr=25.1201dB	2026-02-10 17:35:03 +08:00
qhy	2a6068f9e4	减少了一路视频vae解码	2026-02-10 17:13:45 +08:00
qhy	91a9b0febc	DDIM loop 内小张量分配优化，attention mask 缓存到 GPU	2026-02-10 16:53:00 +08:00
qhy	ed637c972b	tf32推理	2026-02-10 16:39:14 +08:00
				`@@ -1 +0,0 @@`
				`{"framework": "pytorch", "task": "robotics", "allow_remote": true}`