From 257fc5da45acfb928a1d26086b610b0464fe117e Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Sun, 12 Jul 2026 16:50:33 +0000 Subject: [PATCH 1/6] update cosmos3 nano --- ...smos3_nano_omni_action_fd_agibotworld.json | 28 +++++++++++++++++ ...omni_action_fd_agibotworld_multichunk.json | 30 +++++++++++++++++++ .../cosmos3_nano_omni_action_id_av.json | 27 +++++++++++++++++ configs/cosmos3/cosmos3_nano_omni_i2av.json | 23 ++++++++++++++ configs/cosmos3/cosmos3_nano_omni_i2v.json | 22 ++++++++++++++ configs/cosmos3/cosmos3_nano_omni_t2av.json | 23 ++++++++++++++ configs/cosmos3/cosmos3_nano_omni_t2v.json | 22 ++++++++++++++ ...cosmos3_nano_omni_action_fd_agibotworld.sh | 23 ++++++++++++++ ...o_omni_action_fd_agibotworld_multichunk.sh | 23 ++++++++++++++ .../cosmos3/cosmos3_nano_omni_action_id_av.sh | 22 ++++++++++++++ scripts/cosmos3/cosmos3_nano_omni_i2av.sh | 24 +++++++++++++++ scripts/cosmos3/cosmos3_nano_omni_i2v.sh | 24 +++++++++++++++ scripts/cosmos3/cosmos3_nano_omni_t2av.sh | 22 ++++++++++++++ scripts/cosmos3/cosmos3_nano_omni_t2v.sh | 22 ++++++++++++++ 14 files changed, 335 insertions(+) create mode 100644 configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.json create mode 100644 configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.json create mode 100644 configs/cosmos3/cosmos3_nano_omni_action_id_av.json create mode 100644 configs/cosmos3/cosmos3_nano_omni_i2av.json create mode 100644 configs/cosmos3/cosmos3_nano_omni_i2v.json create mode 100644 configs/cosmos3/cosmos3_nano_omni_t2av.json create mode 100644 configs/cosmos3/cosmos3_nano_omni_t2v.json create mode 100755 scripts/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.sh create mode 100755 scripts/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.sh create mode 100755 scripts/cosmos3/cosmos3_nano_omni_action_id_av.sh create mode 100755 scripts/cosmos3/cosmos3_nano_omni_i2av.sh create mode 100755 scripts/cosmos3/cosmos3_nano_omni_i2v.sh create mode 100755 scripts/cosmos3/cosmos3_nano_omni_t2av.sh create mode 100755 scripts/cosmos3/cosmos3_nano_omni_t2v.sh diff --git a/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.json b/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.json new file mode 100644 index 000000000..63724a4a3 --- /dev/null +++ b/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.json @@ -0,0 +1,28 @@ +{ + "infer_steps": 30, + "sample_guide_scale": 1.0, + "sample_shift": 10.0, + "target_height": 720, + "target_width": 640, + "target_video_length": 17, + "target_fps": 10.0, + "enable_cfg": true, + "action_mode": "forward_dynamics", + "domain_name": "agibotworld", + "view_point": "concat_view", + "action_chunk_size": 16, + "action_chunk_index": 0, + "feature_caching": "NoCaching", + "rms_norm_type": "one-pass", + "attn_rms_norm_type": "one-pass", + "rope_type": "triton", + "self_attn_type": "flash_attn3", + "causal_self_attn_type": "flash_attn3", + "add_resolution_template": false, + "add_duration_template": false, + "use_system_prompt": false, + "cosmos3_meta_init": true, + "vae_cpu_offload": false, + "cpu_offload": false, + "offload_granularity": "block" +} diff --git a/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.json b/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.json new file mode 100644 index 000000000..b62d93fd5 --- /dev/null +++ b/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.json @@ -0,0 +1,30 @@ +{ + "infer_steps": 30, + "sample_guide_scale": 1.0, + "sample_shift": 10.0, + "target_height": 720, + "target_width": 640, + "target_video_length": 17, + "target_fps": 10.0, + "enable_cfg": true, + "action_mode": "forward_dynamics", + "domain_name": "agibotworld", + "view_point": "concat_view", + "action_chunk_size": 16, + "action_chunk_index": 0, + "action_multichunk": true, + "action_num_chunks": 4, + "feature_caching": "NoCaching", + "rms_norm_type": "one-pass", + "attn_rms_norm_type": "one-pass", + "rope_type": "triton", + "self_attn_type": "flash_attn3", + "causal_self_attn_type": "flash_attn3", + "add_resolution_template": false, + "add_duration_template": false, + "use_system_prompt": false, + "cosmos3_meta_init": true, + "vae_cpu_offload": false, + "cpu_offload": false, + "offload_granularity": "block" +} diff --git a/configs/cosmos3/cosmos3_nano_omni_action_id_av.json b/configs/cosmos3/cosmos3_nano_omni_action_id_av.json new file mode 100644 index 000000000..05d05ecdf --- /dev/null +++ b/configs/cosmos3/cosmos3_nano_omni_action_id_av.json @@ -0,0 +1,27 @@ +{ + "infer_steps": 30, + "sample_guide_scale": 1.0, + "sample_shift": 10.0, + "target_height": 480, + "target_width": 832, + "target_video_length": 61, + "target_fps": 10.0, + "enable_cfg": true, + "action_mode": "inverse_dynamics", + "domain_name": "av", + "view_point": "ego_view", + "action_chunk_size": 60, + "feature_caching": "NoCaching", + "rms_norm_type": "one-pass", + "attn_rms_norm_type": "one-pass", + "rope_type": "triton", + "self_attn_type": "flash_attn3", + "causal_self_attn_type": "flash_attn3", + "add_resolution_template": false, + "add_duration_template": false, + "use_system_prompt": false, + "cosmos3_meta_init": true, + "vae_cpu_offload": false, + "cpu_offload": false, + "offload_granularity": "block" +} diff --git a/configs/cosmos3/cosmos3_nano_omni_i2av.json b/configs/cosmos3/cosmos3_nano_omni_i2av.json new file mode 100644 index 000000000..48602096e --- /dev/null +++ b/configs/cosmos3/cosmos3_nano_omni_i2av.json @@ -0,0 +1,23 @@ +{ + "infer_steps": 35, + "sample_guide_scale": 6.0, + "sample_shift": 10.0, + "target_height": 720, + "target_width": 1280, + "target_video_length": 189, + "target_fps": 24.0, + "enable_cfg": true, + "enable_sound": true, + "feature_caching": "NoCaching", + "rms_norm_type": "one-pass", + "attn_rms_norm_type": "one-pass", + "rope_type": "triton", + "self_attn_type": "flash_attn3", + "causal_self_attn_type": "flash_attn3", + "add_resolution_template": false, + "add_duration_template": false, + "cosmos3_meta_init": true, + "vae_cpu_offload": false, + "cpu_offload": false, + "offload_granularity": "block" +} diff --git a/configs/cosmos3/cosmos3_nano_omni_i2v.json b/configs/cosmos3/cosmos3_nano_omni_i2v.json new file mode 100644 index 000000000..3072500ad --- /dev/null +++ b/configs/cosmos3/cosmos3_nano_omni_i2v.json @@ -0,0 +1,22 @@ +{ + "infer_steps": 35, + "sample_guide_scale": 6.0, + "sample_shift": 10.0, + "target_height": 720, + "target_width": 1280, + "target_video_length": 189, + "target_fps": 24.0, + "enable_cfg": true, + "feature_caching": "NoCaching", + "rms_norm_type": "one-pass", + "attn_rms_norm_type": "one-pass", + "rope_type": "triton", + "self_attn_type": "flash_attn3", + "causal_self_attn_type": "flash_attn3", + "add_resolution_template": false, + "add_duration_template": false, + "cosmos3_meta_init": true, + "vae_cpu_offload": false, + "cpu_offload": false, + "offload_granularity": "block" +} diff --git a/configs/cosmos3/cosmos3_nano_omni_t2av.json b/configs/cosmos3/cosmos3_nano_omni_t2av.json new file mode 100644 index 000000000..48602096e --- /dev/null +++ b/configs/cosmos3/cosmos3_nano_omni_t2av.json @@ -0,0 +1,23 @@ +{ + "infer_steps": 35, + "sample_guide_scale": 6.0, + "sample_shift": 10.0, + "target_height": 720, + "target_width": 1280, + "target_video_length": 189, + "target_fps": 24.0, + "enable_cfg": true, + "enable_sound": true, + "feature_caching": "NoCaching", + "rms_norm_type": "one-pass", + "attn_rms_norm_type": "one-pass", + "rope_type": "triton", + "self_attn_type": "flash_attn3", + "causal_self_attn_type": "flash_attn3", + "add_resolution_template": false, + "add_duration_template": false, + "cosmos3_meta_init": true, + "vae_cpu_offload": false, + "cpu_offload": false, + "offload_granularity": "block" +} diff --git a/configs/cosmos3/cosmos3_nano_omni_t2v.json b/configs/cosmos3/cosmos3_nano_omni_t2v.json new file mode 100644 index 000000000..3072500ad --- /dev/null +++ b/configs/cosmos3/cosmos3_nano_omni_t2v.json @@ -0,0 +1,22 @@ +{ + "infer_steps": 35, + "sample_guide_scale": 6.0, + "sample_shift": 10.0, + "target_height": 720, + "target_width": 1280, + "target_video_length": 189, + "target_fps": 24.0, + "enable_cfg": true, + "feature_caching": "NoCaching", + "rms_norm_type": "one-pass", + "attn_rms_norm_type": "one-pass", + "rope_type": "triton", + "self_attn_type": "flash_attn3", + "causal_self_attn_type": "flash_attn3", + "add_resolution_template": false, + "add_duration_template": false, + "cosmos3_meta_init": true, + "vae_cpu_offload": false, + "cpu_offload": false, + "offload_granularity": "block" +} diff --git a/scripts/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.sh b/scripts/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.sh new file mode 100755 index 000000000..05b302ddd --- /dev/null +++ b/scripts/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.sh @@ -0,0 +1,23 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/root/yongyang/LightX2V +model_path=/root/yongyang/models/nvidia_cosmos3_models/Cosmos3-Nano +image_path=${model_path}/assets/example_action_fd_agibotworld_first_frame.png +action_path=${model_path}/assets/example_action_fd_agibotworld_action_chunks.json + +export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.infer \ +--model_cls cosmos3 \ +--task i2va \ +--model_path ${model_path} \ +--config_json ${lightx2v_path}/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.json \ +--prompt "" \ +--image_path ${image_path} \ +--action_path ${action_path} \ +--save_result_path ${lightx2v_path}/save_results/cosmos3_nano_action_fd_agibotworld.mp4 \ +--seed 0 diff --git a/scripts/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.sh b/scripts/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.sh new file mode 100755 index 000000000..dca2fa09f --- /dev/null +++ b/scripts/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.sh @@ -0,0 +1,23 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/root/yongyang/LightX2V +model_path=/root/yongyang/models/nvidia_cosmos3_models/Cosmos3-Nano +image_path=${model_path}/assets/example_action_fd_agibotworld_first_frame.png +action_path=${model_path}/assets/example_action_fd_agibotworld_action_chunks.json + +export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.infer \ +--model_cls cosmos3 \ +--task i2va \ +--model_path ${model_path} \ +--config_json ${lightx2v_path}/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.json \ +--prompt "" \ +--image_path ${image_path} \ +--action_path ${action_path} \ +--save_result_path ${lightx2v_path}/save_results/cosmos3_nano_action_fd_agibotworld_multichunk.mp4 \ +--seed 0 diff --git a/scripts/cosmos3/cosmos3_nano_omni_action_id_av.sh b/scripts/cosmos3/cosmos3_nano_omni_action_id_av.sh new file mode 100755 index 000000000..cab4f28e3 --- /dev/null +++ b/scripts/cosmos3/cosmos3_nano_omni_action_id_av.sh @@ -0,0 +1,22 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/root/yongyang/LightX2V +model_path=/root/yongyang/models/nvidia_cosmos3_models/Cosmos3-Nano +video_path=${model_path}/assets/example_action_id_av_0_input.mp4 + +export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.infer \ +--model_cls cosmos3 \ +--task v2av \ +--model_path ${model_path} \ +--config_json ${lightx2v_path}/configs/cosmos3/cosmos3_nano_omni_action_id_av.json \ +--prompt "You are an autonomous vehicle planning system." \ +--video_path ${video_path} \ +--save_result_path ${lightx2v_path}/save_results/cosmos3_nano_action_id_av.mp4 \ +--save_action_path ${lightx2v_path}/save_results/cosmos3_nano_action_id_av.json \ +--seed 0 diff --git a/scripts/cosmos3/cosmos3_nano_omni_i2av.sh b/scripts/cosmos3/cosmos3_nano_omni_i2av.sh new file mode 100755 index 000000000..d3ffb844f --- /dev/null +++ b/scripts/cosmos3/cosmos3_nano_omni_i2av.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/root/yongyang/LightX2V +model_path=/root/yongyang/models/nvidia_cosmos3_models/Cosmos3-Nano +prompt_path=${model_path}/assets/example_i2v_prompt.json +negative_prompt_path=${model_path}/assets/negative_prompt.json +image_path=${model_path}/assets/example_i2v_input.jpg + +export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.infer \ +--model_cls cosmos3 \ +--task i2av \ +--model_path ${model_path} \ +--config_json ${lightx2v_path}/configs/cosmos3/cosmos3_nano_omni_i2av.json \ +--prompt ${prompt_path} \ +--negative_prompt ${negative_prompt_path} \ +--image_path ${image_path} \ +--save_result_path ${lightx2v_path}/save_results/cosmos3_nano_omni_i2av.mp4 \ +--seed 17 diff --git a/scripts/cosmos3/cosmos3_nano_omni_i2v.sh b/scripts/cosmos3/cosmos3_nano_omni_i2v.sh new file mode 100755 index 000000000..f46955423 --- /dev/null +++ b/scripts/cosmos3/cosmos3_nano_omni_i2v.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/root/yongyang/LightX2V +model_path=/root/yongyang/models/nvidia_cosmos3_models/Cosmos3-Nano +prompt_path=${model_path}/assets/example_i2v_prompt.json +negative_prompt_path=${model_path}/assets/negative_prompt.json +image_path=${model_path}/assets/example_i2v_input.jpg + +export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.infer \ +--model_cls cosmos3 \ +--task i2v \ +--model_path ${model_path} \ +--config_json ${lightx2v_path}/configs/cosmos3/cosmos3_nano_omni_i2v.json \ +--prompt ${prompt_path} \ +--negative_prompt ${negative_prompt_path} \ +--image_path ${image_path} \ +--save_result_path ${lightx2v_path}/save_results/cosmos3_nano_omni_i2v.mp4 \ +--seed 17 diff --git a/scripts/cosmos3/cosmos3_nano_omni_t2av.sh b/scripts/cosmos3/cosmos3_nano_omni_t2av.sh new file mode 100755 index 000000000..7f60020ef --- /dev/null +++ b/scripts/cosmos3/cosmos3_nano_omni_t2av.sh @@ -0,0 +1,22 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/root/yongyang/LightX2V +model_path=/root/yongyang/models/nvidia_cosmos3_models/Cosmos3-Nano +prompt_path=${model_path}/assets/example_t2vs_prompt.json +negative_prompt_path=${model_path}/assets/negative_prompt.json + +export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.infer \ +--model_cls cosmos3 \ +--task t2av \ +--model_path ${model_path} \ +--config_json ${lightx2v_path}/configs/cosmos3/cosmos3_nano_omni_t2av.json \ +--prompt ${prompt_path} \ +--negative_prompt ${negative_prompt_path} \ +--save_result_path ${lightx2v_path}/save_results/cosmos3_nano_omni_t2av.mp4 \ +--seed 17 diff --git a/scripts/cosmos3/cosmos3_nano_omni_t2v.sh b/scripts/cosmos3/cosmos3_nano_omni_t2v.sh new file mode 100755 index 000000000..5707a0677 --- /dev/null +++ b/scripts/cosmos3/cosmos3_nano_omni_t2v.sh @@ -0,0 +1,22 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/root/yongyang/LightX2V +model_path=/root/yongyang/models/nvidia_cosmos3_models/Cosmos3-Nano +prompt_path=${model_path}/assets/example_t2v_prompt.json +negative_prompt_path=${model_path}/assets/negative_prompt.json + +export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.infer \ +--model_cls cosmos3 \ +--task t2v \ +--model_path ${model_path} \ +--config_json ${lightx2v_path}/configs/cosmos3/cosmos3_nano_omni_t2v.json \ +--prompt ${prompt_path} \ +--negative_prompt ${negative_prompt_path} \ +--save_result_path ${lightx2v_path}/save_results/cosmos3_nano_omni_t2v.mp4 \ +--seed 123 From a8c77f4de2fa77a6ddc5fbb2defadecf06d0d087 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Sun, 12 Jul 2026 19:00:39 +0000 Subject: [PATCH 2/6] support policy --- ...osmos3_nano_omni_t2v_cfg_ulysses_8gpu.json | 28 ++++ ...s3_nano_policy_droid_cfg_ulysses_8gpu.json | 38 +++++ .../networks/cosmos3/infer/pre_infer.py | 24 +++- lightx2v/models/networks/cosmos3/model.py | 1 + .../models/runners/cosmos3/cosmos3_runner.py | 136 ++++++++++++++++-- .../models/schedulers/cosmos3/scheduler.py | 28 +++- scripts/cosmos3/cosmos3_nano_omni_t2v.sh | 4 +- .../cosmos3_nano_omni_t2v_cfg_ulysses_8gpu.sh | 22 +++ ...mos3_nano_policy_droid_cfg_ulysses_8gpu.sh | 24 ++++ 9 files changed, 289 insertions(+), 16 deletions(-) create mode 100644 configs/cosmos3/cosmos3_nano_omni_t2v_cfg_ulysses_8gpu.json create mode 100644 configs/cosmos3/cosmos3_nano_policy_droid_cfg_ulysses_8gpu.json create mode 100755 scripts/cosmos3/cosmos3_nano_omni_t2v_cfg_ulysses_8gpu.sh create mode 100755 scripts/cosmos3/cosmos3_nano_policy_droid_cfg_ulysses_8gpu.sh diff --git a/configs/cosmos3/cosmos3_nano_omni_t2v_cfg_ulysses_8gpu.json b/configs/cosmos3/cosmos3_nano_omni_t2v_cfg_ulysses_8gpu.json new file mode 100644 index 000000000..139afc807 --- /dev/null +++ b/configs/cosmos3/cosmos3_nano_omni_t2v_cfg_ulysses_8gpu.json @@ -0,0 +1,28 @@ +{ + "infer_steps": 35, + "sample_guide_scale": 6.0, + "sample_shift": 10.0, + "target_height": 720, + "target_width": 1280, + "target_video_length": 189, + "target_fps": 24.0, + "enable_cfg": true, + "feature_caching": "NoCaching", + "rms_norm_type": "one-pass", + "attn_rms_norm_type": "one-pass", + "rope_type": "triton", + "self_attn_type": "flash_attn2", + "causal_self_attn_type": "flash_attn2", + "add_resolution_template": false, + "add_duration_template": false, + "cosmos3_meta_init": true, + "unload_modules": false, + "vae_cpu_offload": true, + "cpu_offload": true, + "offload_granularity": "block", + "parallel": { + "seq_p_size": 4, + "seq_p_attn_type": "ulysses", + "cfg_p_size": 2 + } +} diff --git a/configs/cosmos3/cosmos3_nano_policy_droid_cfg_ulysses_8gpu.json b/configs/cosmos3/cosmos3_nano_policy_droid_cfg_ulysses_8gpu.json new file mode 100644 index 000000000..d071c6a10 --- /dev/null +++ b/configs/cosmos3/cosmos3_nano_policy_droid_cfg_ulysses_8gpu.json @@ -0,0 +1,38 @@ +{ + "infer_steps": 4, + "sample_guide_scale": 3.0, + "sample_shift": 5.0, + "target_height": 544, + "target_width": 736, + "target_video_length": 33, + "target_fps": 15.0, + "enable_cfg": true, + "action_mode": "policy", + "domain_name": "droid_lerobot", + "view_point": "concat_view", + "action_chunk_size": 32, + "raw_action_dim": 8, + "policy_use_state": true, + "policy_history_length": 1, + "policy_flip_gripper": true, + "decode_video": false, + "feature_caching": "NoCaching", + "rms_norm_type": "one-pass", + "attn_rms_norm_type": "one-pass", + "rope_type": "triton", + "self_attn_type": "flash_attn2", + "causal_self_attn_type": "flash_attn2", + "add_resolution_template": false, + "add_duration_template": false, + "use_system_prompt": false, + "cosmos3_meta_init": true, + "unload_modules": false, + "vae_cpu_offload": true, + "cpu_offload": true, + "offload_granularity": "block", + "parallel": { + "seq_p_size": 4, + "seq_p_attn_type": "ulysses", + "cfg_p_size": 2 + } +} diff --git a/lightx2v/models/networks/cosmos3/infer/pre_infer.py b/lightx2v/models/networks/cosmos3/infer/pre_infer.py index e2ab5841a..d1d41b74c 100644 --- a/lightx2v/models/networks/cosmos3/infer/pre_infer.py +++ b/lightx2v/models/networks/cosmos3/infer/pre_infer.py @@ -145,7 +145,16 @@ def _prepare_sound_segment(self, sound_latents, text_segment, vision_segment, de "sound_len": sound_len, } - def _prepare_action_segment(self, action_latents, text_segment, vision_segment, sound_segment, device, condition_frame_indexes=None): + def _prepare_action_segment( + self, + action_latents, + text_segment, + vision_segment, + sound_segment, + device, + condition_frame_indexes=None, + start_frame_offset=1, + ): action_len = int(action_latents.shape[0]) condition_frame_indexes = [] if condition_frame_indexes is None else condition_frame_indexes cond_set = {int(idx) for idx in condition_frame_indexes if 0 <= int(idx) < action_len} @@ -162,7 +171,7 @@ def _prepare_action_segment(self, action_latents, text_segment, vision_segment, base_fps=self.base_fps, temporal_compression_factor=1, base_temporal_compression_factor=self.config.get("vae_scale_factor_temporal", 4), - start_frame_offset=1, + start_frame_offset=int(start_frame_offset), ) sequence_indexes = torch.arange(curr, curr + action_len, dtype=torch.long, device=device) return { @@ -186,6 +195,7 @@ def infer( action_latents=None, action_domain_id=None, action_condition_frame_indexes=None, + action_start_frame_offset=1, raw_action_dim=None, ): device = latents.device @@ -196,7 +206,15 @@ def infer( if sound_latents is not None: sound_segment = self._prepare_sound_segment(sound_latents, text_segment, vision_segment, device) if action_latents is not None: - action_segment = self._prepare_action_segment(action_latents, text_segment, vision_segment, sound_segment, device, condition_frame_indexes=action_condition_frame_indexes) + action_segment = self._prepare_action_segment( + action_latents, + text_segment, + vision_segment, + sound_segment, + device, + condition_frame_indexes=action_condition_frame_indexes, + start_frame_offset=action_start_frame_offset, + ) sequence_length = text_segment["und_len"] + vision_segment["num_vision_tokens"] + sound_segment.get("sound_len", 0) + action_segment.get("action_len", 0) packed_text_embedding = weights.embed_tokens.apply(text_segment["input_ids"]) diff --git a/lightx2v/models/networks/cosmos3/model.py b/lightx2v/models/networks/cosmos3/model.py index 3dfa34bcf..887846114 100644 --- a/lightx2v/models/networks/cosmos3/model.py +++ b/lightx2v/models/networks/cosmos3/model.py @@ -133,6 +133,7 @@ def _infer_cond_uncond(self, input_ids): action_latents=getattr(self.scheduler, "action_latents", None), action_domain_id=getattr(self.scheduler, "action_domain_id", None), action_condition_frame_indexes=getattr(self.scheduler, "action_condition_frame_indexes", None), + action_start_frame_offset=getattr(self.scheduler, "action_start_frame_offset", 1), raw_action_dim=getattr(self.scheduler, "raw_action_dim", None), ) if self.config["seq_parallel"]: diff --git a/lightx2v/models/runners/cosmos3/cosmos3_runner.py b/lightx2v/models/runners/cosmos3/cosmos3_runner.py index c6670abe7..d7502940e 100644 --- a/lightx2v/models/runners/cosmos3/cosmos3_runner.py +++ b/lightx2v/models/runners/cosmos3/cosmos3_runner.py @@ -37,6 +37,10 @@ "concat_view": "This video contains concatenated views from multiple camera perspectives.", } +_DROID_CONCAT_VIEW_DESCRIPTION = ( + "The top row is from the wrist-mounted camera. The bottom row contains two horizontally concatenated third-person perspective views of the scene from opposite sides, with the robot visible." +) + _EMBODIMENT_TO_DOMAIN_ID = { "no_action": 0, "av": 1, @@ -202,7 +206,7 @@ def _prepare_action_context(self): self.input_info.target_fps = float(spec["fps"]) @staticmethod - def _build_action_json_prompt(description, view_point, num_frames, fps, height, width): + def _build_action_json_prompt(description, view_point, num_frames, fps, height, width, additional_view_description=None): duration_seconds = num_frames / fps if fps > 0 else 0.0 duration = int(duration_seconds) if duration_seconds >= 0 and np.isfinite(duration_seconds) else 0 action_end = round(duration_seconds) if duration_seconds >= 0 and np.isfinite(duration_seconds) else 0 @@ -212,6 +216,8 @@ def _build_action_json_prompt(description, view_point, num_frames, fps, height, desc = f"{desc}." prompt = {} framing = _ACTION_VIEWPOINT_TEMPLATES.get(view_point) + if framing and additional_view_description: + framing = f"{framing} {additional_view_description}" if framing: prompt["cinematography"] = {"framing": framing} ratio = width / height if height > 0 else 1.0 @@ -271,13 +277,19 @@ def tokenize_prompt(self, prompt, negative_prompt=None): cond_text = prompt uncond_text = negative_prompt if action_mode: + view_point = self._get_action_value("view_point", "ego_view") + domain_name = self._get_action_value("domain_name", None) + additional_view_description = None + if action_mode == "policy" and domain_name == "droid_lerobot" and view_point == "concat_view": + additional_view_description = _DROID_CONCAT_VIEW_DESCRIPTION cond_text = self._build_action_json_prompt( prompt, - view_point=self._get_action_value("view_point", "ego_view"), + view_point=view_point, num_frames=num_frames, fps=fps, height=height, width=width, + additional_view_description=additional_view_description, ) elif not is_image and self.config.get("add_duration_template", True): cond_text = self._append_prompt_template(cond_text, f"The video is {num_frames / fps:.1f} seconds long and is of {fps:.0f} FPS.") @@ -419,6 +431,70 @@ def _load_image_tensor_by_path(self, image_path, height, width): frame = np.asarray(image).astype(np.float32) / 127.5 - 1.0 return torch.from_numpy(frame).permute(2, 0, 1).unsqueeze(0).to(device=AI_DEVICE, dtype=GET_DTYPE()) + def _load_policy_image_tensor_by_path(self, image_path, height, width): + """Load a policy observation with the aspect-preserving bottom/right padding used for DROID.""" + if not image_path or not os.path.isfile(image_path): + raise FileNotFoundError(f"Cosmos3 policy image_path does not exist: {image_path}") + with Image.open(image_path) as image: + image = image.convert("RGB") + source_width, source_height = image.size + scale = min(width / source_width, height / source_height, 1.0) + resized_width = max(1, int(scale * source_width + 0.5)) + resized_height = max(1, int(scale * source_height + 0.5)) + if (resized_width, resized_height) != image.size: + resample = getattr(Image, "Resampling", Image).BICUBIC + image = image.resize((resized_width, resized_height), resample=resample) + frame = np.asarray(image, dtype=np.uint8) + + padding_right = width - resized_width + padding_bottom = height - resized_height + if padding_right < 0 or padding_bottom < 0: + raise ValueError(f"Cosmos3 policy resize produced {resized_height}x{resized_width} for target {height}x{width}.") + if padding_right or padding_bottom: + pad_mode = "edge" if padding_right >= resized_width or padding_bottom >= resized_height else "reflect" + frame = np.pad(frame, ((0, padding_bottom), (0, padding_right), (0, 0)), mode=pad_mode) + frame = frame.astype(np.float32) / 127.5 - 1.0 + return torch.from_numpy(frame).permute(2, 0, 1).unsqueeze(0).to(device=AI_DEVICE, dtype=GET_DTYPE()) + + def _load_policy_state(self, raw_action_dim): + state_path = getattr(self.input_info, "state_path", None) or self.config.get("state_path", "") + if not state_path or not os.path.isfile(state_path): + raise FileNotFoundError(f"Cosmos3 policy state_path does not exist: {state_path}") + + suffix = os.path.splitext(state_path)[1].lower() + if suffix == ".json": + with open(state_path, "r") as f: + payload = json.load(f) + elif suffix in (".npy", ".npz"): + payload = np.load(state_path, allow_pickle=True) + if isinstance(payload, np.lib.npyio.NpzFile): + payload = {key: payload[key] for key in payload.files} + elif isinstance(payload, np.ndarray) and payload.shape == () and isinstance(payload.item(), dict): + payload = payload.item() + else: + payload = np.loadtxt(state_path, delimiter=",", dtype=np.float32) + + if isinstance(payload, dict): + if "state" in payload: + payload = payload["state"] + elif "joint_position" in payload and "gripper_position" in payload: + payload = np.concatenate( + [ + np.asarray(payload["joint_position"], dtype=np.float32).reshape(-1), + np.asarray(payload["gripper_position"], dtype=np.float32).reshape(-1), + ] + ) + else: + raise ValueError("Cosmos3 policy state JSON/NPZ must contain `state` or joint/gripper fields.") + + state = np.asarray(payload, dtype=np.float32).reshape(-1) + if state.size != raw_action_dim: + raise ValueError(f"Cosmos3 policy state must contain {raw_action_dim} floats, got {state.size}.") + if self.config.get("policy_flip_gripper", True): + state = state.copy() + state[-1] = 1.0 - state[-1] + return torch.as_tensor(state, device=AI_DEVICE, dtype=GET_DTYPE()) + def _get_action_chunk_index(self): return int(getattr(self, "_action_chunk_index", self.config.get("action_chunk_index", 0))) @@ -480,7 +556,7 @@ def _prepare_action_condition_latents(self): raise ValueError("Cosmos3 action generation requires domain_name in config or action JSON.") if domain_name not in _EMBODIMENT_TO_DOMAIN_ID or domain_name not in _EMBODIMENT_TO_RAW_ACTION_DIM: raise ValueError(f"Unsupported Cosmos3 action domain_name={domain_name!r}") - raw_action_dim = int(_EMBODIMENT_TO_RAW_ACTION_DIM[domain_name]) + raw_action_dim = int(self._get_action_value("raw_action_dim", _EMBODIMENT_TO_RAW_ACTION_DIM[domain_name])) if raw_action_dim > action_dim: raise ValueError(f"Cosmos3 raw_action_dim={raw_action_dim} exceeds model action_dim={action_dim}") @@ -496,8 +572,18 @@ def _prepare_action_condition_latents(self): if action_mode == "inverse_dynamics": video = self._load_video_tensor(video_path, num_frames, height, width, keep_first=False) elif image_path: - frame = self._load_image_tensor_by_path(image_path, height, width) - video = frame.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous() + if action_mode == "policy": + frame = self._load_policy_image_tensor_by_path(image_path, height, width) + video = torch.full( + (frame.shape[0], frame.shape[1], num_frames, frame.shape[2], frame.shape[3]), + -1.0, + device=frame.device, + dtype=frame.dtype, + ) + video[:, :, 0] = frame + else: + frame = self._load_image_tensor_by_path(image_path, height, width) + video = frame.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous() else: video = self._load_video_tensor(video_path, num_frames, height, width, keep_first=True) @@ -517,9 +603,19 @@ def _prepare_action_condition_latents(self): raw_actions = torch.cat([raw_actions, padding], dim=-1) self.input_info.action_latents = raw_actions self.input_info.action_condition_frame_indexes = list(range(chunk_size)) - elif action_mode in ("inverse_dynamics", "policy"): + elif action_mode == "inverse_dynamics": self.input_info.action_latent_shape = (chunk_size, action_dim) self.input_info.action_condition_frame_indexes = [] + self.input_info.action_start_frame_offset = 1 + elif action_mode == "policy": + if not self.config.get("policy_use_state", True): + raise ValueError("Cosmos3-Nano-Policy-DROID requires policy_use_state=True.") + state = self._load_policy_state(raw_action_dim) + action_latents = torch.zeros((chunk_size + 1, action_dim), device=AI_DEVICE, dtype=GET_DTYPE()) + action_latents[0, :raw_action_dim] = state + self.input_info.action_latents = action_latents + self.input_info.action_condition_frame_indexes = [0] + self.input_info.action_start_frame_offset = 0 else: raise ValueError(f"Unsupported Cosmos3 action_mode={action_mode!r}") @@ -538,6 +634,7 @@ def _clear_action_condition_state(self): "action_condition_frame_indexes", "action_domain_id", "raw_action_dim", + "action_start_frame_offset", ): if hasattr(self.input_info, name): delattr(self.input_info, name) @@ -641,12 +738,21 @@ def _collect_action_output(self, action_latents): if action_mode not in ("inverse_dynamics", "policy"): return None raw_action_dim = getattr(self.input_info, "raw_action_dim", None) - action = action_latents.detach().cpu() + # NumPy does not support torch.bfloat16. Policy actions are also more + # convenient for downstream robot code when persisted as float32. + action = action_latents.detach().float().cpu() if raw_action_dim is not None: action = action[:, : int(raw_action_dim)] + if action_mode == "policy": + history_length = int(self.config.get("policy_history_length", 1)) + action = action[history_length:] + if self.config.get("policy_flip_gripper", True) and action.shape[-1] > 0: + action[:, -1] = 1.0 - action[:, -1] return action def _save_action_output(self, action): + if dist.is_initialized() and dist.get_rank() != 0: + return if action is None: return save_action_path = getattr(self.input_info, "save_action_path", None) @@ -659,8 +765,11 @@ def _save_action_output(self, action): save_dir = os.path.dirname(save_action_path) if save_dir: os.makedirs(save_dir, exist_ok=True) - with open(save_action_path, "w") as f: - json.dump(action.tolist(), f) + if save_action_path.endswith(".npy"): + np.save(save_action_path, action.numpy()) + else: + with open(save_action_path, "w") as f: + json.dump(action.tolist(), f) logger.info(f"Action saved: {save_action_path}") def run(self, total_steps=None): @@ -843,6 +952,15 @@ def run_pipeline(self, input_info): action_latents = getattr(self.model.scheduler, "action_latents", None) if hasattr(self, "model") else None sound = self.run_sound_decoder(sound_latents) action = self._collect_action_output(action_latents) + if self._get_action_mode() == "policy" and not self.config.get("decode_video", False): + self._save_action_output(action) + del latents, generator + self.end_run() + torch_device_module.empty_cache() + gc.collect() + if input_info.return_result_tensor or not getattr(input_info, "save_action_path", None): + return {"action": action} + return {"action": None} images = self.run_vae_decoder(latents) self._save_images(images, input_info, log_prefix="Image saved", sound=sound) self._save_action_output(action) diff --git a/lightx2v/models/schedulers/cosmos3/scheduler.py b/lightx2v/models/schedulers/cosmos3/scheduler.py index edafdd9e2..b7ae7beec 100644 --- a/lightx2v/models/schedulers/cosmos3/scheduler.py +++ b/lightx2v/models/schedulers/cosmos3/scheduler.py @@ -565,12 +565,31 @@ def prepare_latents(self, input_info): self.action_latents = None self.action_domain_id = getattr(input_info, "action_domain_id", None) self.action_condition_frame_indexes = getattr(input_info, "action_condition_frame_indexes", None) + self.action_start_frame_offset = getattr(input_info, "action_start_frame_offset", 1) self.raw_action_dim = getattr(input_info, "raw_action_dim", None) + self.action_condition_mask = None + self.action_condition_reference = None action_latents = getattr(input_info, "action_latents", None) if action_latents is not None: - self.action_latents = action_latents.to(device=AI_DEVICE, dtype=GET_DTYPE()) + action_reference = action_latents.to(device=AI_DEVICE, dtype=GET_DTYPE()) if self.raw_action_dim is not None: - self.action_latents[:, int(self.raw_action_dim) :] = 0 + action_reference[:, int(self.raw_action_dim) :] = 0 + condition_indexes = self.action_condition_frame_indexes + if condition_indexes is None: + self.action_latents = action_reference + else: + action_len = int(action_reference.shape[0]) + condition_indexes = [int(idx) for idx in condition_indexes if 0 <= int(idx) < action_len] + mask = torch.zeros((action_len, 1), device=AI_DEVICE, dtype=GET_DTYPE()) + if condition_indexes: + mask[condition_indexes] = 1.0 + noise = torch.randn(tuple(action_reference.shape), generator=self.generator, device=AI_DEVICE, dtype=GET_DTYPE()) + if self.raw_action_dim is not None: + noise[:, int(self.raw_action_dim) :] = 0 + self.action_latents = mask * action_reference + (1.0 - mask) * noise + self.action_condition_frame_indexes = condition_indexes + self.action_condition_mask = mask + self.action_condition_reference = action_reference else: action_shape = getattr(input_info, "action_latent_shape", None) if action_shape is not None: @@ -623,6 +642,8 @@ def step_post(self): self.action_latents.unsqueeze(0), return_dict=False, )[0].squeeze(0) + if self.action_condition_mask is not None and self.action_condition_reference is not None: + self.action_latents = self.action_condition_mask * self.action_condition_reference + (1.0 - self.action_condition_mask) * self.action_latents if self.raw_action_dim is not None: self.action_latents[:, int(self.raw_action_dim) :] = 0 self.noise_pred = None @@ -645,4 +666,7 @@ def clear(self): self.vision_condition_mask = None self.action_domain_id = None self.action_condition_frame_indexes = None + self.action_start_frame_offset = 1 + self.action_condition_mask = None + self.action_condition_reference = None self.raw_action_dim = None diff --git a/scripts/cosmos3/cosmos3_nano_omni_t2v.sh b/scripts/cosmos3/cosmos3_nano_omni_t2v.sh index 5707a0677..d48b454ea 100755 --- a/scripts/cosmos3/cosmos3_nano_omni_t2v.sh +++ b/scripts/cosmos3/cosmos3_nano_omni_t2v.sh @@ -1,8 +1,8 @@ #!/bin/bash # set path firstly -lightx2v_path=/root/yongyang/LightX2V -model_path=/root/yongyang/models/nvidia_cosmos3_models/Cosmos3-Nano +lightx2v_path=/app/LightX2V +model_path=/app/nvidia_cosmos3_models/Cosmos3-Nano prompt_path=${model_path}/assets/example_t2v_prompt.json negative_prompt_path=${model_path}/assets/negative_prompt.json diff --git a/scripts/cosmos3/cosmos3_nano_omni_t2v_cfg_ulysses_8gpu.sh b/scripts/cosmos3/cosmos3_nano_omni_t2v_cfg_ulysses_8gpu.sh new file mode 100755 index 000000000..268f7ebe2 --- /dev/null +++ b/scripts/cosmos3/cosmos3_nano_omni_t2v_cfg_ulysses_8gpu.sh @@ -0,0 +1,22 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/app/LightX2V +model_path=/app/nvidia_cosmos3_models/Cosmos3-Nano +prompt_path=${model_path}/assets/example_t2v_prompt.json +negative_prompt_path=${model_path}/assets/negative_prompt.json + +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +torchrun --nproc_per_node=8 -m lightx2v.infer \ +--model_cls cosmos3 \ +--task t2v \ +--model_path ${model_path} \ +--config_json ${lightx2v_path}/configs/cosmos3/cosmos3_nano_omni_t2v_cfg_ulysses_8gpu.json \ +--prompt ${prompt_path} \ +--negative_prompt ${negative_prompt_path} \ +--save_result_path ${lightx2v_path}/save_results/cosmos3_nano_omni_t2v_cfg_ulysses_8gpu.mp4 \ +--seed 123 diff --git a/scripts/cosmos3/cosmos3_nano_policy_droid_cfg_ulysses_8gpu.sh b/scripts/cosmos3/cosmos3_nano_policy_droid_cfg_ulysses_8gpu.sh new file mode 100755 index 000000000..8dc1b174d --- /dev/null +++ b/scripts/cosmos3/cosmos3_nano_policy_droid_cfg_ulysses_8gpu.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/app/LightX2V +model_path=/app/nvidia_cosmos3_models/Cosmos3-Nano-Policy-DROID +input_path=/app/lightx2v_examples/i2va/robolab/banana_in_bowl +image_path=${input_path}/observation.png +state_path=${input_path}/state.npy + +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +torchrun --nproc_per_node=8 -m lightx2v.infer \ +--model_cls cosmos3 \ +--task i2va \ +--model_path ${model_path} \ +--config_json ${lightx2v_path}/configs/cosmos3/cosmos3_nano_policy_droid_cfg_ulysses_8gpu.json \ +--prompt "Pick up the banana and place it in the bowl" \ +--image_path ${image_path} \ +--state_path ${state_path} \ +--save_action_path ${lightx2v_path}/save_results/cosmos3_nano_policy_droid_action.npy \ +--seed 0 From 21a7e3541aac7368d0ecb366221fff7787784f44 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 13 Jul 2026 13:47:33 +0000 Subject: [PATCH 3/6] fix libero --- configs/fastwam/libero_i2va.json | 4 +- configs/fastwam/robotwin_i2va.json | 1 - lightx2v/utils/set_config.py | 2 +- lightx2v_ros/scripts/run_robotwin_service.sh | 236 ------------------- scripts/fastwam/run_libero_i2va.sh | 10 +- 5 files changed, 8 insertions(+), 245 deletions(-) delete mode 100755 lightx2v_ros/scripts/run_robotwin_service.sh diff --git a/configs/fastwam/libero_i2va.json b/configs/fastwam/libero_i2va.json index 2c17b506f..c493bb562 100644 --- a/configs/fastwam/libero_i2va.json +++ b/configs/fastwam/libero_i2va.json @@ -1,6 +1,6 @@ { - "adapter_model_path": "/data/nvme1/yongyang/nb/FastWAM/checkpoints/fastwam_release/libero_uncond_2cam224.pt", - "dataset_stats_path": "/data/nvme1/yongyang/nb/FastWAM/checkpoints/fastwam_release/libero_uncond_2cam224_dataset_stats.json", + "adapter_model_path": "/data/nvme7/yongyang/fastwam_release/libero_uncond_2cam224.pt", + "dataset_stats_path": "/data/nvme7/yongyang/fastwam_release/libero_uncond_2cam224_dataset_stats.json", "camera_size": 224, "action_chunk_size": 32, "actions_per_plan": 10, diff --git a/configs/fastwam/robotwin_i2va.json b/configs/fastwam/robotwin_i2va.json index 17bea6a97..b80c77a5c 100644 --- a/configs/fastwam/robotwin_i2va.json +++ b/configs/fastwam/robotwin_i2va.json @@ -12,7 +12,6 @@ "robot_state_dim": 14, "policy_profile": "robotwin", "normalize_mode": "z-score", - "target_video_length": 9, "binarize_gripper": false, "gripper_postprocess": false, "default_prompt": "A video recorded from a robot's point of view executing the following instruction: {task_prompt}" diff --git a/lightx2v/utils/set_config.py b/lightx2v/utils/set_config.py index 97155fa10..603f0ca1a 100755 --- a/lightx2v/utils/set_config.py +++ b/lightx2v/utils/set_config.py @@ -228,7 +228,7 @@ def auto_calc_config(config): if "infer_steps" not in config and "num_inference_steps" in config: config["infer_steps"] = config["num_inference_steps"] - if config["task"] in ["i2v", "t2av", "i2av", "i2va", "s2v", "rs2v", "ltx2_s2v", "v2av"]: + if config["task"] in ["i2v", "t2av", "i2av", "i2va", "s2v", "rs2v", "ltx2_s2v", "v2av"] and "target_video_length" in config and "vae_stride" in config: if config["target_video_length"] % config["vae_stride"][0] != 1: logger.warning(f"`num_frames - 1` has to be divisible by {config['vae_stride'][0]}. Rounding to the nearest number.") config["target_video_length"] = config["target_video_length"] // config["vae_stride"][0] * config["vae_stride"][0] + 1 diff --git a/lightx2v_ros/scripts/run_robotwin_service.sh b/lightx2v_ros/scripts/run_robotwin_service.sh deleted file mode 100755 index 8e8090bac..000000000 --- a/lightx2v_ros/scripts/run_robotwin_service.sh +++ /dev/null @@ -1,236 +0,0 @@ -#!/usr/bin/env bash -# ============================================================================= -# One-click launcher for the RoboTwin continuous-eval ROS service. -# -# Brings up the three nodes that make up the demo: -# 1. simulator/robotwin_node - SAPIEN dual-arm sim -# 2. inference/fastwam_node - FastWAM policy inference -# 3. visualization/image_web_viewer - web dashboard (default 0.0.0.0:6061) -# -# By default the simulator starts in "ready" state and runs a SINGLE episode -# when you press 开始 in the web dashboard (or publish a `start` control -# command). Start/pause/restart and task/scenario switching are all available -# from the dashboard. Set LOOP=true for the old continuous-eval behaviour and -# AUTOSTART=true to begin evaluating immediately on launch. -# -# Usage: -# run_robotwin_service.sh start # launch all three nodes (default) -# run_robotwin_service.sh stop # stop all nodes -# run_robotwin_service.sh restart # stop then start -# run_robotwin_service.sh status # show pid / running state + log tails -# run_robotwin_service.sh logs [sim|fastwam|viewer] # follow a node log -# -# Everything below is overridable via environment variables, e.g. -# VIEWER_PORT=8080 TASK_NAME=click_alarmclock ./run_robotwin_service.sh start -# ============================================================================= -set -uo pipefail - -# ------------------------------- configuration -------------------------------- -LIGHTX2V="${LIGHTX2V:-/home/fuhaiwen/LightX2V}" -ROS_WS="${ROS_WS:-$LIGHTX2V/lightx2v_ros}" -ROS_UNDERLAY="${ROS_UNDERLAY:-/root/ros2_lyrical/install/setup.bash}" -ROS_OVERLAY="${ROS_OVERLAY:-$ROS_WS/install/setup.bash}" -# RoboTwin uses relative asset paths (e.g. ./assets/...), so the simulator node must -# run with its working directory set to the vendored RoboTwin root. -ROBOTWIN_ROOT="${ROBOTWIN_ROOT:-$ROS_WS/src/simulator/simulator/robotwin_node/RoboTwin}" - -# Task / embodiment (simulator node params) -TASK_NAME="${TASK_NAME:-lift_pot}" -TASK_CONFIG="${TASK_CONFIG:-demo_clean}" -EMBODIMENT="${EMBODIMENT:-aloha-agilex}" -SEED="${SEED:-0}" -# <=0 means "use RoboTwin's per-task step limit"; hitting the cap = FAILURE. -MAX_EPISODE_STEPS="${MAX_EPISODE_STEPS:-0}" -# Single-episode by default; control via the web dashboard. -LOOP="${LOOP:-false}" -AUTOSTART="${AUTOSTART:-false}" -# Skip seeds the scripted expert cannot solve + generate real instructions -# (requires curobo; falls back gracefully when unavailable). -EXPERT_CHECK="${EXPERT_CHECK:-true}" -# Stream an intermediate viewer frame every N physics steps during an action -# (higher FPS in the dashboard at some simulation-speed cost; 0 disables). -RENDER_PUBLISH_EVERY="${RENDER_PUBLISH_EVERY:-15}" - -# Policy (fastwam node params) -CONFIG_JSON="${CONFIG_JSON:-$LIGHTX2V/configs/fastwam/robotwin_i2va.json}" -MODEL_PATH="${MODEL_PATH:-/data/nvme7/yongyang/models/Wan-AI/Wan2.2-TI2V-5B}" - -# Web viewer -VIEWER_HOST="${VIEWER_HOST:-0.0.0.0}" -VIEWER_PORT="${VIEWER_PORT:-6061}" - -# GPU assignment (sim is light, policy is heavy) -SIM_GPU="${SIM_GPU:-5}" -POLICY_GPU="${POLICY_GPU:-4}" - -# Runtime dirs -RUN_DIR="${RUN_DIR:-/tmp/robotwin_service}" -LOG_SIM="$RUN_DIR/sim.log" -LOG_FASTWAM="$RUN_DIR/fastwam.log" -LOG_VIEWER="$RUN_DIR/viewer.log" -PID_SIM="$RUN_DIR/sim.pid" -PID_FASTWAM="$RUN_DIR/fastwam.pid" -PID_VIEWER="$RUN_DIR/viewer.pid" - -# ------------------------------- helpers -------------------------------------- -c_info() { printf '\033[1;34m[service]\033[0m %s\n' "$*"; } -c_ok() { printf '\033[1;32m[service]\033[0m %s\n' "$*"; } -c_warn() { printf '\033[1;33m[service]\033[0m %s\n' "$*"; } -c_err() { printf '\033[1;31m[service]\033[0m %s\n' "$*" >&2; } - -source_ros() { - # ROS setup scripts reference unbound vars; disable -u while sourcing them. - set +u - # shellcheck disable=SC1090 - source "$ROS_UNDERLAY" - # shellcheck disable=SC1090 - source "$ROS_OVERLAY" - set -u - export PYTHONPATH="$LIGHTX2V:${PYTHONPATH:-}" -} - -is_running() { # $1 = pidfile - local pf="$1" pid - [[ -f "$pf" ]] || return 1 - pid="$(cat "$pf" 2>/dev/null)" - [[ -n "$pid" ]] || return 1 - kill -0 "$pid" 2>/dev/null -} - -# Launch a command in its own process group so we can kill the whole subtree. -# $1=name $2=logfile $3=pidfile $4=command string (run through bash -c) -launch() { - local name="$1" log="$2" pidfile="$3" cmd="$4" - if is_running "$pidfile"; then - c_warn "$name already running (pid $(cat "$pidfile")); skipping" - return 0 - fi - c_info "starting $name -> $log" - setsid bash -c "$cmd" >"$log" 2>&1 & - echo "$!" >"$pidfile" - c_ok "$name started (pid $(cat "$pidfile"))" -} - -# Wait until $2 (regex) shows up in log $1, or timeout $3 seconds. -wait_for_log() { - local log="$1" pattern="$2" timeout="${3:-180}" waited=0 - c_info "waiting for '$pattern' in $(basename "$log") (<=${timeout}s)..." - while (( waited < timeout )); do - if grep -Eq "$pattern" "$log" 2>/dev/null; then - c_ok "ready: $(basename "$log")" - return 0 - fi - sleep 2; waited=$((waited + 2)) - done - c_warn "timed out waiting for '$pattern' in $(basename "$log") (continuing anyway)" - return 1 -} - -stop_one() { # $1=name $2=pidfile - local name="$1" pf="$2" pid - if ! is_running "$pf"; then - c_info "$name not running" - rm -f "$pf" - return 0 - fi - pid="$(cat "$pf")" - c_info "stopping $name (pgid $pid)" - kill -TERM -- "-$pid" 2>/dev/null || kill -TERM "$pid" 2>/dev/null - for _ in $(seq 1 10); do - kill -0 "$pid" 2>/dev/null || break - sleep 1 - done - if kill -0 "$pid" 2>/dev/null; then - c_warn "$name did not exit; sending SIGKILL" - kill -KILL -- "-$pid" 2>/dev/null || kill -KILL "$pid" 2>/dev/null - fi - rm -f "$pf" - c_ok "$name stopped" -} - -# ------------------------------- commands ------------------------------------- -cmd_start() { - mkdir -p "$RUN_DIR" - [[ -f "$ROS_UNDERLAY" ]] || { c_err "ROS underlay not found: $ROS_UNDERLAY"; exit 1; } - [[ -f "$ROS_OVERLAY" ]] || { c_err "ROS overlay not found: $ROS_OVERLAY"; exit 1; } - [[ -d "$MODEL_PATH" ]] || c_warn "MODEL_PATH does not exist: $MODEL_PATH" - [[ -f "$CONFIG_JSON" ]] || c_warn "CONFIG_JSON does not exist: $CONFIG_JSON" - - c_info "config: task=$TASK_NAME/$TASK_CONFIG embodiment=$EMBODIMENT seed=$SEED" - c_info "gpus: sim=$SIM_GPU policy=$POLICY_GPU | viewer=http://$VIEWER_HOST:$VIEWER_PORT" - - # NOTE: PYTHONPATH must use double quotes so ${PYTHONPATH:-} expands *inside* the - # launched shell (after sourcing ROS), otherwise it clobbers the ros2 paths. - local ros_prelude - ros_prelude="set +u; source '$ROS_UNDERLAY'; source '$ROS_OVERLAY'; set -u; export PYTHONPATH=\"$LIGHTX2V:\${PYTHONPATH:-}\";" - - # 1) simulator (continuous loop). cd into RoboTwin root for its relative asset paths. - launch "simulator" "$LOG_SIM" "$PID_SIM" \ - "$ros_prelude cd '$ROBOTWIN_ROOT'; CUDA_VISIBLE_DEVICES=$SIM_GPU exec ros2 run simulator robotwin_node --ros-args \ - -p task_name:=$TASK_NAME -p task_config:=$TASK_CONFIG -p embodiment:=$EMBODIMENT \ - -p seed:=$SEED -p loop:=$LOOP -p autostart:=$AUTOSTART -p max_episode_steps:=$MAX_EPISODE_STEPS \ - -p expert_check:=$EXPERT_CHECK -p render_publish_every:=$RENDER_PUBLISH_EVERY" - wait_for_log "$LOG_SIM" "control on" 300 || true - - # 2) web viewer - launch "viewer" "$LOG_VIEWER" "$PID_VIEWER" \ - "$ros_prelude exec ros2 run visualization image_web_viewer --ros-args \ - -p env:=robotwin -p host:=$VIEWER_HOST -p port:=$VIEWER_PORT" - wait_for_log "$LOG_VIEWER" "image web viewer on" 60 || true - - # 3) policy (heavy model load; can take a couple of minutes) - launch "fastwam" "$LOG_FASTWAM" "$PID_FASTWAM" \ - "$ros_prelude CUDA_VISIBLE_DEVICES=$POLICY_GPU exec ros2 run inference fastwam_node --ros-args \ - -p env:=robotwin -p config_json:=$CONFIG_JSON -p model_path:=$MODEL_PATH" - wait_for_log "$LOG_FASTWAM" "fastwam_node ready" 600 || true - - echo - c_ok "RoboTwin service is up." - c_ok "Dashboard: http://$VIEWER_HOST:$VIEWER_PORT (press 开始 to run an episode)" - c_info "Follow logs: $0 logs [sim|fastwam|viewer]" - c_info "Stop: $0 stop" - cmd_status -} - -cmd_stop() { - # Stop policy first so it stops feeding actions, then sim, then viewer. - stop_one "fastwam" "$PID_FASTWAM" - stop_one "simulator" "$PID_SIM" - stop_one "viewer" "$PID_VIEWER" -} - -cmd_status() { - echo - printf '%-12s %-10s %s\n' "NODE" "STATE" "PID / LOG" - for entry in "simulator:$PID_SIM:$LOG_SIM" "fastwam:$PID_FASTWAM:$LOG_FASTWAM" "viewer:$PID_VIEWER:$LOG_VIEWER"; do - IFS=':' read -r name pf log <<<"$entry" - if is_running "$pf"; then - printf '%-12s \033[1;32m%-10s\033[0m %s\n' "$name" "running" "$(cat "$pf") ($log)" - else - printf '%-12s \033[1;31m%-10s\033[0m %s\n' "$name" "stopped" "- ($log)" - fi - done - echo -} - -cmd_logs() { - local which="${1:-sim}" log - case "$which" in - sim|simulator) log="$LOG_SIM" ;; - fastwam|policy) log="$LOG_FASTWAM" ;; - viewer|web) log="$LOG_VIEWER" ;; - *) c_err "unknown log '$which' (use: sim|fastwam|viewer)"; exit 1 ;; - esac - [[ -f "$log" ]] || { c_err "log not found: $log"; exit 1; } - exec tail -n 100 -f "$log" -} - -# ------------------------------- dispatch ------------------------------------- -case "${1:-start}" in - start) cmd_start ;; - stop) cmd_stop ;; - restart) cmd_stop; sleep 2; cmd_start ;; - status) cmd_status ;; - logs) shift; cmd_logs "${1:-sim}" ;; - *) c_err "usage: $0 {start|stop|restart|status|logs [sim|fastwam|viewer]}"; exit 1 ;; -esac diff --git a/scripts/fastwam/run_libero_i2va.sh b/scripts/fastwam/run_libero_i2va.sh index fd90e7632..599b9999e 100644 --- a/scripts/fastwam/run_libero_i2va.sh +++ b/scripts/fastwam/run_libero_i2va.sh @@ -1,14 +1,14 @@ #!/bin/bash -lightx2v_path=/data/nvme1/yongyang/nb/LightX2V +lightx2v_path=/data/nvme7/yongyang/LightX2V config_json=${lightx2v_path}/configs/fastwam/libero_i2va.json -model_path=/data/nvme0/models/Wan-AI/Wan2.2-TI2V-5B +model_path=/data/nvme7/yongyang/models/Wan-AI/Wan2.2-TI2V-5B -image_path=/data/nvme1/yongyang/nb/lightx2v_examples/i2va/libero_spatial/task0_init0 -state_path=/data/nvme1/yongyang/nb/lightx2v_examples/i2va/libero_spatial/task0_init0/state.npy +image_path=/data/nvme7/yongyang/lightx2v_examples/i2va/libero_spatial/task0_init0 +state_path=/data/nvme7/yongyang/lightx2v_examples/i2va/libero_spatial/task0_init0/state.npy prompt="pick up the black bowl between the plate and the ramekin and place it on the plate" -export CUDA_VISIBLE_DEVICES=0 +export CUDA_VISIBLE_DEVICES=6 source "${lightx2v_path}/scripts/base/base.sh" From 1ccee3c1b966d01e2489fe1871f7420360dd0596 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Tue, 14 Jul 2026 04:38:04 +0000 Subject: [PATCH 4/6] fix curobo --- .../simulator/simulator/robotwin_node/env.py | 121 +++++++++++++++++- 1 file changed, 117 insertions(+), 4 deletions(-) diff --git a/lightx2v_ros/src/simulator/simulator/robotwin_node/env.py b/lightx2v_ros/src/simulator/simulator/robotwin_node/env.py index f05638b32..3a6d83042 100644 --- a/lightx2v_ros/src/simulator/simulator/robotwin_node/env.py +++ b/lightx2v_ros/src/simulator/simulator/robotwin_node/env.py @@ -20,7 +20,9 @@ on machines where the RoboTwin runtime is not installed yet. """ +import contextlib import importlib +import io import os import sys from pathlib import Path @@ -37,6 +39,14 @@ def default_robotwin_root() -> Path: return Path(__file__).resolve().parent / "RoboTwin" +def resolve_robotwin_root(path=None) -> Path: + """Return a normalized RoboTwin root supplied by ROS or the default.""" + raw_path = str(path).strip() if path is not None else "" + if not raw_path: + return default_robotwin_root().resolve() + return Path(os.path.expandvars(raw_path)).expanduser().resolve() + + def _add_python_path(path) -> None: path = str(Path(path)) if path not in sys.path: @@ -63,7 +73,7 @@ def __init__( logger=None, ): super().__init__(contract) - self.robotwin_root = Path(robotwin_root or default_robotwin_root()).expanduser() + self.robotwin_root = resolve_robotwin_root(robotwin_root) self.task_name = str(task_name) self.task_config = str(task_config) self.embodiment = str(embodiment).strip() @@ -98,12 +108,32 @@ def _log(self, message: str) -> None: # ------------------------------------------------------------------ setup def _prepare_runtime(self) -> None: root = self.robotwin_root - if not (root / "envs").is_dir(): - raise FileNotFoundError(f"RoboTwin is not vendored at {root}. See robotwin_node/RoboTwin/README and run the RoboTwin install/asset-download steps.") + required_paths = { + "envs/": (root / "envs").is_dir(), + "task_config/": (root / "task_config").is_dir(), + "assets/": (root / "assets").is_dir(), + "assets/objects/objaverse/list.json": ( + root / "assets" / "objects" / "objaverse" / "list.json" + ).is_file(), + "assets/objects/same.json": (root / "assets" / "objects" / "same.json").is_file(), + } + missing = [name for name, exists in required_paths.items() if not exists] + if missing: + raise FileNotFoundError( + f"Invalid or incomplete RoboTwin root '{root}': missing {', '.join(missing)}. " + "Pass the directory containing envs/, task_config/, and assets/ " + "with '--ros-args -p robotwin_root:=/path/to/RoboTwin'." + ) + # RoboTwin source uses root-relative imports such as `from envs import ...` - # and `from generate_episode_instructions import *`. + # and `from generate_episode_instructions import *`. It also resolves many + # resources from the process working directory (for example + # `./assets/objects/objaverse/list.json`), so the dedicated simulator node + # must run from the selected RoboTwin root for its entire lifetime. + os.chdir(root) _add_python_path(root) _add_python_path(root / "description" / "utils") + self._log(f"using RoboTwin root: {root}") def _require_config(self, *parts) -> Path: path = self._configs_path.joinpath(*parts) @@ -111,6 +141,88 @@ def _require_config(self, *parts) -> Path: raise FileNotFoundError(f"Missing RoboTwin config: {path}. Populate `task_config/` (and `assets/`) from the official RoboTwin repo (see robotwin_node/RoboTwin/script).") return path + def _prepare_planner_runtime(self) -> None: + """Make RoboTwin importable when its optional Curobo planner is absent. + + RoboTwin catches the Curobo import error in ``planner.py``, but then + ``robot.py`` unconditionally imports ``CuroboPlanner`` from that module. + The ROS adapter only executes qpos actions, so it can safely use a small + compatibility planner for scene initialization and gripper interpolation. + The scripted expert remains disabled through ``CUROBO_AVAILABLE=False``. + """ + planner_module = sys.modules.get("envs.robot.planner") + if planner_module is not None and hasattr(planner_module, "CuroboPlanner"): + if not hasattr(planner_module, "CUROBO_AVAILABLE"): + planner_module.CUROBO_AVAILABLE = True + return + + # Importing this submodule first executes envs.robot.__init__, whose + # unconditional CuroboPlanner import is precisely the upstream bug. Keep + # its expected warning/traceback out of the ROS log and inspect the + # successfully loaded planner submodule after that import fails. + captured_stdout = io.StringIO() + captured_stderr = io.StringIO() + import_error = None + try: + with contextlib.redirect_stdout(captured_stdout), contextlib.redirect_stderr(captured_stderr): + importlib.import_module("envs.robot.planner") + except ImportError as exc: + import_error = exc + + planner_module = sys.modules.get("envs.robot.planner") + if planner_module is None: + captured = captured_stdout.getvalue() + captured_stderr.getvalue() + if captured: + print(captured, file=sys.stderr, end="") + if import_error is not None: + raise import_error + raise ImportError("RoboTwin did not load envs.robot.planner") + + if hasattr(planner_module, "CuroboPlanner"): + planner_module.CUROBO_AVAILABLE = True + return + + class CuroboPlanner: + """No-Curobo compatibility planner for qpos-only simulation.""" + + def __init__(self, *args, **kwargs): + pass + + @staticmethod + def plan_grippers(now_val, target_val): + num_step = 200 + return { + "num_step": num_step, + "per_step": (target_val - now_val) / num_step, + "result": np.linspace(now_val, target_val, num_step), + } + + @staticmethod + def plan_path(*args, **kwargs): + return {"status": "Fail"} + + @staticmethod + def plan_batch(*args, **kwargs): + return {"status": np.array(["Failure"], dtype=object)} + + @staticmethod + def update_point_cloud(*args, **kwargs): + return None + + CuroboPlanner.__module__ = planner_module.__name__ + planner_module.CuroboPlanner = CuroboPlanner + planner_module.CUROBO_AVAILABLE = False + + # Python removes the failed parent imports automatically. Clear any + # remaining partial modules so the next task import sees the patched + # planner and constructs envs.robot normally. + sys.modules.pop("envs.robot.robot", None) + sys.modules.pop("envs.robot", None) + envs_module = sys.modules.get("envs") + if envs_module is not None and hasattr(envs_module, "robot"): + delattr(envs_module, "robot") + self._log("curobo is unavailable: using qpos-only planner compatibility mode") + def _build_task_args(self) -> dict: """Replicates third_party/RoboTwin/script/eval_policy.py:main() arg assembly.""" import yaml @@ -168,6 +280,7 @@ def embodiment_config(robot_file): return args def _instantiate_task(self): + self._prepare_planner_runtime() module = importlib.import_module(f"envs.{self.task_name}") task_cls = getattr(module, self.task_name) task = task_cls() From ebff0265329ff6a7cfb1f5fcc840741577b3aa22 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Tue, 14 Jul 2026 05:22:48 +0000 Subject: [PATCH 5/6] update ros --- .../simulator/simulator/libero_node/env.py | 1 + .../simulator/libero_node/observer.py | 11 ++++++-- .../src/simulator/simulator/sim/node.py | 27 ++++++++++++++++--- 3 files changed, 33 insertions(+), 6 deletions(-) diff --git a/lightx2v_ros/src/simulator/simulator/libero_node/env.py b/lightx2v_ros/src/simulator/simulator/libero_node/env.py index 49f7a384c..d1aa67ddf 100644 --- a/lightx2v_ros/src/simulator/simulator/libero_node/env.py +++ b/lightx2v_ros/src/simulator/simulator/libero_node/env.py @@ -53,6 +53,7 @@ def task_description(self) -> str: return self.observer.task_description def reset(self) -> Observation: + self.observer.reset() return self._observation() def step(self, action): diff --git a/lightx2v_ros/src/simulator/simulator/libero_node/observer.py b/lightx2v_ros/src/simulator/simulator/libero_node/observer.py index 585ef3d50..b0f6f60c0 100644 --- a/lightx2v_ros/src/simulator/simulator/libero_node/observer.py +++ b/lightx2v_ros/src/simulator/simulator/libero_node/observer.py @@ -78,7 +78,9 @@ def __init__( task = task_suite.get_task(task_id) self.task_description = task.language bddl_file = Path(get_libero_path("bddl_files")) / task.problem_folder / task.bddl_file - init_state = load_init_states(get_libero_path, task, init_state_id) + # Keep an owned copy: every restart must restore this exact MuJoCo state, + # rather than returning the observation cached after the last action. + self.init_state = np.asarray(load_init_states(get_libero_path, task, init_state_id)).copy() self.env = env_cls( bddl_file_name=str(bddl_file), @@ -87,8 +89,13 @@ def __init__( camera_names=["robot0_eye_in_hand", "agentview", "frontview", "galleryview"], ) self.env.seed(seed) + self.reset() + + def reset(self): + """Reset simulator internals and restore the configured initial state.""" self.env.reset() - self.obs = self.env.set_init_state(init_state) + self.obs = self.env.set_init_state(self.init_state.copy()) + return self.obs def step(self, action): action = np.asarray(action, dtype=np.float32) diff --git a/lightx2v_ros/src/simulator/simulator/sim/node.py b/lightx2v_ros/src/simulator/simulator/sim/node.py index 0267c67d4..6bba8b177 100644 --- a/lightx2v_ros/src/simulator/simulator/sim/node.py +++ b/lightx2v_ros/src/simulator/simulator/sim/node.py @@ -20,7 +20,8 @@ ``observation_ready`` indices) goes quiet. The scene itself stays put. - ``resume``/``start`` bumps the observation counter so the policy sees a fresh observation and continues from the current physical state. -- ``restart`` tears the episode down and rebuilds it with a new seed. +- ``restart`` tears the episode down, rebuilds it, and waits in ``ready`` until + the user explicitly starts the new episode. - ``set_task`` rebuilds the env with a different task/scenario (may take tens of seconds; the node publishes a "switching" status first). """ @@ -254,7 +255,7 @@ def _finish_episode(self, outcome: str): if self.loop: self._start_next_episode() - def _start_next_episode(self): + def _start_next_episode(self, *, wait_for_start=False): self.state = SWITCHING self.publish_status() try: @@ -265,7 +266,25 @@ def _start_next_episode(self): self.state = FAILURE self.publish_status() raise - self._begin_episode() + if wait_for_start: + self._ready_episode() + else: + self._begin_episode() + + def _ready_episode(self): + """Publish a fresh episode without allowing the policy to act yet.""" + self.episode_index += 1 + self.episode_step = 0 + self.step_index += 1 + self.success = False + self.state = READY + self._refresh_max_steps() + self.get_logger().info( + f"episode {self.episode_index} reset and waiting for start " + f"(global step {self.step_index}): {self.env.task_description!r}" + ) + self.publish_observation() + self.publish_status() def _begin_episode(self): self.episode_index += 1 @@ -346,7 +365,7 @@ def _cmd_resume(self): self.publish_status() def _cmd_restart(self): - self._start_next_episode() + self._start_next_episode(wait_for_start=True) def _cmd_set_task(self, command): if not self.env.supports_task_switch: From aec13d0379f26a72b2b6375ee25c615bfe994083 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Tue, 14 Jul 2026 09:07:22 +0000 Subject: [PATCH 6/6] update --- .../simulator/simulator/libero_node/env.py | 62 +++++++- .../simulator/libero_node/observer.py | 69 +++++++- .../simulator/simulator/robotwin_node/env.py | 4 +- .../src/simulator/simulator/sim/node.py | 5 +- .../image_web_viewer_node/page.py | 150 ++++++++++++++++-- 5 files changed, 262 insertions(+), 28 deletions(-) diff --git a/lightx2v_ros/src/simulator/simulator/libero_node/env.py b/lightx2v_ros/src/simulator/simulator/libero_node/env.py index d1aa67ddf..7128a6863 100644 --- a/lightx2v_ros/src/simulator/simulator/libero_node/env.py +++ b/lightx2v_ros/src/simulator/simulator/libero_node/env.py @@ -6,7 +6,7 @@ from common.contract import EnvContract from ..sim.base_env import BaseSimEnv, Observation -from .observer import LiberoActionObserver, default_libero_root +from .observer import LiberoActionObserver, build_task_catalog, default_libero_root def quat_to_axis_angle(quat): @@ -39,14 +39,26 @@ def __init__( libero_root=None, ): super().__init__(contract) + self.image_size = int(image_size) + self.libero_root = libero_root self.observer = LiberoActionObserver( benchmark_name=benchmark, task_id=int(task_id), init_state_id=int(init_state_id), - image_size=int(image_size), + image_size=self.image_size, seed=int(seed), - libero_root=libero_root, + libero_root=self.libero_root, ) + self._task_catalog = build_task_catalog(self.observer.benchmark_module) + self._sync_metadata() + + def _sync_metadata(self): + self.benchmark = self.observer.benchmark_name + self.task_id = self.observer.task_id + self.init_state_id = self.observer.init_state_id + self.task_name = self.observer.task_key + self.task_config = str(self.init_state_id) + self.seed = self.observer.seed @property def task_description(self) -> str: @@ -72,6 +84,50 @@ def _state(self, obs) -> np.ndarray: gripper = np.asarray(obs["robot0_gripper_qpos"], dtype=np.float32) return np.concatenate([pos, axis_angle, gripper]).astype(np.float32) + @property + def supports_task_switch(self) -> bool: + return True + + def list_tasks(self): + return [ + { + "value": key, + "label": f"[{item['benchmark']} {item['task_id']}] {item['language']}", + } + for key, item in self._task_catalog.items() + ] + + def list_task_configs(self): + return [str(index) for index in range(self.observer.num_init_states)] + + def set_task(self, task_name: str, task_config: str = "", seed=None) -> Observation: + task_key = str(task_name).strip() + task = self._task_catalog.get(task_key) + if task is None: + raise ValueError(f"unknown LIBERO task {task_key!r}") + + init_state_id = self.init_state_id if str(task_config).strip() == "" else int(task_config) + new_seed = self.seed + 1 if seed is None or str(seed).strip() == "" else int(seed) + + # Construct the replacement first so an invalid task/config leaves the + # currently displayed environment alive and usable. + new_observer = LiberoActionObserver( + benchmark_name=task["benchmark"], + task_id=task["task_id"], + init_state_id=init_state_id, + image_size=self.image_size, + seed=new_seed, + libero_root=self.libero_root, + ) + old_observer = self.observer + self.observer = new_observer + self._sync_metadata() + try: + old_observer.close() + except Exception: + pass + return self._observation() + def close(self) -> None: self.observer.close() diff --git a/lightx2v_ros/src/simulator/simulator/libero_node/observer.py b/lightx2v_ros/src/simulator/simulator/libero_node/observer.py index b0f6f60c0..a26d5827b 100644 --- a/lightx2v_ros/src/simulator/simulator/libero_node/observer.py +++ b/lightx2v_ros/src/simulator/simulator/libero_node/observer.py @@ -4,6 +4,14 @@ import numpy as np +LIBERO_BENCHMARKS = ( + "libero_spatial", + "libero_object", + "libero_goal", + "libero_10", + "libero_90", +) + def default_libero_root(): return Path(__file__).resolve().parent / "LIBERO" @@ -55,10 +63,40 @@ def load_libero(libero_root): def load_init_states(get_libero_path, task, init_state_id): + state, _ = load_init_state(get_libero_path, task, init_state_id) + return state + + +def load_init_state(get_libero_path, task, init_state_id): import torch init_states_path = Path(get_libero_path("init_states")) / task.problem_folder / task.init_states_file - return torch.load(init_states_path, weights_only=False)[init_state_id] + init_states = torch.load(init_states_path, map_location="cpu", weights_only=False) + index = int(init_state_id) + if index < 0 or index >= len(init_states): + raise ValueError(f"init_state_id {index} is out of range for {task.name!r}; expected 0..{len(init_states) - 1}") + return init_states[index], len(init_states) + + +def build_task_catalog(benchmark_module): + """Return stable UI task ids mapped to their LIBERO suite/task metadata.""" + factories = benchmark_module.get_benchmark_dict() + catalog = {} + for benchmark_name in LIBERO_BENCHMARKS: + factory = factories.get(benchmark_name) + if factory is None: + continue + task_suite = factory() + for task_id in range(task_suite.get_num_tasks()): + task = task_suite.get_task(task_id) + key = f"{benchmark_name}/{task_id}" + catalog[key] = { + "benchmark": benchmark_name, + "task_id": task_id, + "task_name": task.name, + "language": task.language, + } + return catalog class LiberoActionObserver: @@ -74,23 +112,40 @@ def __init__( self.libero_root = Path(libero_root or default_libero_root()).expanduser() benchmark, get_libero_path, env_cls = load_libero(self.libero_root) - task_suite = benchmark.get_benchmark_dict()[benchmark_name.lower()]() - task = task_suite.get_task(task_id) + self.benchmark_module = benchmark + self.benchmark_name = str(benchmark_name).strip().lower() + factories = benchmark.get_benchmark_dict() + if self.benchmark_name not in factories or self.benchmark_name not in LIBERO_BENCHMARKS: + raise ValueError(f"unknown LIBERO benchmark {benchmark_name!r}; available: {', '.join(LIBERO_BENCHMARKS)}") + task_suite = factories[self.benchmark_name]() + self.task_id = int(task_id) + if self.task_id < 0 or self.task_id >= task_suite.get_num_tasks(): + raise ValueError(f"task_id {self.task_id} is out of range for {self.benchmark_name!r}; expected 0..{task_suite.get_num_tasks() - 1}") + task = task_suite.get_task(self.task_id) + self.task = task self.task_description = task.language bddl_file = Path(get_libero_path("bddl_files")) / task.problem_folder / task.bddl_file # Keep an owned copy: every restart must restore this exact MuJoCo state, # rather than returning the observation cached after the last action. - self.init_state = np.asarray(load_init_states(get_libero_path, task, init_state_id)).copy() + self.init_state_id = int(init_state_id) + init_state, self.num_init_states = load_init_state(get_libero_path, task, self.init_state_id) + self.init_state = np.asarray(init_state).copy() + self.image_size = int(image_size) + self.seed = int(seed) self.env = env_cls( bddl_file_name=str(bddl_file), - camera_heights=image_size, - camera_widths=image_size, + camera_heights=self.image_size, + camera_widths=self.image_size, camera_names=["robot0_eye_in_hand", "agentview", "frontview", "galleryview"], ) - self.env.seed(seed) + self.env.seed(self.seed) self.reset() + @property + def task_key(self): + return f"{self.benchmark_name}/{self.task_id}" + def reset(self): """Reset simulator internals and restore the configured initial state.""" self.env.reset() diff --git a/lightx2v_ros/src/simulator/simulator/robotwin_node/env.py b/lightx2v_ros/src/simulator/simulator/robotwin_node/env.py index 3a6d83042..acd5c4d53 100644 --- a/lightx2v_ros/src/simulator/simulator/robotwin_node/env.py +++ b/lightx2v_ros/src/simulator/simulator/robotwin_node/env.py @@ -112,9 +112,7 @@ def _prepare_runtime(self) -> None: "envs/": (root / "envs").is_dir(), "task_config/": (root / "task_config").is_dir(), "assets/": (root / "assets").is_dir(), - "assets/objects/objaverse/list.json": ( - root / "assets" / "objects" / "objaverse" / "list.json" - ).is_file(), + "assets/objects/objaverse/list.json": (root / "assets" / "objects" / "objaverse" / "list.json").is_file(), "assets/objects/same.json": (root / "assets" / "objects" / "same.json").is_file(), } missing = [name for name, exists in required_paths.items() if not exists] diff --git a/lightx2v_ros/src/simulator/simulator/sim/node.py b/lightx2v_ros/src/simulator/simulator/sim/node.py index 6bba8b177..b56fceff1 100644 --- a/lightx2v_ros/src/simulator/simulator/sim/node.py +++ b/lightx2v_ros/src/simulator/simulator/sim/node.py @@ -279,10 +279,7 @@ def _ready_episode(self): self.success = False self.state = READY self._refresh_max_steps() - self.get_logger().info( - f"episode {self.episode_index} reset and waiting for start " - f"(global step {self.step_index}): {self.env.task_description!r}" - ) + self.get_logger().info(f"episode {self.episode_index} reset and waiting for start (global step {self.step_index}): {self.env.task_description!r}") self.publish_observation() self.publish_status() diff --git a/lightx2v_ros/src/visualization/visualization/image_web_viewer_node/page.py b/lightx2v_ros/src/visualization/visualization/image_web_viewer_node/page.py index 80f5ea8fb..0696e1018 100644 --- a/lightx2v_ros/src/visualization/visualization/image_web_viewer_node/page.py +++ b/lightx2v_ros/src/visualization/visualization/image_web_viewer_node/page.py @@ -83,6 +83,12 @@ font-size: 13px; } input[type=number] { width: 90px; } + #task-select { width: min(560px, 70vw); } + #config-select { min-width: 90px; } + .config-feedback { color: var(--muted); font-size: 13px; } + .config-feedback.pending { color: var(--yellow); } + .config-feedback.success { color: var(--green); } + .config-feedback.error { color: var(--red); } button { border: 1px solid var(--panel-border); border-radius: 6px; @@ -167,6 +173,8 @@ const POLICY_CAMS = __POLICY_CAMS__; let lastState = null; let statusData = {}; + let configDirty = false; + let applyingConfig = false; function el(id) { return document.getElementById(id); } @@ -184,19 +192,127 @@ async function post(cmd, extra) { const body = Object.assign({ cmd: cmd }, extra || {}); - try { - await fetch("/control", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(body) }); - } catch (e) { console.error(e); } + const response = await fetch("/control", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }); + const result = await response.json(); + if (!response.ok || !result.ok) throw new Error(result.error || `HTTP ${response.status}`); setTimeout(refreshStatus, 300); + return result; } - function applyTask() { + function selectedConfig() { const task = el("task-select").value; const config = el("config-select").value; const seed = el("seed-input").value; const extra = { task_name: task, task_config: config }; if (seed !== "") extra.seed = parseInt(seed, 10); - post("set_task", extra); + return { task, config, seed: seed === "" ? null : parseInt(seed, 10), extra }; + } + + function setConfigFeedback(text, kind) { + const feedback = el("config-feedback"); + feedback.textContent = text; + feedback.className = `config-feedback${kind ? ` ${kind}` : ""}`; + } + + function markConfigDirty() { + configDirty = true; + setConfigFeedback("配置尚未应用", "pending"); + } + + async function waitForConfig(expected, timeoutMs = 60000) { + const deadline = Date.now() + timeoutMs; + let sawSwitching = false; + while (Date.now() < deadline) { + await refreshStatus(); + sawSwitching ||= statusData.state === "switching"; + const seedMatches = expected.seed === null || Number(statusData.seed) === expected.seed; + if (statusData.state !== "switching" && + statusData.task_name === expected.task && + String(statusData.task_config) === String(expected.config) && seedMatches) { + return statusData; + } + if (sawSwitching && statusData.state === "failure") { + throw new Error("模拟器切换失败，请查看 libero_node 日志"); + } + await new Promise(resolve => setTimeout(resolve, 250)); + } + throw new Error("等待模拟器应用配置超时"); + } + + async function waitForRestart(expected, previousEpisode, timeoutMs = 60000) { + const deadline = Date.now() + timeoutMs; + let sawSwitching = false; + while (Date.now() < deadline) { + await refreshStatus(); + sawSwitching ||= statusData.state === "switching"; + const configMatches = + statusData.task_name === expected.task && + String(statusData.task_config) === String(expected.config) && + Number(statusData.seed) === Number(expected.seed); + if (configMatches && statusData.state === "ready" && + Number(statusData.episode) > Number(previousEpisode)) { + return statusData; + } + if (sawSwitching && statusData.state === "failure") { + throw new Error("配置已应用，但自动重启失败，请查看 libero_node 日志"); + } + await new Promise(resolve => setTimeout(resolve, 250)); + } + throw new Error("配置已应用，但等待自动重启超时"); + } + + async function applyTask() { + if (applyingConfig) return false; + const selected = selectedConfig(); + applyingConfig = true; + updateButtons(statusData.state); + setConfigFeedback("正在重建仿真环境…", "pending"); + try { + await post("set_task", selected.extra); + let applied = await waitForConfig(selected); + let restarted = false; + if (applied.env === "libero") { + const appliedEpisode = Number(applied.episode); + const restartExpected = { + task: applied.task_name, + config: String(applied.task_config), + seed: Number(applied.seed), + }; + setConfigFeedback("配置已应用，正在自动重启…", "pending"); + await post("restart"); + applied = await waitForRestart(restartExpected, appliedEpisode); + restarted = true; + } + configDirty = false; + setConfigFeedback( + `${restarted ? "已应用并重启" : "已应用"} ${applied.task_name} · 场景 ${applied.task_config} · seed ${applied.seed}`, + "success", + ); + return true; + } catch (error) { + console.error(error); + setConfigFeedback(`应用失败：${error.message}`, "error"); + return false; + } finally { + applyingConfig = false; + updateButtons(statusData.state); + } + } + + async function startEvaluation() { + if (lastState === "paused") { + await post("resume"); + return; + } + // A common failure mode was selecting a new task and pressing Start while + // the simulator still held the old task. Apply and verify pending config + // first; this Start click remains the explicit authorization to run. + if (configDirty && !(await applyTask())) return; + await post("start"); } function fillSelect(select, options, current) { @@ -205,10 +321,18 @@ select.innerHTML = ""; for (const opt of options) { const o = document.createElement("option"); - o.value = opt; o.textContent = opt; + if (opt && typeof opt === "object") { + o.value = String(opt.value); + o.textContent = String(opt.label || opt.value); + } else { + o.value = String(opt); + o.textContent = String(opt); + } select.appendChild(o); } - if (current && options.includes(current)) select.value = current; + if (current && Array.from(select.options).some(o => o.value === String(current))) { + select.value = String(current); + } select.dataset.filled = "1"; } @@ -232,10 +356,10 @@ } function updateButtons(s) { - el("btn-start").disabled = (s === "running" || s === "switching"); + el("btn-start").disabled = (s === "running" || s === "switching" || applyingConfig); el("btn-pause").disabled = (s !== "running"); el("btn-restart").disabled = (s === "switching"); - el("btn-apply").disabled = (s === "switching"); + el("btn-apply").disabled = (s === "switching" || applyingConfig); el("btn-start").textContent = (s === "paused") ? "继续" : "开始"; } @@ -293,10 +417,13 @@ } document.addEventListener("DOMContentLoaded", () => { - el("btn-start").addEventListener("click", () => post(lastState === "paused" ? "resume" : "start")); + el("btn-start").addEventListener("click", () => startEvaluation().catch(console.error)); el("btn-pause").addEventListener("click", () => post("pause")); el("btn-restart").addEventListener("click", () => post("restart")); el("btn-apply").addEventListener("click", applyTask); + el("task-select").addEventListener("change", markConfigDirty); + el("config-select").addEventListener("change", markConfigDirty); + el("seed-input").addEventListener("input", markConfigDirty); document.querySelectorAll(".cam-toggles input").forEach((box) => { box.addEventListener("change", () => toggleCam(box.value, box.checked)); }); @@ -342,7 +469,8 @@ def render_index(cameras, title="LightX2V ROS", policy_cameras=None): - + + 修改后请应用配置