diff --git "a/docs/2026-08-07_G2\351\200\202\351\205\215\350\256\255\347\273\203\346\216\250\347\220\206\350\257\264\346\230\216.md" "b/docs/2026-08-07_G2\351\200\202\351\205\215\350\256\255\347\273\203\346\216\250\347\220\206\350\257\264\346\230\216.md" new file mode 100644 index 00000000..b7fb1d09 --- /dev/null +++ "b/docs/2026-08-07_G2\351\200\202\351\205\215\350\256\255\347\273\203\346\216\250\347\220\206\350\257\264\346\230\216.md" @@ -0,0 +1,234 @@ +# G2 适配 DreamZero 训练与推理说明 + +本文只说明本 PR 保留的最小 G2 适配主路径：G2 joint-space 数据转换、relative action 训练、checkpoint 预检、G2 推理服务和真机客户端解码。 + +## 适配目标 + +G2 适配让 DreamZero 使用 G2 的三路视觉、双臂关节状态和双臂动作进行训练，并在推理时输出 G2 可执行的 16 维关节目标。 + +核心数据契约： + +- 视觉输入：`top_head`、`hand_left`、`hand_right` 三路相机。 +- 状态输入：16 维，布局为左臂 7 维关节、左夹爪、右臂 7 维关节、右夹爪。 +- 动作输出：16 维，布局与状态一致。 +- action horizon：24。 +- 训练帧数：`num_frames=33`。 +- 图像尺寸：`320x176`。 +- embodiment tag：`g2`。 + +## 数据链路 + +G2 原始 LeRobot 数据先转换为 DreamZero/GEAR 格式： + +```bash +python scripts/data/convert_lerobot_g2_to_gear.py \ + --source /path/to/g2_lerobot_dataset \ + --output /path/to/g2_gear_dataset \ + --test-episodes 10 \ + --video-width 320 \ + --video-height 176 +``` + +转换脚本会： + +- 校验 state/action 是否符合固定 16 维 G2 joint layout。 +- 把三路相机映射为 `video.top_head`、`video.hand_left`、`video.hand_right`。 +- 生成 `train/` 和 `test/` 两个物理隔离 split。 +- 写入 `meta/info.json`、`meta/modality.json`、`meta/embodiment.json`、`meta/stats.json`。 +- 写入 `meta/relative_stats_dreamzero.json`，用于 relative action 归一化和 motion mask。 + +训练时 `G2_DATA_ROOT` 必须指向转换后的 `train/` 目录。 + +## 配置适配 + +G2 modality 注册在： + +```text +groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml +``` + +G2 的 state/action key： + +```yaml +state: + - state.left_joint_position + - state.left_gripper_position + - state.right_joint_position + - state.right_gripper_position +action: + - action.left_joint_position + - action.left_gripper_position + - action.right_joint_position + - action.right_gripper_position +``` + +所有 state/action key 使用 `q99` 归一化，最后通过 `ConcatTransform` 拼成模型输入。 + +G2 relative action 数据配置在： + +```text +groot/vla/configs/data/dreamzero/g2_relative.yaml +``` + +关键设置： + +```yaml +relative_action: true +relative_action_keys: + - left_joint_position + - right_joint_position +active_hold_index_path: ${g2_data_root}/meta/g2_active_hold_windows.json +active_window_ratio: 0.8 +``` + +含义是：双臂关节学习相对当前状态的增量，夹爪保持 policy space 数值；训练采样时用 active/hold window 提高有动作片段的比例。 + +## 模型适配 + +G2 embodiment tag 注册在： + +```text +groot/vla/data/schema/embodiment_tags.py +``` + +```python +G2 = "g2" +``` + +transform 侧增加 G2 embedding id： + +```text +groot/vla/configs/model/dreamzero/transform/base.yaml +``` + +训练脚本显式覆盖： + +```bash +++model_specific_transform.embodiment_tag_mapping.g2=33 +``` + +action head 的核心改动在： + +```text +groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py +``` + +本 PR 保留的核心能力： + +- 支持 G2 的 16 维 action 输出契约。 +- 支持 action loss / dynamics loss 权重配置。 +- 支持 motion mask，根据 `relative_stats_dreamzero.json` 把归一化 action delta 还原到弧度尺度后判断是否为有效运动。 +- 支持 LoRA checkpoint 按训练时 base model 加载，避免 LoRA delta 套到错误底座。 + +checkpoint 加载相关逻辑在： + +```text +groot/vla/model/dreamzero/base_vla.py +groot/vla/model/n1_5/sim_policy.py +``` + +LoRA-only checkpoint 推理时会读取训练配置里的 `pretrained_model_path`，先恢复 DreamZero base，再加载 LoRA 权重。 + +## 训练主路径 + +主训练入口： + +```text +scripts/train/train_dreamzero_g2_joint_lora.sh +``` + +典型运行： + +```bash +G2_DATA_ROOT=/data/.../g2_gear/train \ +OUTPUT_DIR=/data/.../dreamzero_g2_joint_lora \ +PRETRAINED_MODEL_PATH=/data/.../DreamZero-AgiBot \ +WAN_CKPT_DIR=/data/.../Wan2.1-I2V-14B-480P \ +TOKENIZER_DIR=/data/.../umt5-xxl \ +GPU_IDS=4,5,6,7 \ +bash scripts/train/train_dreamzero_g2_joint_lora.sh +``` + +关键 Hydra 参数： + +```bash +data=dreamzero/g2_relative +train_architecture=lora +num_frames=33 +action_horizon=24 +num_views=3 +image_resolution_width=320 +image_resolution_height=176 +save_lora_only=true +max_chunk_size=4 +``` + +训练脚本会在启动前检查： + +- G2 数据目录和 Wan/T5/CLIP/VAE checkpoint 是否存在。 +- `meta/info.json`、`meta/modality.json`、`meta/embodiment.json`、`meta/stats.json`、`meta/relative_stats_dreamzero.json` 是否存在。 +- episode 数、parquet 数、视频数、G2 16 维 state/action、三路相机是否符合预期。 + +## 推理主路径 + +G2 推理服务入口： + +```text +scripts/run_g2_server_9443_final.sh +``` + +典型运行： + +```bash +MODEL_PATH=/data/.../checkpoint-1500 \ +WAN_CKPT_DIR=/data/.../Wan2.1-I2V-14B-480P \ +TOKENIZER_PATH=/data/.../umt5-xxl \ +PORT=9443 \ +bash scripts/run_g2_server_9443_final.sh +``` + +启动前脚本会执行： + +```bash +python scripts/audit_g2_checkpoint.py "$MODEL_PATH" +``` + +预检会确认 checkpoint 的 action horizon、`num_frames`、action 维度、base model 路径和 LoRA/action 权重结构是否满足 G2 契约。 + +服务端使用： + +```bash +python -m torch.distributed.run \ + --standalone \ + --nproc_per_node=2 \ + socket_optimized_AR_g2.py \ + --port "$PORT" \ + --model-path "$MODEL_PATH" \ + --wan-ckpt-dir "$WAN_CKPT_DIR" \ + --tokenizer-path "$TOKENIZER_PATH" \ + --embodiment-tag g2 +``` + +真机客户端在： + +```text +robot_live_client_g2.py +``` + +relative action checkpoint 的输出不是直接下发的绝对目标。客户端会用当前执行边界的机器人状态还原： + +```text +absolute_target = current_state + decoded_relative_delta +``` + +然后再做 GDK 关节限位裁剪并下发到 G2。 + +## 最小闭环 + +1. 用 `scripts/data/convert_lerobot_g2_to_gear.py` 生成 GEAR train/test 数据。 +2. 用 `scripts/data/build_g2_active_hold_windows.py` 生成 `meta/g2_active_hold_windows.json`。 +3. 确认 `meta/embodiment.json` 为 `{"embodiment_tag": "g2"}`。 +4. 用 `scripts/train/train_dreamzero_g2_joint_lora.sh` 训练 relative-action LoRA。 +5. 用 `scripts/audit_g2_checkpoint.py` 校验 checkpoint。 +6. 用 `scripts/run_g2_server_9443_final.sh` 启动 G2 server。 +7. 用 `robot_live_client_g2.py` 按 relative action 模式连接服务端并执行。 diff --git a/groot/vla/configs/conf.yaml b/groot/vla/configs/conf.yaml index 48da05bc..2aef043a 100644 --- a/groot/vla/configs/conf.yaml +++ b/groot/vla/configs/conf.yaml @@ -33,6 +33,7 @@ trainer: profile_record_shapes: false # record tensor shapes (adds overhead) profile_with_stack: false # record Python stack traces profile_memory: false # record memory allocation + milestone_save_steps: ${milestone_save_steps} # === Training Arguments === @@ -67,6 +68,7 @@ num_train_epochs: 1000 max_steps: -1 save_strategy: steps save_steps: 500 +milestone_save_steps: null eval_strategy: "no" # there has to be a double quote; otherwise a bare `no` will be interpreted as False save_total_limit: 8 report_to: wandb @@ -80,6 +82,7 @@ eval_bf16: true torch_compile_mode: null pretrained_model_path: null +pretrained_lora_path: null only_tune_projectors: false save_llm: false diff --git a/groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml b/groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml index cfcdd449..6998adf5 100644 --- a/groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml +++ b/groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml @@ -345,10 +345,86 @@ transform_yam: # Modality Configs ################################################################################ +################################################################################ +# g2 (Unitree G2: dual-arm joints + grippers, state/action 16, 3 views) +################################################################################ + +modality_config_g2: + video: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] + eval_delta_indices: [-3,-2,-1,0] + modality_keys: + - video.top_head + - video.hand_left + - video.hand_right + state: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - state.left_joint_position + - state.left_gripper_position + - state.right_joint_position + - state.right_gripper_position + action: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23] + modality_keys: + - action.left_joint_position + - action.left_gripper_position + - action.right_joint_position + - action.right_gripper_position + language: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - annotation.language.action_text + +transform_g2: + _target_: groot.vla.data.transform.ComposedModalityTransform + transforms: + - <<: *totensor_cfg + apply_to: ${modality_config_g2.video.modality_keys} + - <<: *crop_cfg + apply_to: ${modality_config_g2.video.modality_keys} + - <<: *resize_cfg + apply_to: ${modality_config_g2.video.modality_keys} + - <<: *color_jitter_cfg + apply_to: ${modality_config_g2.video.modality_keys} + - <<: *to_numpy_cfg + apply_to: ${modality_config_g2.video.modality_keys} + + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_g2.state.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_g2.state.modality_keys} + normalization_modes: + state.left_joint_position: q99 + state.left_gripper_position: q99 + state.right_joint_position: q99 + state.right_gripper_position: q99 + + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_g2.action.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_g2.action.modality_keys} + normalization_modes: + action.left_joint_position: q99 + action.left_gripper_position: q99 + action.right_joint_position: q99 + action.right_gripper_position: q99 + + - _target_: groot.vla.data.transform.ConcatTransform + video_concat_order: ${modality_config_g2.video.modality_keys} + state_concat_order: ${modality_config_g2.state.modality_keys} + action_concat_order: ${modality_config_g2.action.modality_keys} + - ${model_specific_transform} + modality_configs: oxe_droid: ${modality_config_oxe_droid} agibot: ${modality_config_agibot} yam: ${modality_config_yam} + g2: ${modality_config_g2} ################################################################################ # Transforms @@ -358,6 +434,7 @@ transforms: oxe_droid: ${transform_oxe_droid} agibot: ${transform_agibot} yam: ${transform_yam} + g2: ${transform_g2} ################################################################################ # Metadata Versions @@ -367,6 +444,7 @@ metadata_versions: oxe_droid: '0221' agibot: '0221' yam: '0221' + g2: '0221' ################################################################################ # FPS (per embodiment, null means use dataset default) @@ -374,3 +452,4 @@ metadata_versions: fps: yam: 30 + g2: 30 diff --git a/groot/vla/configs/data/dreamzero/g2_relative.yaml b/groot/vla/configs/data/dreamzero/g2_relative.yaml new file mode 100644 index 00000000..d935912a --- /dev/null +++ b/groot/vla/configs/data/dreamzero/g2_relative.yaml @@ -0,0 +1,52 @@ +# @package _global_ + +defaults: + - dreamzero/base_48_wan_fine_aug_relative + - _self_ + +max_state_dim: 64 +use_global_metadata: false +relative_action: true +relative_action_per_horizon: false +relative_action_keys: + - left_joint_position + - right_joint_position +max_chunk_size: 5 +dataset_shard_sampling_rate: 0.1 +active_hold_index_path: ${g2_data_root}/meta/g2_active_hold_windows.json +active_window_ratio: 0.8 +mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec +single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID + +# Override with g2_data_root=/path/to/the/GEAR/train/directory. +g2_data_root: ??? + +train_dataset: + _target_: ${mixture_dataset_cls} + _convert_: object + mixture_spec: + - dataset_path: + g2: + - ${g2_data_root} + dataset_weight: 1.0 + distribute_weights: true + + dataset_class: ${single_dataset_cls} + all_modality_configs: ${modality_configs} + all_transforms: ${transforms} + metadata_versions: ${metadata_versions} + fps: ${fps} + dataset_kwargs: + video_backend: decord + use_global_metadata: ${use_global_metadata} + max_chunk_size: ${max_chunk_size} + relative_action: ${relative_action} + relative_action_keys: ${relative_action_keys} + relative_action_per_horizon: ${relative_action_per_horizon} + mixture_kwargs: + training: true + balance_dataset_weights: false + seed: 42 + shard_sampling_rate: ${dataset_shard_sampling_rate} + active_hold_index_path: ${active_hold_index_path} + active_window_ratio: ${active_window_ratio} diff --git a/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf.yaml b/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf.yaml index 8cccdc63..fb133969 100644 --- a/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf.yaml +++ b/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf.yaml @@ -38,7 +38,7 @@ action_head_cfg: model_dtype: float32 max_state_dim: ${max_state_dim} max_action_dim: ${max_action_dim} - action_loss_embodiment_ids: [26, 17, 32] + action_loss_embodiment_ids: [26, 17, 32, 33] hidden_size: ${hidden_size} input_embedding_dim: 1536 backbone_embedding_dim: ${backbone_hidden_size} @@ -106,3 +106,12 @@ action_head_cfg: tune_projector: true tune_diffusion_model: true + action_only_training: false + dynamics_loss_weight: 1.0 + action_loss_weight: 1.0 + # Auxiliary action constraints (keep small to avoid disrupting flow-matching): + # L_start: first action of the chunk anchors to the current joint state. + # L_video_to_action: actions recovered from predicted video latent must + # match the action chunk (action depends on implicit dynamics). + action_start_loss_weight: 0.05 + action_video_consistency_loss_weight: 0.10 diff --git a/groot/vla/configs/model/dreamzero/transform/base.yaml b/groot/vla/configs/model/dreamzero/transform/base.yaml index 3873b9ab..cd6439dd 100644 --- a/groot/vla/configs/model/dreamzero/transform/base.yaml +++ b/groot/vla/configs/model/dreamzero/transform/base.yaml @@ -33,6 +33,9 @@ embodiment_tag_to_projector_index: oxe_plex: 30 dream: 31 yam: 32 + # Logical G2 identity. The current DreamZero-AgiBot action adapter is + # single-embodiment internally and maps this to local adapter slot 0. + g2: 33 xdof: 22 gr1_unified_segmentation: 14 language_table_sim: 7 diff --git a/groot/vla/data/dataset/lerobot_sharded.py b/groot/vla/data/dataset/lerobot_sharded.py old mode 100755 new mode 100644 index f8aa338c..3d8df304 --- a/groot/vla/data/dataset/lerobot_sharded.py +++ b/groot/vla/data/dataset/lerobot_sharded.py @@ -1291,6 +1291,8 @@ def __init__( shard_sampling_rate: float = 0.5, num_shards_to_sample: int = 2**20, allow_padding_at_end: bool = False, + active_hold_index_path: str | None = None, + active_window_ratio: float = 0.8, ): """ Initialize the mixture dataset. @@ -1317,6 +1319,25 @@ def __init__( # Set properties self.shard_sampling_rate = shard_sampling_rate self.num_shards_to_sample = num_shards_to_sample + self.active_window_ratio = float(active_window_ratio) + if not 0.0 <= self.active_window_ratio <= 1.0: + raise ValueError("active_window_ratio must be in [0, 1]") + self._active_steps: set[tuple[int, int]] | None = None + if active_hold_index_path: + index_path = Path(active_hold_index_path) + with index_path.open() as stream: + active_hold = json.load(stream) + self._active_steps = { + (int(episode), int(step)) + for episode, steps in active_hold["active"].items() + for step in steps + } + print( + "[G2 ACTIVE/HOLD SAMPLING] " + f"index={index_path} active={len(self._active_steps)} " + f"target_ratio={self.active_window_ratio:.2f} " + f"threshold={active_hold['arm_motion_threshold']:.8f}" + ) # Calculate shard sampling weights all_shard_sampling_weights = [] @@ -1501,9 +1522,37 @@ def __iter__(self): allowed_indices = allowed_indices[allowed_indices <= allowed_length] for i in allowed_indices: all_steps.append((trajectory_id, i)) + sample_count = int(dataset.num_steps_per_shard * self.shard_sampling_rate) if self.training: rng.shuffle(all_steps) - sampled_steps = all_steps[: int(dataset.num_steps_per_shard * self.shard_sampling_rate)] + if self._active_steps is None: + sampled_steps = all_steps[:sample_count] + else: + active_steps = [ + step for step in all_steps if step in self._active_steps + ] + hold_steps = [ + step for step in all_steps if step not in self._active_steps + ] + if self.training: + rng.shuffle(active_steps) + rng.shuffle(hold_steps) + active_count = min( + len(active_steps), + int(round(sample_count * self.active_window_ratio)), + ) + hold_count = min(len(hold_steps), sample_count - active_count) + sampled_steps = ( + active_steps[:active_count] + hold_steps[:hold_count] + ) + if len(sampled_steps) < sample_count: + used = set(sampled_steps) + sampled_steps.extend( + step for step in all_steps if step not in used + ) + sampled_steps = sampled_steps[:sample_count] + if self.training: + rng.shuffle(sampled_steps) for trajectory_id, step_index in sampled_steps: # print( # f"Loading step data from rank {self.rank}, worker {self.worker_id}: {dataset_index} {trajectory_id}, {step_index}" diff --git a/groot/vla/data/schema/embodiment_tags.py b/groot/vla/data/schema/embodiment_tags.py old mode 100755 new mode 100644 index e7eaec04..dd09e353 --- a/groot/vla/data/schema/embodiment_tags.py +++ b/groot/vla/data/schema/embodiment_tags.py @@ -153,6 +153,8 @@ class EmbodimentTag(Enum): """ AGIBOT = "agibot" + G2 = "g2" + YAM = "yam" DREAM = "dream" diff --git a/groot/vla/experiment/base.py b/groot/vla/experiment/base.py index 5e9a7d28..faaf8dc4 100644 --- a/groot/vla/experiment/base.py +++ b/groot/vla/experiment/base.py @@ -144,6 +144,11 @@ def on_save(self, args, state, control, **kwargs): print(f"Copying wandb_config.json from {wandb_config_src} to {wandb_config_dst}") shutil.copy2(wandb_config_src, wandb_config_dst) + trainer_state = checkpoint_dir / "trainer_state.json" + training_state = checkpoint_dir / "training_state.json" + if trainer_state.exists(): + shutil.copy2(trainer_state, training_state) + class ProfCallback(transformers.TrainerCallback): """Callback to manage PyTorch profiler during training. @@ -481,14 +486,77 @@ def save_model(self, output_dir: Optional[str], _internal_call: bool): else: state_dict = self.model.state_dict() - if self.base_cfg.save_lora_only: - # Save only the trainable parameters - train_key = [k for k, v in self.model.named_parameters() if v.requires_grad] - lora_state_dict = {k: v for k, v in self.model.state_dict().items() if k in train_key} + action_head_config = getattr( + getattr(self.model, "action_head", None), + "config", + None, + ) + action_only_training = bool( + getattr(action_head_config, "action_only_training", False) + ) + + if action_only_training: + adapter_prefixes = ( + "action_head.model.state_encoder.", + "action_head.model.action_encoder.", + "action_head.model.action_decoder.", + ) + adapter_state_dict = { + key: value + for key, value in state_dict.items() + if key.startswith(adapter_prefixes) + } + expected_adapter_keys = { + key + for key in self.model.state_dict() + if key.startswith(adapter_prefixes) + } + if set(adapter_state_dict) != expected_adapter_keys: + missing = sorted( + expected_adapter_keys - set(adapter_state_dict) + ) + unexpected = sorted( + set(adapter_state_dict) - expected_adapter_keys + ) + raise RuntimeError( + "Action adapter save contract mismatch: " + f"missing={missing[:20]} unexpected={unexpected[:20]}" + ) + state_dict = adapter_state_dict + elif self.base_cfg.save_lora_only: + # Save trainable parameters. During the action-only second stage, + # also retain the frozen stage-1 LoRA tensors: inference needs + # those adapters to preserve the already-recovered video model. + train_key = { + k for k, v in self.model.named_parameters() if v.requires_grad + } + lora_state_dict = { + k: v + for k, v in self.model.state_dict().items() + if k in train_key + } state_dict = lora_state_dict if self.args.should_save: ret = self.model.save_pretrained(output_dir, state_dict=state_dict) + if action_only_training: + from safetensors.torch import save_file + + action_adapter_path = os.path.join( + output_dir, + "action_expert.safetensors", + ) + save_file( + { + key: value.detach().cpu().contiguous() + for key, value in state_dict.items() + }, + action_adapter_path, + ) + print( + "[ACTION ADAPTER SAVE] " + f"{action_adapter_path} keys={len(state_dict)}" + ) # can separately save the VLM model for downstream evalualtion if self.base_cfg.save_llm: @@ -701,6 +769,20 @@ def create_model(self, cfg, training_args): safetensors_index_path = os.path.join(ckpt_dir, "model.safetensors.index.json") safetensors_path = os.path.join(ckpt_dir, "model.safetensors") + model_state = model.state_dict() + loaded_base_keys = set() + skipped_base_keys = set() + + def load_compatible_base_weights(weights): + compatible = {} + for key, value in weights.items(): + if key in model_state and model_state[key].shape == value.shape: + compatible[key] = value + loaded_base_keys.add(key) + else: + skipped_base_keys.add(key) + model.load_state_dict(compatible, strict=False) + if os.path.exists(safetensors_index_path): with open(safetensors_index_path, 'r') as f: index = json.load(f) @@ -708,23 +790,74 @@ def create_model(self, cfg, training_args): shard_path = os.path.join(ckpt_dir, shard_file) mprint(f"Loading shard: {shard_path}") shard_state_dict = load_file(shard_path) - model.load_state_dict(shard_state_dict, strict=False) + load_compatible_base_weights(shard_state_dict) del shard_state_dict gc.collect() elif os.path.exists(safetensors_path): state_dict = load_file(safetensors_path) - model.load_state_dict(state_dict, strict=False) + load_compatible_base_weights(state_dict) else: raise FileNotFoundError( f"No weights found at '{ckpt_dir}'. " "Expected 'model.safetensors' or 'model.safetensors.index.json'." ) + if getattr( + model.action_head.config, + "action_only_training", + False, + ): + adapter_prefixes = ( + "action_head.model.state_encoder.", + "action_head.model.action_encoder.", + "action_head.model.action_decoder.", + ) + required_shared_dit_keys = { + key + for key in model_state + if key.startswith("action_head.model.") + and not key.startswith(adapter_prefixes) + } + missing_shared_dit = sorted( + required_shared_dit_keys - loaded_base_keys + ) + if missing_shared_dit: + raise RuntimeError( + "Incomplete DreamZero-AgiBot shared DiT load: " + + ", ".join(missing_shared_dit[:20]) + ) + mprint( + "[BASE LOAD] DreamZero-AgiBot loaded | " + f"shared_dit_keys={len(required_shared_dit_keys)} " + f"skipped_incompatible={len(skipped_base_keys)}" + ) + mprint( + "[EMBODIMENT] g2 logical_id=33 " + "action_adapter_local_slot=0" + ) + if (hasattr(model, 'action_head') and hasattr(model.action_head, 'inject_lora_after_loading') and model.action_head.config.defer_lora_injection): model.action_head.inject_lora_after_loading() + if cfg.pretrained_lora_path is not None: + mprint( + f"Loading pretrained LoRA/action weights from: " + f"{cfg.pretrained_lora_path}" + ) + model.load_lora_weight(cfg.pretrained_lora_path) + + if ( + hasattr(model, "action_head") + and getattr( + model.action_head.config, + "action_only_training", + False, + ) + ): + model.action_head.configure_action_only_training() + mprint("Successfully loaded pretrained weights") model.config.resume_path = model.config._name_or_path = training_args.output_dir diff --git a/groot/vla/experiment/experiment.py b/groot/vla/experiment/experiment.py index e02d36c5..f76d9f69 100644 --- a/groot/vla/experiment/experiment.py +++ b/groot/vla/experiment/experiment.py @@ -7,6 +7,7 @@ import numpy as np from omegaconf import DictConfig import torch +from transformers import TrainerCallback from groot.vla.experiment.base import BaseExperiment, BaseTrainer from groot.vla.utils.action_args_override_utils import apply_action_overrides @@ -17,6 +18,18 @@ INITIAL_ACTIONS_FILENAME = "initial_actions.npz" +class MilestoneSaveCallback(TrainerCallback): + """Request checkpoints only at configured optimizer-step milestones.""" + + def __init__(self, milestones): + self.milestones = frozenset(int(step) for step in milestones) + + def on_step_end(self, args, state, control, **kwargs): + if state.global_step in self.milestones: + control.should_save = True + return control + + class ForceRestart(ValueError): pass @@ -35,8 +48,16 @@ def __init__(self, **kwargs): self.rank = dist.get_rank() self.micro_global_step = 0 + milestones = kwargs.pop("milestone_save_steps", None) super().__init__(**kwargs) + if milestones: + self.add_callback(MilestoneSaveCallback(milestones)) + if self.rank == 0: + print( + "[CHECKPOINT MILESTONES] " + + ",".join(str(int(step)) for step in milestones) + ) def training_step(self, model, inputs, *args, **kwargs): self.micro_global_step += 1 diff --git a/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py b/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py index db8fe0f8..27e5e7e0 100644 --- a/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py +++ b/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py @@ -128,6 +128,60 @@ class WANPolicyHeadConfig(PretrainedConfig): tune_diffusion_model: bool = field( default=True, metadata={"help": "Whether to tune the diffusion model."} ) + action_only_training: bool = field( + default=False, + metadata={"help": "Freeze the shared DiT and train only state/action adapters."}, + ) + dynamics_loss_weight: float = field( + default=1.0, metadata={"help": "Multiplier for the video dynamics loss."} + ) + action_loss_weight: float = field( + default=1.0, metadata={"help": "Multiplier for the action flow-matching loss."} + ) + action_x0_loss_weight: float = field( + default=0.0, + metadata={ + "help": "Multiplier for direct decoded-action reconstruction loss." + }, + ) + action_endpoint_loss_weight: float = field( + default=0.0, + metadata={ + "help": "Multiplier for direct final-step action reconstruction loss." + }, + ) + action_start_loss_weight: float = field( + default=0.0, + metadata={ + "help": "Multiplier for the current-state start consistency loss " + "(first action of each chunk should be close to the current joint " + "state). Soft anchor: state is a training target, not a model input." + }, + ) + action_video_consistency_loss_weight: float = field( + default=0.0, + metadata={ + "help": "Multiplier for the video->action consistency loss " + "(actions recovered from the predicted future video latent should " + "match the action chunk). Increases the action's dependence on the " + "video/dynamics representation. Keep small initially to avoid " + "disrupting flow-matching." + }, + ) + motion_mask_enabled: bool = field( + default=False, + metadata={"help": "Mask stationary action timesteps out of the action loss."}, + ) + motion_mask_threshold_rad: float = field( + default=0.03, + metadata={"help": "Joint movement threshold in radians for motion masking."}, + ) + motion_mask_stats_path: str | None = field( + default=None, + metadata={ + "help": "Relative-action stats used to convert normalized action deltas to radians." + }, + ) load_pretrained_det_decode_layer_path: str = field( default=None, metadata={"help": "Path to pretrained detection model."} ) @@ -238,6 +292,13 @@ def __init__( self.model = instantiate(config.diffusion_model_cfg) self.action_dim = config.action_dim self.action_horizon = config.action_horizon + + # Optional video->action consistency head (created lazily on first + # forward once the video-latent channel count is known). Recovering + # actions from the predicted future video latent makes the action + # depend on the learned implicit dynamics. + self.video_to_action_head: nn.Module | None = None + self._video_head_latent_channels: int | None = None self.num_inference_timesteps = config.num_inference_timesteps text_enc_path = ensure_file( @@ -313,10 +374,129 @@ def __init__( # self.num_timestep_buckets = config.num_timestep_buckets self.config = config self._noise_logged = False + self._loss_contract_logged = False + self._motion_mask_logged = False self.defer_lora_injection = config.defer_lora_injection + self._motion_mask_arm_ranges = self._load_motion_mask_arm_ranges(config) print("defer_lora_injection@@", self.defer_lora_injection) self.set_trainable_parameters(config.tune_projector, config.tune_diffusion_model) + @staticmethod + def _load_motion_mask_arm_ranges(config: WANPolicyHeadConfig) -> torch.Tensor | None: + """Load q99 normalization ranges for the 14 arm joints. + + The action tensor reaching this head is q99-normalized. For a q99 + transform, normalized_delta * (q99-q01) / 2 is the original action + delta, so the configured radian threshold remains meaningful without + changing the dataset or model inputs. + """ + if not config.motion_mask_enabled: + return None + + stats_path = config.motion_mask_stats_path + if not stats_path: + raise ValueError( + "motion_mask_stats_path is required when motion_mask_enabled=true" + ) + if not os.path.isfile(stats_path): + raise FileNotFoundError( + f"Motion-mask stats file does not exist: {stats_path}" + ) + + with open(stats_path, "r", encoding="utf-8") as f: + stats = json.load(f) + + ranges = [] + for key in ("left_joint_position", "right_joint_position"): + key_stats = stats.get(key) + if key_stats is None: + raise KeyError(f"Missing {key} in motion-mask stats: {stats_path}") + q01 = key_stats.get("q01") + q99 = key_stats.get("q99") + if q01 is None or q99 is None or len(q01) != 7 or len(q99) != 7: + raise ValueError( + f"Expected 7-dimensional q01/q99 stats for {key}: {stats_path}" + ) + ranges.extend(float(high) - float(low) for low, high in zip(q01, q99)) + + ranges_tensor = torch.tensor(ranges, dtype=torch.float32) + if not torch.all(ranges_tensor > 0): + raise ValueError( + f"Motion-mask q99 ranges must be positive, got {ranges_tensor.tolist()}" + ) + print( + "[MOTION MASK] enabled=true " + f"threshold_rad={config.motion_mask_threshold_rad:.4f} " + f"stats={stats_path}" + ) + return ranges_tensor + + def _build_motion_loss_mask(self, actions: torch.Tensor) -> torch.Tensor: + """Return a [B, T] mask for action timesteps involved in motion. + + Actions are normalized before reaching this head, so use the loaded + q99 ranges to evaluate the 14-arm-joint delta in radians. The action + stream can contain multiple 24-step chunks; deltas never cross a + chunk boundary. + """ + if self._motion_mask_arm_ranges is None: + raise RuntimeError("Motion-mask ranges were not initialized") + if actions.ndim != 3: + raise RuntimeError(f"Expected actions with shape [B,T,D], got {actions.shape}") + + batch_size, action_steps, _ = actions.shape + if action_steps % self.action_horizon != 0: + raise RuntimeError( + f"Motion mask requires action length divisible by horizon " + f"{self.action_horizon}, got {action_steps}" + ) + + arm_indices = torch.tensor( + [0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14], + device=actions.device, + ) + num_chunks = action_steps // self.action_horizon + arm_actions = actions.index_select(dim=2, index=arm_indices).reshape( + batch_size, num_chunks, self.action_horizon, 14 + ) + delta_normalized = arm_actions[:, :, 1:] - arm_actions[:, :, :-1] + ranges = self._motion_mask_arm_ranges.to( + device=actions.device, dtype=actions.dtype + ).reshape(1, 1, 1, 14) + delta_rad = delta_normalized * ranges / 2.0 + max_joint_move = delta_rad.abs().max(dim=-1).values + moving_transition = max_joint_move > float(self.config.motion_mask_threshold_rad) + + # A transition contributes to both endpoints: this preserves the + # target immediately before movement and the target at the new pose. + motion_mask = torch.zeros( + batch_size, + num_chunks, + self.action_horizon, + dtype=actions.dtype, + device=actions.device, + ) + motion_mask[:, :, :-1] = moving_transition.to(dtype=actions.dtype) + motion_mask[:, :, 1:] = torch.maximum( + motion_mask[:, :, 1:], moving_transition.to(dtype=actions.dtype) + ) + motion_mask = motion_mask.reshape(batch_size, action_steps) + + if not self._motion_mask_logged: + active_steps = motion_mask.sum().detach() + total_steps = torch.tensor( + motion_mask.numel(), dtype=motion_mask.dtype, device=motion_mask.device + ) + print( + "[MOTION MASK] action_steps=" + f"{action_steps} chunks={num_chunks} " + f"active_ratio={float(active_steps / total_steps.clamp_min(1.0)):.4f} " + f"active_steps={int(active_steps.item())}/{motion_mask.numel()}" + ) + self._motion_mask_logged = True + + return motion_mask + def reset_inference_cache(self) -> None: self.kv_cache1 = None self.kv_cache_neg = None @@ -344,7 +524,9 @@ def set_trainable_parameters(self, tune_projector: bool, tune_diffusion_model: b if not any(p.requires_grad for p in self.parameters()): print("Warning: No action head trainable parameters found.") - if self.train_architecture == "lora" and not self.defer_lora_injection: + if self.train_architecture == "action_only": + self.configure_action_only_training() + elif self.train_architecture == "lora" and not self.defer_lora_injection: print("Adding LoRA to model") for p in self.parameters(): p.requires_grad = False @@ -416,6 +598,53 @@ def inject_lora_after_loading(self): else: print("LoRA injection not needed (train_architecture != 'lora')") + def configure_action_only_training(self): + """Freeze the shared video DiT and tune only the three action adapters.""" + if self.train_architecture != "action_only": + raise RuntimeError( + "action_only_training requires train_architecture=action_only, " + f"got {self.train_architecture!r}" + ) + lora_parameters = [ + name for name, _ in self.model.named_parameters() + if "lora_" in name + ] + if lora_parameters: + raise RuntimeError( + "Action-adapter-only training forbids shared DiT LoRA; " + f"found {len(lora_parameters)} LoRA parameters" + ) + for parameter in self.parameters(): + parameter.requires_grad = False + self.model.state_encoder.requires_grad_(True) + self.model.action_encoder.requires_grad_(True) + self.model.action_decoder.requires_grad_(True) + self.text_encoder.requires_grad_(False) + self.image_encoder.requires_grad_(False) + self.vae.requires_grad_(False) + allowed_fragments = ( + "action_head.model.state_encoder.", + "action_head.model.action_encoder.", + "action_head.model.action_decoder.", + ) + unexpected = [ + name + for name, parameter in self.named_parameters() + if parameter.requires_grad + and not any(fragment in f"action_head.{name}" for fragment in allowed_fragments) + ] + if unexpected: + raise RuntimeError( + "Action-adapter trainable whitelist violation: " + + ", ".join(unexpected[:20]) + ) + print( + "[TRAINABLE WHITELIST] state_encoder, action_encoder, " + "action_decoder only" + ) + print("[SHARED DIT LORA] disabled") + self.print_trainable_params() + def set_frozen_modules_to_eval_mode(self): """ Huggingface will call model.train() at each training_step. To ensure @@ -761,17 +990,298 @@ def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> weighted_dynamics_loss = weight_dynamics.mean() if actions.numel() > 0: - action_loss_per_sample = torch.nn.functional.mse_loss( + action_loss_per_element = torch.nn.functional.mse_loss( action_noise_pred.float(), training_target_action.float(), reduction='none' - ) * action_mask # shape: [B, ...] - action_loss_per_sample = has_real_action[:, None].float() * action_loss_per_sample # apply has_real_action - weight_action = action_loss_per_sample.mean(dim=2) * self.scheduler.training_weight( + ) + valid_action_mask = action_mask.to( + dtype=action_loss_per_element.dtype, + device=action_loss_per_element.device, + ) + # Do not divide G2's 16 valid dimensions by max_action_dim=32. + # The previous masked-then-mean reduction silently halved the + # G2 action objective because its padded dimensions are zero. + valid_dim_count = valid_action_mask.sum(dim=2).clamp_min(1.0) + valid_dims = int(valid_dim_count.max().item()) + if self.config.action_only_training and valid_dims != 16: + raise RuntimeError( + "G2 action adapter requires exactly 16 valid action " + f"dimensions, got {valid_dims}" + ) + + arm_indices = torch.tensor( + [0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14], + device=action_loss_per_element.device, + ) + gripper_indices = torch.tensor( + [7, 15], + device=action_loss_per_element.device, + ) + arm_element_loss = action_loss_per_element.index_select( + dim=2, index=arm_indices + ) + arm_mask = valid_action_mask.index_select( + dim=2, index=arm_indices + ) + gripper_element_loss = action_loss_per_element.index_select( + dim=2, index=gripper_indices + ) + gripper_mask = valid_action_mask.index_select( + dim=2, index=gripper_indices + ) + arm_loss = ( + arm_element_loss * arm_mask + ).sum(dim=2) / arm_mask.sum(dim=2).clamp_min(1.0) + gripper_loss = ( + gripper_element_loss * gripper_mask + ).sum(dim=2) / gripper_mask.sum(dim=2).clamp_min(1.0) + action_loss_per_timestep = ( + (14.0 / 16.0) * arm_loss + + (2.0 / 16.0) * gripper_loss + ) + action_loss_per_timestep = ( + has_real_action[:, None].float() + * action_loss_per_timestep + ) + if not self._loss_contract_logged: + old_reduction = ( + action_loss_per_element * valid_action_mask + ).mean(dim=2).mean().detach() + corrected_reduction = ( + action_loss_per_timestep.mean().detach() + ) + ratio = corrected_reduction / old_reduction.clamp_min(1e-12) + print( + "[ACTION LOSS CONTRACT] padded_action_dim=" + f"{action_loss_per_element.shape[2]} " + f"valid_action_dim={valid_dims} " + f"corrected_loss/old_loss={float(ratio):.4f}" + ) + self._loss_contract_logged = True + weight_action = action_loss_per_timestep * self.scheduler.training_weight( timestep_action.flatten(0, 1), ).unflatten(0, (noise_action.shape[0], noise_action.shape[1])).to(self._device) - weighted_action_loss = weight_action.mean() - loss = weighted_dynamics_loss + weighted_action_loss + action_time_mask = has_real_action[:, None].float() + if self.config.motion_mask_enabled: + action_time_mask = action_time_mask * self._build_motion_loss_mask(actions) + weighted_action_loss = ( + weight_action * action_time_mask + ).sum() / action_time_mask.sum().clamp_min(1.0) + else: + weighted_action_loss = weight_action.mean() + + # The flow-matching objective supervises the velocity/noise + # field. Add an optional direct x0 objective so the decoded + # action trajectory and its final pose are also constrained. + action_x0_loss = torch.tensor(0.0, device=self._device) + action_endpoint_loss = torch.tensor(0.0, device=self._device) + action_start_loss = torch.tensor(0.0, device=self._device) + if ( + float(self.config.action_x0_loss_weight) > 0.0 + or float(self.config.action_endpoint_loss_weight) > 0.0 + or float(self.config.action_start_loss_weight) > 0.0 + ): + sigma_action = self.scheduler.sigmas.to( + device=noisy_actions.device, dtype=noisy_actions.dtype + )[timestep_action_id.to(device=noisy_actions.device)].unsqueeze(-1) + decoded_actions = ( + noisy_actions.float() + - sigma_action.float() * action_noise_pred.float() + ) + direct_action_element_loss = torch.nn.functional.smooth_l1_loss( + decoded_actions, + actions.float(), + beta=0.05, + reduction="none", + ) + direct_arm_loss = ( + direct_action_element_loss.index_select(dim=2, index=arm_indices) + * valid_action_mask.index_select(dim=2, index=arm_indices) + ).sum(dim=2) / valid_action_mask.index_select( + dim=2, index=arm_indices + ).sum(dim=2).clamp_min(1.0) + direct_gripper_loss = ( + direct_action_element_loss.index_select( + dim=2, index=gripper_indices + ) + * valid_action_mask.index_select(dim=2, index=gripper_indices) + ).sum(dim=2) / valid_action_mask.index_select( + dim=2, index=gripper_indices + ).sum(dim=2).clamp_min(1.0) + direct_action_per_timestep = ( + (14.0 / 16.0) * direct_arm_loss + + (2.0 / 16.0) * direct_gripper_loss + ) + direct_action_per_timestep = ( + has_real_action[:, None].float() + * direct_action_per_timestep + ) + action_x0_loss = ( + direct_action_per_timestep * action_time_mask + ).sum() / action_time_mask.sum().clamp_min(1.0) + + if float(self.config.action_endpoint_loss_weight) > 0.0: + if actions.shape[1] % self.action_horizon != 0: + raise RuntimeError( + "Endpoint action loss requires action length divisible " + f"by horizon {self.action_horizon}, got {actions.shape[1]}" + ) + num_chunks = actions.shape[1] // self.action_horizon + decoded_chunks = decoded_actions.reshape( + actions.shape[0], num_chunks, self.action_horizon, -1 + ) + target_chunks = actions.float().reshape( + actions.shape[0], num_chunks, self.action_horizon, -1 + ) + valid_chunks = valid_action_mask.reshape( + actions.shape[0], num_chunks, self.action_horizon, -1 + ) + endpoint_element_loss = torch.nn.functional.smooth_l1_loss( + decoded_chunks[:, :, -1], + target_chunks[:, :, -1], + beta=0.05, + reduction="none", + ) + endpoint_valid = valid_chunks[:, :, -1] + endpoint_arm_loss = ( + endpoint_element_loss.index_select(dim=2, index=arm_indices) + * endpoint_valid.index_select(dim=2, index=arm_indices) + ).sum(dim=2) / endpoint_valid.index_select( + dim=2, index=arm_indices + ).sum(dim=2).clamp_min(1.0) + endpoint_gripper_loss = ( + endpoint_element_loss.index_select( + dim=2, index=gripper_indices + ) + * endpoint_valid.index_select(dim=2, index=gripper_indices) + ).sum(dim=2) / endpoint_valid.index_select( + dim=2, index=gripper_indices + ).sum(dim=2).clamp_min(1.0) + endpoint_per_chunk = ( + (14.0 / 16.0) * endpoint_arm_loss + + (2.0 / 16.0) * endpoint_gripper_loss + ) + endpoint_chunk_mask = action_time_mask.reshape( + actions.shape[0], num_chunks, self.action_horizon + ).amax(dim=2) + action_endpoint_loss = ( + endpoint_per_chunk * endpoint_chunk_mask + ).sum() / endpoint_chunk_mask.sum().clamp_min(1.0) + + # Current-state start consistency (L_start): the FIRST predicted + # action should be close to the current joint state. state is a + # TRAINING TARGET here, not a model input, so this anchors the + # model's trajectory to the true start pose without re-enabling + # the state-input shortcut. Only the sequence-start state is + # available, so the anchor applies to the first action step. + if float(self.config.action_start_loss_weight) > 0.0: + # action is padded to action_dim (32), state to + # max_state_dim (64); compare only the valid G2 dims (16). + start_pred = decoded_actions[:, 0, :16] # (B, 16) + state_0 = state_features[:, 0, :16] # (B, 16) + start_element = torch.nn.functional.smooth_l1_loss( + start_pred.float(), state_0.float(), beta=0.05, reduction="none" + ) + start_valid = valid_action_mask[:, 0] + start_arm = ( + start_element.index_select(dim=1, index=arm_indices) + * start_valid.index_select(dim=1, index=arm_indices) + ).sum(dim=1) / start_valid.index_select( + dim=1, index=arm_indices + ).sum(dim=1).clamp_min(1.0) + start_grip = ( + start_element.index_select(dim=1, index=gripper_indices) + * start_valid.index_select(dim=1, index=gripper_indices) + ).sum(dim=1) / start_valid.index_select( + dim=1, index=gripper_indices + ).sum(dim=1).clamp_min(1.0) + start_per_sample = ( + (14.0 / 16.0) * start_arm + (2.0 / 16.0) * start_grip + ) + action_start_loss = ( + start_per_sample * has_real_action.float() + ).sum() / has_real_action.sum().clamp_min(1.0) + + # Video->action consistency (L_video_to_action): recover the + # action from the PREDICTED clean video latent so the action + # chunk depends on the learned implicit dynamics encoded in the + # video. Implemented defensively: if the latent/head shapes are + # not as expected the term is skipped (logged) instead of + # crashing the run. + action_video_consistency_loss = torch.tensor(0.0, device=self._device) + if float(self.config.action_video_consistency_loss_weight) > 0.0 and actions.numel() > 0: + try: + sigma_video = self.scheduler.sigmas.to( + device=noisy_latents.device, dtype=noisy_latents.dtype + )[timestep_id.to(device=noisy_latents.device)] + while sigma_video.ndim < noisy_latents.ndim: + sigma_video = sigma_video.unsqueeze(-1) + decoded_latents = ( + noisy_latents.float() - sigma_video.float() * video_noise_pred.float() + ) + # pool spatial + time, keep channel dim (dim 1) + pooled = decoded_latents.mean(dim=tuple(range(2, decoded_latents.ndim))) + if self.video_to_action_head is None: + self._video_head_latent_channels = int(pooled.shape[1]) + self.video_to_action_head = nn.Linear( + self._video_head_latent_channels, self.action_dim + ).to(device=pooled.device, dtype=pooled.dtype) + pred_action = self.video_to_action_head(pooled.float()) + vid_element = torch.nn.functional.mse_loss( + pred_action, actions[:, 0].float(), reduction="none" + ) + vid_valid = valid_action_mask[:, 0] + vid_arm = ( + vid_element.index_select(dim=1, index=arm_indices) + * vid_valid.index_select(dim=1, index=arm_indices) + ).sum(dim=1) / vid_valid.index_select( + dim=1, index=arm_indices + ).sum(dim=1).clamp_min(1.0) + vid_grip = ( + vid_element.index_select(dim=1, index=gripper_indices) + * vid_valid.index_select(dim=1, index=gripper_indices) + ).sum(dim=1) / vid_valid.index_select( + dim=1, index=gripper_indices + ).sum(dim=1).clamp_min(1.0) + vid_per_sample = ( + (14.0 / 16.0) * vid_arm + (2.0 / 16.0) * vid_grip + ) + action_video_consistency_loss = ( + vid_per_sample * has_real_action.float() + ).sum() / has_real_action.sum().clamp_min(1.0) + except Exception as exc: # noqa: BLE001 - defensive skip + logging.warning( + "video->action consistency loss skipped (%s); " + "check latent/head shapes.", + exc, + ) + + action_objective = ( + float(self.config.action_loss_weight) * weighted_action_loss + + float(self.config.action_x0_loss_weight) * action_x0_loss + + float(self.config.action_endpoint_loss_weight) + * action_endpoint_loss + + float(self.config.action_start_loss_weight) * action_start_loss + + float(self.config.action_video_consistency_loss_weight) + * action_video_consistency_loss + ) + if self.config.action_only_training: + loss = action_objective + weighted_dynamics_loss = weighted_dynamics_loss.detach() + else: + loss = ( + float(self.config.dynamics_loss_weight) + * weighted_dynamics_loss + + action_objective + ) else: weighted_action_loss = torch.tensor(0.0, device=self._device) + action_x0_loss = torch.tensor(0.0, device=self._device) + action_endpoint_loss = torch.tensor(0.0, device=self._device) + if self.config.action_only_training: + raise RuntimeError( + "Action-adapter-only training received an empty action " + "tensor; refusing to optimize video loss" + ) loss = weighted_dynamics_loss # loss = dynamics_loss_per_sample.mean() @@ -780,6 +1290,10 @@ def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> "loss": loss, "dynamics_loss": weighted_dynamics_loss, "action_loss": weighted_action_loss, + "action_x0_loss": action_x0_loss, + "action_endpoint_loss": action_endpoint_loss, + "action_start_loss": action_start_loss, + "action_video_consistency_loss": action_video_consistency_loss, } return BatchFeature(data=output_dict) diff --git a/groot/vla/model/dreamzero/base_vla.py b/groot/vla/model/dreamzero/base_vla.py index d81401a4..f89dd185 100644 --- a/groot/vla/model/dreamzero/base_vla.py +++ b/groot/vla/model/dreamzero/base_vla.py @@ -102,6 +102,106 @@ def validate_inputs(self, inputs): if detected_error: raise ValueError(error_msg) + @classmethod + def load_action_adapter( + cls, + adapter_model_path: str, + pretrained_base_model_path: str, + config: VLAConfig | None = None, + ): + """Load a clean DreamZero base and then a strict G2 action adapter.""" + import json + import os + from safetensors.torch import load_file + + if config is None: + config = cls.config_class.from_pretrained(adapter_model_path) + model = cls(config) + model_state = model.state_dict() + adapter_prefixes = ( + "action_head.model.state_encoder.", + "action_head.model.action_encoder.", + "action_head.model.action_decoder.", + ) + + index_path = os.path.join( + pretrained_base_model_path, + "model.safetensors.index.json", + ) + single_path = os.path.join( + pretrained_base_model_path, + "model.safetensors", + ) + if os.path.exists(index_path): + with open(index_path) as stream: + index = json.load(stream) + base_files = [ + os.path.join(pretrained_base_model_path, filename) + for filename in sorted(set(index["weight_map"].values())) + ] + elif os.path.exists(single_path): + base_files = [single_path] + else: + raise FileNotFoundError( + f"No DreamZero base weights at {pretrained_base_model_path}" + ) + + loaded_base_keys: set[str] = set() + skipped_base_keys: set[str] = set() + for base_file in base_files: + shard = load_file(base_file) + compatible = {} + for key, value in shard.items(): + if key in model_state and model_state[key].shape == value.shape: + compatible[key] = value + loaded_base_keys.add(key) + else: + skipped_base_keys.add(key) + model.load_state_dict(compatible, strict=False) + + required_shared_keys = { + key + for key in model_state + if not key.startswith(adapter_prefixes) + } + missing_shared = sorted(required_shared_keys - loaded_base_keys) + if missing_shared: + raise RuntimeError( + "Incomplete shared DreamZero base load; missing shared keys: " + + ", ".join(missing_shared[:20]) + ) + print( + "[BASE LOAD] DreamZero-AgiBot loaded " + f"shared_keys={len(required_shared_keys)} " + f"skipped_incompatible={len(skipped_base_keys)}" + ) + print("[EMBODIMENT] g2 (logical id), action-adapter local slot=0") + + adapter_path = os.path.join( + adapter_model_path, + "action_expert.safetensors", + ) + adapter_state = load_file(adapter_path) + expected_adapter_keys = { + key for key in model_state if key.startswith(adapter_prefixes) + } + if set(adapter_state) != expected_adapter_keys: + missing = sorted(expected_adapter_keys - set(adapter_state)) + unexpected = sorted(set(adapter_state) - expected_adapter_keys) + raise RuntimeError( + "Action adapter key contract mismatch: " + f"missing={missing[:20]} unexpected={unexpected[:20]}" + ) + model.load_state_dict(adapter_state, strict=False) + for module_name in ( + "state_encoder", + "action_encoder", + "action_decoder", + ): + print(f"[ACTION ADAPTER] {module_name} loaded") + print("[SHARED DIT LORA] disabled") + return model + def validate_data(self, action_head_outputs, backbone_outputs, is_training): fail_backbone = ( @@ -336,78 +436,170 @@ def from_pretrained_for_tuning( @classmethod def load_lora( - cls, - pretrained_model_name_or_path: str - ): + cls, + pretrained_model_name_or_path: str, + pretrained_base_model_path: str | None = None, + ): + """Load a LoRA checkpoint on the exact full-model training base. + + A LoRA delta is not a standalone model. G2 training first instantiated + the checkpoint architecture, loaded ``DreamZero-AgiBot``, and only then + injected LoRA adapters. Inference must mirror that order. Loading the + raw Wan component and applying the LoRA delta to it produces a + numerically valid but semantically wrong model. + """ from safetensors.torch import load_file import os import json - print("loading lora@@@@@") + import gc + + if pretrained_base_model_path is None: + raise ValueError( + "A LoRA-only checkpoint requires pretrained_base_model_path. " + "Use the same full checkpoint recorded as " + "pretrained_model_path in experiment_cfg/conf.yaml." + ) + if not os.path.isdir(pretrained_base_model_path): + raise FileNotFoundError( + f"LoRA base model directory does not exist: " + f"{pretrained_base_model_path}" + ) + + print( + "Loading LoRA checkpoint " + f"{pretrained_model_name_or_path} on base " + f"{pretrained_base_model_path}" + ) - # Check for different checkpoint formats safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") - - state_dict = {} + + lora_state_dict = {} if os.path.exists(safetensors_index_path): - # Handle sharded safetensors - print(f"Loading sharded safetensors using index: {safetensors_index_path}") - with open(safetensors_index_path, 'r') as f: index = json.load(f) - - # Load each shard - for shard_file in set(index["weight_map"].values()): + for shard_file in sorted(set(index["weight_map"].values())): shard_path = os.path.join(pretrained_model_name_or_path, shard_file) - print(f"Loading shard: {shard_path}") - shard_state_dict = load_file(shard_path) - state_dict.update(shard_state_dict) - + lora_state_dict.update(load_file(shard_path)) elif os.path.exists(safetensors_path): - # Handle single safetensors file - print(f"Loading weights from safetensors: {safetensors_path}") - state_dict.update(load_file(safetensors_path)) - - # Load config - print("loading config@@") + lora_state_dict.update(load_file(safetensors_path)) + else: + raise FileNotFoundError( + f"No LoRA weights found at {pretrained_model_name_or_path}" + ) + config_path = os.path.join(pretrained_model_name_or_path, "config.json") with open(config_path, "r") as f: config_dict = json.load(f) config = VLAConfig(**config_dict) - print("loading model") - # Disable defer_lora_injection so LoRA layers are created during init, - # matching the PEFT key hierarchy (base_model.model.*) in the checkpoint. + # Mirror training: build the target G2 architecture without loading raw + # Wan component weights, load the full DreamZero base, then inject LoRA. ah_cfg = config.action_head_cfg inner = ah_cfg.get('config', ah_cfg) if isinstance(ah_cfg.get('config'), dict) else ah_cfg - if 'defer_lora_injection' in inner: - inner['defer_lora_injection'] = False - print("defer_lora_injection disabled for load_lora") - # Enable component loading so DiT base weights are loaded from pretrained - if 'skip_component_loading' in inner: - inner['skip_component_loading'] = False - print("skip_component_loading disabled for load_lora") - - # Instantiate model (LoRA layers now exist from init) + inner['defer_lora_injection'] = True + inner['skip_component_loading'] = True model = cls(config) - # Remove .base_layer from keys if present - has_base_layer = any(".base_layer." in key for key in state_dict.keys()) + base_index_path = os.path.join( + pretrained_base_model_path, + "model.safetensors.index.json", + ) + base_single_path = os.path.join( + pretrained_base_model_path, + "model.safetensors", + ) + loaded_base_keys: set[str] = set() + unexpected_base_keys: set[str] = set() + + def load_base_state_dict(base_state_dict: dict) -> None: + model_keys = set(model.state_dict()) + loaded_base_keys.update(set(base_state_dict) & model_keys) + _, unexpected = model.load_state_dict( + base_state_dict, + strict=False, + ) + unexpected_base_keys.update(unexpected) + + if os.path.exists(base_index_path): + with open(base_index_path, "r") as f: + base_index = json.load(f) + for shard_file in sorted(set(base_index["weight_map"].values())): + shard_path = os.path.join( + pretrained_base_model_path, + shard_file, + ) + print(f"Loading DreamZero base shard: {shard_path}") + base_state_dict = load_file(shard_path) + load_base_state_dict(base_state_dict) + del base_state_dict + gc.collect() + elif os.path.exists(base_single_path): + base_state_dict = load_file(base_single_path) + load_base_state_dict(base_state_dict) + del base_state_dict + gc.collect() + else: + raise FileNotFoundError( + "No full base weights found at " + f"{pretrained_base_model_path}" + ) + + if not loaded_base_keys: + raise RuntimeError( + "The full DreamZero base checkpoint did not match any " + "parameters in the G2 model architecture" + ) + if unexpected_base_keys: + print( + "Ignoring base-only keys that are not present in the G2 " + f"architecture: count={len(unexpected_base_keys)}" + ) + + if not ( + hasattr(model, "action_head") + and hasattr(model.action_head, "inject_lora_after_loading") + ): + raise RuntimeError( + "G2 action head does not support deferred LoRA injection" + ) + model.action_head.inject_lora_after_loading() + + has_base_layer = any( + ".base_layer." in key for key in lora_state_dict + ) if has_base_layer: print("Removing '.base_layer' from state dict keys") - state_dict = {k.replace(".base_layer.", "."): v for k, v in state_dict.items()} + lora_state_dict = { + key.replace(".base_layer.", "."): value + for key, value in lora_state_dict.items() + } + + model_keys = set(model.state_dict()) + unknown_lora_keys = set(lora_state_dict) - model_keys + if unknown_lora_keys: + sample = sorted(unknown_lora_keys)[:20] + raise RuntimeError( + "LoRA checkpoint keys do not match the model constructed on " + f"the DreamZero base: count={len(unknown_lora_keys)}, " + f"sample={sample}" + ) - # Load weights - missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) - - if missing_keys: - print(f"Missing keys when loading pretrained weights: {missing_keys}") + missing_keys, unexpected_keys = model.load_state_dict( + lora_state_dict, + strict=False, + ) if unexpected_keys: - print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") - - print("Successfully loaded pretrained weights") + raise RuntimeError( + f"Unexpected LoRA keys after validation: {unexpected_keys}" + ) - print(f"{cls}\n") + print( + "Successfully loaded full DreamZero base and LoRA delta | " + f"base_keys={len(loaded_base_keys)} " + f"lora_keys={len(lora_state_dict)} " + f"model_only_keys={len(missing_keys)}" + ) return model def load_lora_weight(self, pretrained_model_name_or_path: str): diff --git a/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py b/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py index 02009be9..78def9b4 100644 --- a/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py +++ b/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py @@ -197,7 +197,8 @@ def __init__(self, qk_norm=True, eps=1e-6, num_action_per_block=32, - num_state_per_block=1): + num_state_per_block=1, + cut_state_attention=True): assert dim % num_heads == 0 super().__init__() self.dim = dim @@ -212,6 +213,7 @@ def __init__(self, self.frame_seqlen = frame_seqlen self.num_action_per_block = num_action_per_block self.num_state_per_block = num_state_per_block + self.cut_state_attention = cut_state_attention # layers self.q = nn.Linear(dim, dim) self.k = nn.Linear(dim, dim) @@ -479,7 +481,11 @@ def _blockwise_causal_flash_attn(self, q, k, v, frame_seqlen, num_frame_per_bloc state_block_starts = [state_start + i * num_state_per_block for i in range(num_state_blocks)] state_block_ends = [state_start + (i + 1) * num_state_per_block for i in range(num_state_blocks)] - # Process each image block + # Process each image block. + # Context: first image + image blocks + current action block (joint + # video<->action dynamics). When cut_state_attention=True (default) state + # is NOT in the context: it only conditions the shared blocks via the e0 + # modulation. When False, state joins the context. for block_idx in range(num_image_blocks): block_start = image_block_starts[block_idx] block_end = image_block_ends[block_idx] @@ -488,53 +494,61 @@ def _blockwise_causal_flash_attn(self, q, k, v, frame_seqlen, num_frame_per_bloc action_block_end = action_block_ends[block_idx] state_block_start = state_block_starts[block_idx] state_block_end = state_block_ends[block_idx] - - # Build context: first image + relevant image blocks + current action + current state - k_context = torch.cat([ + + k_parts = [ k[:, first_image_start:first_image_end], # First image k[:, image_kv_start:block_end], # Image blocks k[:, action_block_start:action_block_end], # Current action block - k[:, state_block_start:state_block_end] # Current state block - ], dim=1) - v_context = torch.cat([ + ] + v_parts = [ v[:, first_image_start:first_image_end], v[:, image_kv_start:block_end], v[:, action_block_start:action_block_end], - v[:, state_block_start:state_block_end] - ], dim=1) - + ] + if not self.cut_state_attention: + k_parts.append(k[:, state_block_start:state_block_end]) + v_parts.append(v[:, state_block_start:state_block_end]) + k_context = torch.cat(k_parts, dim=1) + v_context = torch.cat(v_parts, dim=1) + output[:, block_start:block_end] = self.attn( q[:, block_start:block_end], k_context, v_context ) - - # Process each action block + + # Process each action block. + # Context: first image + image blocks + current action block. When + # cut_state_attention=True (default) state is NOT in the action context — + # the action chunk must be driven by the future video, not by a direct + # current-state copy. When False, state joins the context. for block_idx in range(num_action_blocks): action_block_start = action_block_starts[block_idx] action_block_end = action_block_ends[block_idx] image_block_end = image_block_ends[block_idx] state_block_start = state_block_starts[block_idx] state_block_end = state_block_ends[block_idx] - + # Determine image context range if self.local_attn_size != -1: image_kv_start = max(image_blocks_start, image_block_end - self.local_attn_size * frame_seqlen) else: image_kv_start = image_blocks_start - - # Build context - k_context = torch.cat([ + + k_parts = [ k[:, first_image_start:first_image_end], # First image k[:, image_kv_start:image_block_end], # Image blocks k[:, action_block_start:action_block_end], # Current action block - k[:, state_block_start:state_block_end] # Current state block - ], dim=1) - v_context = torch.cat([ + ] + v_parts = [ v[:, first_image_start:first_image_end], v[:, image_kv_start:image_block_end], v[:, action_block_start:action_block_end], - v[:, state_block_start:state_block_end] - ], dim=1) - + ] + if not self.cut_state_attention: + k_parts.append(k[:, state_block_start:state_block_end]) + v_parts.append(v[:, state_block_start:state_block_end]) + k_context = torch.cat(k_parts, dim=1) + v_context = torch.cat(v_parts, dim=1) + output[:, action_block_start:action_block_end] = self.attn( q[:, action_block_start:action_block_end], k_context, v_context ) @@ -693,8 +707,13 @@ def _process_noisy_image_blocks(self, noisy_image_q, noisy_image_k, noisy_image_ action_block_ends = [start + self.num_action_per_block for start in action_block_starts] state_block_starts = [i * self.num_state_per_block for i in range(num_blocks)] state_block_ends = [start + self.num_state_per_block for start in state_block_starts] - - # Process noisy image blocks + + # Process noisy image blocks. + # Context: first_clean_frame + clean_blocks[0:i] + current_noisy_block + action[i]. + # When cut_state_attention=True (default), state is NOT part of the video + # context: the video branch is its own dynamics model conditioned on the + # shared blocks (state enters via e0), and state must not leak into the + # future-video prediction either. When False, state joins the context. for block_idx in range(num_blocks): noisy_start = noisy_block_starts[block_idx] noisy_end = noisy_block_ends[block_idx] @@ -703,23 +722,25 @@ def _process_noisy_image_blocks(self, noisy_image_q, noisy_image_k, noisy_image_ action_end = action_block_ends[block_idx] state_start = state_block_starts[block_idx] state_end = state_block_ends[block_idx] - + q_block = noisy_image_q[:, noisy_start:noisy_end] - - # Build context: first_clean_frame + clean_blocks[0:i] + current_noisy_block + action[i] + state[i] - k_context = torch.cat([ + + k_parts = [ clean_image_k[:, :clean_end], noisy_image_k[:, noisy_start:noisy_end], noisy_action_k[:, action_start:action_end], - noisy_state_k[:, state_start:state_end] - ], dim=1) - v_context = torch.cat([ + ] + v_parts = [ clean_image_v[:, :clean_end], noisy_image_v[:, noisy_start:noisy_end], noisy_action_v[:, action_start:action_end], - noisy_state_v[:, state_start:state_end] - ], dim=1) - + ] + if not self.cut_state_attention: + k_parts.append(noisy_state_k[:, state_start:state_end]) + v_parts.append(noisy_state_v[:, state_start:state_end]) + k_context = torch.cat(k_parts, dim=1) + v_context = torch.cat(v_parts, dim=1) + output[:, noisy_start:noisy_end] = self.attn(q_block, k_context, v_context) return output @@ -752,8 +773,13 @@ def _process_noisy_action_blocks(self, noisy_action_q, noisy_action_k, noisy_act noisy_image_block_ends = [start + self.frame_seqlen * self.num_frame_per_block for start in noisy_image_block_starts] state_block_starts = [i * self.num_state_per_block for i in range(num_blocks)] state_block_ends = [start + self.num_state_per_block for start in state_block_starts] - - # Process noisy action blocks + + # Process noisy action blocks. + # Context: first_clean_frame + clean_blocks[0:i] + noisy_image[i] + action[i]. + # When cut_state_attention=True (default), state is NOT part of the action + # context: the action chunk must derive "what happens next" from the future + # video, not from a raw current-state copy. State only conditions the shared + # blocks via e0. When False, the original state-in-context behavior returns. for block_idx in range(num_blocks): action_start = action_block_starts[block_idx] action_end = action_block_ends[block_idx] @@ -762,23 +788,25 @@ def _process_noisy_action_blocks(self, noisy_action_q, noisy_action_k, noisy_act noisy_img_end = noisy_image_block_ends[block_idx] state_start = state_block_starts[block_idx] state_end = state_block_ends[block_idx] - + q_block = noisy_action_q[:, action_start:action_end] - - # Build context: first_clean_frame + clean_blocks[0:i] + noisy_image[i] + action[i] + state[i] - k_context = torch.cat([ + + k_parts = [ clean_image_k[:, :clean_end], noisy_image_k[:, noisy_img_start:noisy_img_end], noisy_action_k[:, action_start:action_end], - noisy_state_k[:, state_start:state_end] - ], dim=1) - v_context = torch.cat([ + ] + v_parts = [ clean_image_v[:, :clean_end], noisy_image_v[:, noisy_img_start:noisy_img_end], noisy_action_v[:, action_start:action_end], - noisy_state_v[:, state_start:state_end] - ], dim=1) - + ] + if not self.cut_state_attention: + k_parts.append(noisy_state_k[:, state_start:state_end]) + v_parts.append(noisy_state_v[:, state_start:state_end]) + k_context = torch.cat(k_parts, dim=1) + v_context = torch.cat(v_parts, dim=1) + output[:, action_start:action_end] = self.attn(q_block, k_context, v_context) return output @@ -1099,7 +1127,8 @@ def __init__(self, cross_attn_norm=False, eps=1e-6, num_action_per_block=32, - num_state_per_block=1): + num_state_per_block=1, + cut_state_attention=True): super().__init__() self.dim = dim self.ffn_dim = ffn_dim @@ -1122,6 +1151,7 @@ def __init__(self, eps=eps, num_action_per_block=num_action_per_block, num_state_per_block=num_state_per_block, + cut_state_attention=cut_state_attention, ) self.norm3 = WanLayerNorm( dim, eps, @@ -1253,7 +1283,9 @@ def __init__(self, hidden_size=1024, diffusion_model_pretrained_path=None, num_action_per_block=32, - num_state_per_block=1): + num_state_per_block=1, + state_dropout=0.0, + cut_state_attention=True): r""" Initialize the diffusion model backbone. @@ -1342,6 +1374,24 @@ def __init__(self, output_dim=action_dim, ) + # State conditioning (AdaLN-style). The current proprioceptive state is + # encoded once into a moderate bottleneck and added to the per-block + # condition embedding `e0`. This symmetrically conditions BOTH video and + # action tokens through the shared DiT blocks. State never appears as + # register tokens that action can attend to, and there is no additive + # state->action path in the output heads. + self.state_cond_proj = nn.Sequential( + nn.Linear(max_state_dim, 256), + nn.SiLU(), + nn.Linear(256, dim * 6), + ) + self.state_dropout = state_dropout + # When True, action/video attention excludes the state register segment + # (state only conditions via the e0 AdaLN modulation). The state token + # structure/shape is preserved; the register state slots are zeroed so + # they cannot act as an attention shortcut. + self.cut_state_attention = cut_state_attention + # embeddings self.patch_embedding = nn.Conv3d( in_dim, dim, kernel_size=patch_size, stride=patch_size) @@ -1359,7 +1409,8 @@ def __init__(self, self.blocks = nn.ModuleList([ CausalWanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads, frame_seqlen, self.local_attn_size, sink_size, num_frame_per_block, qk_norm, cross_attn_norm, eps, - num_action_per_block, num_state_per_block) + num_action_per_block, num_state_per_block, + cut_state_attention=self.cut_state_attention) for _ in range(num_layers) ]) @@ -1383,6 +1434,12 @@ def __init__(self, # initialize weights self.init_weights() + # Zero-init the state conditioning's output projection so the initial + # AdaLN state modulation is ~0 (training starts video-driven; the state + # conditioning grows only if it helps). + nn.init.zeros_(self.state_cond_proj[-1].weight) + nn.init.zeros_(self.state_cond_proj[-1].bias) + self.gradient_checkpointing = True self.independent_first_frame = False if self.num_frame_per_block == 1 else True @@ -1393,7 +1450,8 @@ def _set_gradient_checkpointing(self, module, value=False): @staticmethod def _prepare_blockwise_causal_attn_mask( device: torch.device | str, num_frames: int = 21, - frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1, action_horizon=1, state_horizon=1, num_action_per_block=30, num_state_per_block=1 + frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1, action_horizon=1, state_horizon=1, num_action_per_block=30, num_state_per_block=1, + cut_state_attention: bool = True ) -> BlockMask: """ We will divide the token sequence into the following format: @@ -1498,16 +1556,25 @@ def attention_mask(b, h, q_idx, kv_idx): image_to_first = q_is_image_block & kv_is_first_image # Image block to first image: always allowed image_to_image = q_is_image_block & kv_is_image_block & (kv_block <= q_block) # Image block to image block: can attend to current and previous image blocks image_to_action = q_is_image_block & kv_is_action & (kv_block == q_block) # Image block to action: can attend to current action block - image_to_state = q_is_image_block & kv_is_state & (kv_block == q_block) # Image block to state: can attend to current state block - + # When cut_state_attention=True, the state edges are removed: neither + # video nor action may read the state register directly; state only + # conditions the shared blocks via the e0 AdaLN modulation. + image_to_state = ( + (not cut_state_attention) + and q_is_image_block and kv_is_state and (kv_block == q_block) + ) + image_block_mask = image_to_first | image_to_image | image_to_action | image_to_state - + # Action query action_to_image = q_is_action & kv_is_image_block & (kv_block <= q_block) # Action to image block: can attend to current and all previous image blocks action_to_action = q_is_action & kv_is_action & (kv_block == q_block) # Action to action: only same block - action_to_state = q_is_action & kv_is_state & (kv_block == q_block) # Action to state: only same block + action_to_state = ( + (not cut_state_attention) + and q_is_action and kv_is_state and (kv_block == q_block) + ) # Action to state: only same block (disabled when cut_state_attention) action_to_first = q_is_action & kv_is_first_image # Action to first image: always allowed - + action_mask = action_to_image | action_to_action | action_to_state | action_to_first # State query (conditioning) - cannot attend to anything @@ -1714,7 +1781,14 @@ def _forward_blocks( if action is not None: embodiment_id = torch.tensor([0], device=x.device).repeat(x.shape[0]) action_features = self.action_encoder(action, timestep_action, embodiment_id) - state_features = self.state_encoder(state, embodiment_id) + # State no longer enters the DiT sequence as register tokens that + # video/action could attend to (that was the shortcut). The register + # state slots are zeroed; real state influence is injected as an + # AdaLN-style modulation on `e0` below. + state_features = torch.zeros( + (state.shape[0], state.shape[1], self.dim), + device=action_features.device, dtype=action_features.dtype, + ) action_register = torch.cat([action_features, state_features], dim=1) action_length = action_features.shape[1] action_register_length = action_register.shape[1] @@ -1741,6 +1815,11 @@ def _forward_blocks( e0 = self.time_projection(e) e0 = e0.unflatten(dim=2, sizes=(6, self.dim)) + # Proprioceptive state is intentionally NOT used in this model variant: + # the register state slots are zeroed, attention edges are cut + # (cut_state_attention), and no AdaLN state modulation / dropout is + # applied. The action chunk must be predicted purely from the video. + # context context = self.text_embedding(context) @@ -2011,7 +2090,14 @@ def _forward_train( embodiment_id = torch.tensor([0]).repeat(x.shape[0]).to(device=embodiment_id.device) action_features = self.action_encoder(action, timestep_action, embodiment_id) action_length = action_features.shape[1] - state_features = self.state_encoder(state, embodiment_id) + # State no longer enters the DiT sequence as register tokens that + # video/action could attend to (that was the shortcut). The register + # state slots are zeroed; real state influence is injected as an + # AdaLN-style modulation on `e0` below. + state_features = torch.zeros( + (state.shape[0], state.shape[1], self.dim), + device=action_features.device, dtype=action_features.dtype, + ) action_register = torch.cat([action_features, state_features], dim=1) action_register_length = action_register.shape[1] x = torch.cat([x, action_register], dim=1) @@ -2039,6 +2125,11 @@ def _forward_train( e0 = self.time_projection(e) e0 = e0.unflatten(dim=2, sizes=(6, self.dim)) + # Proprioceptive state is intentionally NOT used in this model variant: + # the register state slots are zeroed, attention edges are cut + # (cut_state_attention), and no AdaLN state modulation / dropout is + # applied. The action chunk must be predicted purely from the video. + # context assert context.shape[1] == self.text_len context = self.text_embedding(context) diff --git a/groot/vla/model/dreamzero/transform/dreamzero_cotrain.py b/groot/vla/model/dreamzero/transform/dreamzero_cotrain.py index 0e303b29..b1d6ef16 100644 --- a/groot/vla/model/dreamzero/transform/dreamzero_cotrain.py +++ b/groot/vla/model/dreamzero/transform/dreamzero_cotrain.py @@ -101,7 +101,10 @@ def collate(features: List[dict], tokenizer: AutoTokenizer, num_views=3, embodim # If it's already a scalar (string, float, int, etc.), convert to string processed_item = str(parsed_item) - if num_views > 1 and elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.AGIBOT.value]: + if num_views > 1 and elem["embodiment_id"] in ( + embodiment_tag_mapping.get(EmbodimentTag.AGIBOT.value), + embodiment_tag_mapping.get(EmbodimentTag.G2.value), + ): processed_item = "A multi-view video shows that a robot " + processed_item.lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right hand, the bottom-left view shows the camera view from the left hand, and the bottom-right view is a black screen (inactive view). The robot " + processed_item.lower() elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.OXE_DROID.value]: processed_item = ( @@ -123,7 +126,10 @@ def collate(features: List[dict], tokenizer: AutoTokenizer, num_views=3, embodim output_values.append(processed_item) except (ValueError, SyntaxError, TypeError): # If parsing fails or item is already a string, use it directly - if num_views > 1 and elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.AGIBOT.value]: + if num_views > 1 and elem["embodiment_id"] in ( + embodiment_tag_mapping.get(EmbodimentTag.AGIBOT.value), + embodiment_tag_mapping.get(EmbodimentTag.G2.value), + ): item = "A multi-view video shows that a robot " + str(item).lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right hand, the bottom-left view shows the camera view from the left hand, and the bottom-right view is a black screen (inactive view). The robot " + str(item).lower() elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.OXE_DROID.value]: item = ( diff --git a/groot/vla/model/n1_5/sim_policy.py b/groot/vla/model/n1_5/sim_policy.py old mode 100755 new mode 100644 index 77ff9296..3adf73a7 --- a/groot/vla/model/n1_5/sim_policy.py +++ b/groot/vla/model/n1_5/sim_policy.py @@ -254,6 +254,7 @@ def __init__( model_target = train_cfg.model._target_ self.model_target = model_target + action_adapter_path = model_dir / "action_expert.safetensors" if model_config_overrides is not None and len(model_config_overrides) != 0: print(f"Applying model config overrides: {model_config_overrides}") @@ -277,9 +278,24 @@ def __init__( model_config = model_config_class.from_dict(model_config) # Instantiate the model - if hasattr(train_cfg, "save_lora_only") and train_cfg.save_lora_only is True: + if action_adapter_path.exists(): + print("Loading G2 action adapter on clean DreamZero base") + base_model_path = train_cfg.get("pretrained_model_path", None) + model = model_class.load_action_adapter( + str(model_dir), + pretrained_base_model_path=base_model_path, + config=model_config, + ) + elif hasattr(train_cfg, "save_lora_only") and train_cfg.save_lora_only is True: print(f"Loading LoRA weights from pretrained") - model = model_class.load_lora(model_path) + base_model_path = train_cfg.get( + "pretrained_model_path", + None, + ) + model = model_class.load_lora( + model_path, + pretrained_base_model_path=base_model_path, + ) else: print(f"Loading model from pretrained directly") model = model_class.from_pretrained(model_path, config=model_config) @@ -289,10 +305,25 @@ def __init__( cls_module, cls_name = model_target.rsplit(".", 1) if 'lora' in cls_name: cls_module, cls_name = cls_module.rsplit(".", 1) - if hasattr(train_cfg, "save_lora_only") and train_cfg.save_lora_only is True: + if action_adapter_path.exists(): + print("Loading G2 action adapter on clean DreamZero base") + cls = getattr(importlib.import_module(cls_module), cls_name) + base_model_path = train_cfg.get("pretrained_model_path", None) + model = cls.load_action_adapter( + str(model_dir), + pretrained_base_model_path=base_model_path, + ) + elif hasattr(train_cfg, "save_lora_only") and train_cfg.save_lora_only is True: print(f"Loading LoRA weights from pretrained") cls = getattr(importlib.import_module(cls_module), cls_name) - model = cls.load_lora(model_path) + base_model_path = train_cfg.get( + "pretrained_model_path", + None, + ) + model = cls.load_lora( + model_path, + pretrained_base_model_path=base_model_path, + ) else: print(f"Loading model from pretrained directly") cls = getattr(importlib.import_module(cls_module), cls_name) diff --git a/robot_live_client_g2.py b/robot_live_client_g2.py new file mode 100644 index 00000000..a8d4d6f7 --- /dev/null +++ b/robot_live_client_g2.py @@ -0,0 +1,677 @@ +#!/usr/bin/env python3 +"""DreamZero live-policy client for AgiBot G2. + +The policy/inference loop is shared with :mod:`robot_live_client`. This file +adapts the current ``agibot_gdk`` G2 API to the small camera/robot interface +used by that loop. Run after sourcing ``~/.cache/agibot/app/env.sh``. + +G2 has a three-axis head while the trained policy has two head dimensions. +They are mapped to yaw (idx11) and pitch (idx13); roll (idx12) is preserved. +G2's five-motor parallel waist needs an inverse-kinematics API which is not +exposed by the supplied examples, so waist commands are deliberately ignored. + +python robot_live_client_g2.py \ + --host 111.0.22.33 \ + --port 30001 \ + --prompt "机器人右臂先从网口扩展坞物料区抓取网口扩展坞，左臂随后从网线物料区依次抓取两根网线并逐一插入" \ + --sdk-arm-order left_right \ + --observation-fps 5 \ + --observation-history 1 \ + --image-transport jpeg \ + --arm-execution-mode direct-48 \ + --direct-control-hz 15 \ + --arm-delta-limit 0.04 \ + --arm-velocity-limit 0.35 \ + --arm-acceleration-limit 0.80 \ + --arm-close-timeout 0.15 \ + --apply-actions + + +""" + +from __future__ import annotations + +import logging +import sys +import threading +import time +from typing import Any + +import cv2 +import numpy as np + +try: + import agibot_gdk +except Exception as exc: # pragma: no cover - only available in G2 runtime + raise RuntimeError( + "Failed to import agibot_gdk. Source ~/.cache/agibot/app/env.sh and " + "use the Python version shipped with the G2 GDK." + ) from exc + +import robot_live_client as live + +_shared_execute_arm_trajectory = live._execute_arm_trajectory +_shared_execute_arm_direct48 = live._execute_arm_direct48 +_shared_maybe_smooth_actions = live._maybe_smooth_actions +_SharedWebsocketClientPolicy = live.WebsocketClientPolicy + + +ARM_JOINT_NAMES = [ + *(f"idx2{i}_arm_l_joint{i}" for i in range(1, 8)), + *(f"idx6{i}_arm_r_joint{i}" for i in range(1, 8)), +] +HEAD_JOINT_NAMES = ["idx11_head_joint1", "idx12_head_joint2", "idx13_head_joint3"] +BODY_JOINT_NAMES = [f"idx0{i}_body_joint{i}" for i in range(1, 6)] +G2_ACTION_DIM = 16 + + +def _decode_g2_relative_action( + actions: np.ndarray, + current_arm: np.ndarray, + current_gripper: np.ndarray, +) -> np.ndarray: + """ + Convert DreamZero relative action into absolute G2 joint targets. + + Training: + action = target_joint - current_joint + + Deployment: + target_joint = predicted_delta + current_joint + + G2 action layout: + [left_arm7, + left_gripper, + right_arm7, + right_gripper] + """ + + actions = np.asarray(actions, dtype=np.float32) + if actions.ndim == 1: + actions = actions.reshape(1, -1) + if actions.ndim != 2 or actions.shape[1] != G2_ACTION_DIM: + raise ValueError( + f"Expected G2 relative action shape (T, {G2_ACTION_DIM}), " + f"got {actions.shape}" + ) + + current_arm = np.asarray(current_arm, dtype=np.float32).reshape(14) + current_gripper = np.asarray(current_gripper, dtype=np.float32).reshape(2) + decoded = actions.copy() + + decoded[:, 0:7] += current_arm[0:7] + decoded[:, 7] += current_gripper[0] + decoded[:, 8:15] += current_arm[7:14] + decoded[:, 15] += current_gripper[1] + return decoded + + +G2_GRIPPER_OPEN_POSITION = -0.785 +G2_GRIPPER_CLOSED_POSITION = 0.0 +G2_ARM_GRIPPER_MIN_INTERVAL_S = 0.050 +# Source: ~/.cache/agibot/app/gdk/config/mc_impl_config.json. +# Keep targets just inside the GDK boundary to avoid float32 round-off at the +# exact limit. Values are ordered like ARM_JOINT_NAMES (left 7, then right 7). +G2_ARM_JOINT_LIMIT_EPSILON = 1e-3 +G2_ARM_JOINT_MIN = np.asarray( + [ + -3.071796, + -2.059505, + -3.071796, + -2.495838, + -3.071796, + -1.012308, + -1.535907, + -3.071796, + -2.059505, + -3.071796, + -2.495838, + -3.071796, + -1.012308, + -1.535907, + ], + dtype=np.float32, +) +G2_ARM_JOINT_MAX = np.asarray( + [ + 3.071796, + 2.059505, + 3.071796, + 1.012308, + 3.071796, + 1.012308, + 1.535907, + 3.071796, + 2.059505, + 3.071796, + 1.012308, + 3.071796, + 1.012308, + 1.535907, + ], + dtype=np.float32, +) + + +class G2WebsocketClientPolicy(_SharedWebsocketClientPolicy): + """WebSocket policy client used by the G2 live client.""" + + +def _position_by_name(robot: Any) -> tuple[dict[str, float], int]: + response = robot.get_joint_states() + states = response.get("states", []) + positions = { + str(state["name"]): float(state.get("motor_position", state.get("position", 0.0))) + for state in states + } + timestamp = int(response.get("timestamp", time.time_ns())) + return positions, timestamp + + +class G2Camera: + """Present the G2 Camera API as the three named streams used by DreamZero.""" + + _TYPES: dict[str, Any] = { + "head": agibot_gdk.CameraType.kHeadColor, + "hand_left": agibot_gdk.CameraType.kHandLeftColor, + "hand_right": agibot_gdk.CameraType.kHandRightColor, + } + + def __init__(self, names: list[str]) -> None: + unknown = set(names) - self._TYPES.keys() + if unknown: + raise ValueError(f"Unsupported G2 camera names: {sorted(unknown)}") + self._camera = agibot_gdk.Camera() + + @staticmethod + def _decode(image: Any) -> np.ndarray | None: + if image is None or not hasattr(image, "data"): + return None + data = image.data + if data is None or np.asarray(data).size == 0: + return None + if image.encoding in (agibot_gdk.Encoding.JPEG, agibot_gdk.Encoding.PNG): + return cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR) + if image.encoding != agibot_gdk.Encoding.UNCOMPRESSED: + raise RuntimeError(f"Unsupported G2 image encoding: {image.encoding}") + + raw = np.frombuffer(data, dtype=np.uint8) + if image.color_format == agibot_gdk.ColorFormat.GRAY8: + gray = raw.reshape((image.height, image.width)) + return cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR) + frame = raw.reshape((image.height, image.width, 3)) + if image.color_format == agibot_gdk.ColorFormat.RGB: + frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) + elif image.color_format != agibot_gdk.ColorFormat.BGR: + raise RuntimeError(f"Unsupported G2 color format: {image.color_format}") + return frame + + def get_latest_image(self, name: str) -> tuple[np.ndarray | None, int]: + image = self._camera.get_latest_image(self._TYPES[name], 1000.0) + if image is None: + return None, 0 + timestamp = int(getattr(image, "timestamp_ns", getattr(image, "timestamp", 0))) + return self._decode(image), timestamp + + def close(self) -> None: + self._camera.close_camera() + + +class G2Robot: + """Compatibility adapter around ``agibot_gdk.Robot`` and ``Pnc``.""" + + def __init__(self) -> None: + self._robot = agibot_gdk.Robot() + self._pnc = None + self._head_roll = 0.0 + self._waist_warning_emitted = False + self._last_arm_gripper_command_at: float | None = None + self._last_arm_gripper_command_kind: str | None = None + # Arm and gripper commands share one G2 control channel. Serialize both + # paths so commands cannot overlap across client worker threads. + self._arm_gripper_lock = threading.Lock() + + def _wait_arm_gripper_switch(self, command_kind: str) -> None: + if ( + self._last_arm_gripper_command_at is not None + and self._last_arm_gripper_command_kind != command_kind + ): + remaining = ( + G2_ARM_GRIPPER_MIN_INTERVAL_S + - (time.monotonic() - self._last_arm_gripper_command_at) + ) + if remaining > 0.0: + time.sleep(remaining) + + def _mark_arm_gripper_command(self, command_kind: str) -> None: + self._last_arm_gripper_command_at = time.monotonic() + self._last_arm_gripper_command_kind = command_kind + + def arm_joint_states(self) -> tuple[list[float], int]: + positions, timestamp = _position_by_name(self._robot) + missing = [name for name in ARM_JOINT_NAMES if name not in positions] + if missing: + raise RuntimeError(f"G2 joint-state response is missing arm joints: {missing}") + return [positions[name] for name in ARM_JOINT_NAMES], timestamp + + def head_joint_states(self) -> tuple[list[float], int]: + positions, timestamp = _position_by_name(self._robot) + missing = [name for name in HEAD_JOINT_NAMES if name not in positions] + if missing: + raise RuntimeError(f"G2 joint-state response is missing head joints: {missing}") + self._head_roll = positions[HEAD_JOINT_NAMES[1]] + return [positions[HEAD_JOINT_NAMES[0]], positions[HEAD_JOINT_NAMES[2]]], timestamp + + def waist_joint_states(self) -> tuple[list[float], int]: + # Raw body motor angles are not equivalent to policy [pitch, lift]. + _, timestamp = _position_by_name(self._robot) + return [0.0, 0.0], timestamp + + def gripper_states(self) -> tuple[list[float], int]: + response = self._robot.get_end_state() + values: list[float] = [] + for side in ("left", "right"): + end = response.get(f"{side}_end_state", {}) + states = end.get("end_states", []) + position = float(states[0].get("position", 0.0)) if states else 0.0 + # G2 training data, policy output, and the omnipicker SDK all use + # the same physical range: -0.785 is open and 0 is closed. + values.append( + float( + np.clip( + position, + G2_GRIPPER_OPEN_POSITION, + G2_GRIPPER_CLOSED_POSITION, + ) + ) + ) + return values, int(time.time_ns()) + + def _joint_request(self, names: list[str], positions: list[float], speed: float = 0.3) -> None: + request = agibot_gdk.JointControlReq() + request.life_time = 1.0 + request.joint_names = names + logging.info( + "FINAL SDK ARM CMD=%s", + positions + ) + request.joint_positions = [float(value) for value in positions] + request.joint_velocities = [float(speed)] * len(names) + result = self._robot.joint_control_request(request) + if result not in (None, 0): + raise RuntimeError(f"G2 joint_control_request failed with result {result}") + + def move_arm(self, positions: list[float]) -> None: + if len(positions) != 14: + raise ValueError(f"G2 arm command must contain 14 positions, got {len(positions)}") + + # The shared executor applies interpolation and delta limits after the + # policy action was clipped. Clamp the final values again immediately + # before the SDK call so a boundary joint cannot drift out of range. + values = np.asarray(positions, dtype=np.float64) + safe_min = ( + G2_ARM_JOINT_MIN.astype(np.float64) + + G2_ARM_JOINT_LIMIT_EPSILON + ) + safe_max = ( + G2_ARM_JOINT_MAX.astype(np.float64) + - G2_ARM_JOINT_LIMIT_EPSILON + ) + clipped = np.clip(values, safe_min, safe_max) + clipped_mask = clipped != values + if np.any(clipped_mask): + affected = [ + ARM_JOINT_NAMES[index] + for index in np.flatnonzero(clipped_mask) + ] + logging.warning( + "Clipped final G2 SDK arm command inside joint limits; affected joints=%s", + affected, + ) + + with self._arm_gripper_lock: + self._wait_arm_gripper_switch("arm") + try: + self._joint_request(ARM_JOINT_NAMES, clipped.tolist()) + finally: + self._mark_arm_gripper_command("arm") + + def move_head(self, positions: list[float]) -> None: + if len(positions) != 2: + raise ValueError(f"G2 policy head command must contain 2 positions, got {len(positions)}") + # G2 order is yaw, roll, pitch. Preserve the unmodelled roll axis. + self._robot.move_head_joint( + [float(positions[0]), self._head_roll, float(positions[1])], + [0.3, 0.3, 0.3], + ) + + def move_waist(self, positions: list[float]) -> None: + if not self._waist_warning_emitted: + logging.warning( + "Ignoring waist action: G2 uses five coupled body motors and the supplied SDK " + "examples do not expose a safe pitch/lift command API." + ) + self._waist_warning_emitted = True + + def move_gripper(self, positions: list[float]) -> None: + if len(positions) != 2: + raise ValueError(f"G2 gripper command must contain 2 positions, got {len(positions)}") + command = agibot_gdk.JointStates() + command.group = "dual_tool" + command.target_type = "omnipicker" + states: list[Any] = [] + for value in positions: + state = agibot_gdk.JointState() + state.position = float( + np.clip( + value, + G2_GRIPPER_OPEN_POSITION, + G2_GRIPPER_CLOSED_POSITION, + ) + ) + states.append(state) + command.states = states + command.nums = len(states) + with self._arm_gripper_lock: + self._wait_arm_gripper_switch("gripper") + try: + result = self._robot.move_ee_pos(command) + finally: + self._mark_arm_gripper_command("gripper") + if result != 0: + raise RuntimeError(f"G2 dual gripper move_ee_pos failed with result {result}") + + def move_wheel(self, linear: float, angular: float) -> None: + if self._pnc is None: + self._pnc = agibot_gdk.Pnc() + self._pnc.request_chassis_control(0) + time.sleep(0.5) + twist = agibot_gdk.Twist() + twist.linear = agibot_gdk.Vector3() + twist.angular = agibot_gdk.Vector3() + twist.linear.x = float(linear) + twist.angular.z = float(angular) + self._pnc.move_chassis(twist) + + def shutdown(self) -> None: + # G2 objects have no shutdown method in the supplied SDK examples. + if self._pnc is not None: + self.move_wheel(0.0, 0.0) + + +def _build_g2_obs( + head_img: np.ndarray, + left_img: np.ndarray, + right_img: np.ndarray, + arm_pos: list[float], + head_pos: list[float], + waist_pos: list[float], + gripper_pos: list[float], + prompt: str, + session_id: str, + sdk_arm_order: live.ArmOrder, + obs_flip_config: live.ObsFlipConfig, + image_transport: live.ImageTransportMode, + image_jpeg_quality: int, +) -> dict[str, object]: + """Build the observation keys declared by modality_config_g2.""" + del head_pos, waist_pos + policy_arm = live._sdk_to_policy_arm(arm_pos, sdk_arm_order) + policy_arm = live._apply_obs_joint_sign_flips(policy_arm, obs_flip_config) + gripper = np.asarray(gripper_pos, dtype=np.float32) + if gripper.shape != (2,): + raise ValueError(f"G2 gripper state must contain 2 values, got {gripper.shape}") + return { + "observation/top_head": live._encode_video_observation( + head_img, image_transport=image_transport, jpeg_quality=image_jpeg_quality + ), + "observation/hand_left": live._encode_video_observation( + left_img, image_transport=image_transport, jpeg_quality=image_jpeg_quality + ), + "observation/hand_right": live._encode_video_observation( + right_img, image_transport=image_transport, jpeg_quality=image_jpeg_quality + ), + "observation/left_joint_position": policy_arm[:7], + "observation/left_gripper_position": gripper[:1], + "observation/right_joint_position": policy_arm[7:], + "observation/right_gripper_position": gripper[1:], + "prompt": prompt, + "session_id": session_id, + } + + +def _parse_g2_action_row(row: np.ndarray, horizon: int) -> dict[str, np.ndarray]: + """Parse [left arm 7, left grip 1, right arm 7, right grip 1].""" + row = np.asarray(row, dtype=np.float32).reshape(-1) + if row.shape[0] == 22: + # Internal representation used only by the shared 22-D executor. + return { + "left_arm": row[0:7], + "right_arm": row[7:14], + "gripper": row[14:16], + "head": row[16:18], + "waist": row[18:20], + "wheel": row[20:22], + "horizon": np.asarray([horizon], dtype=np.int32), + } + if row.shape[0] != G2_ACTION_DIM: + raise ValueError(f"Expected G2 action dimension 16, got {row.shape[0]}") + return { + "left_arm": row[0:7], + "right_arm": row[8:15], + "gripper": np.asarray([row[7], row[15]], dtype=np.float32), + # NaN marks modalities which do not exist in the G2 policy output, so + # the shared executor skips their command paths. + "head": np.full(2, np.nan, dtype=np.float32), + "waist": np.full(2, np.nan, dtype=np.float32), + "wheel": np.full(2, np.nan, dtype=np.float32), + "horizon": np.asarray([horizon], dtype=np.int32), + } + + +def _parse_g2_action_first(actions: np.ndarray) -> dict[str, np.ndarray]: + actions = np.asarray(actions, dtype=np.float32) + if actions.ndim == 1: + actions = actions.reshape(1, -1) + if actions.ndim != 2 or actions.shape[1] != G2_ACTION_DIM: + raise ValueError(f"Expected G2 action shape (T, 16), got {actions.shape}") + return _parse_g2_action_row(actions[0], actions.shape[0]) + + +def _select_g2_gripper_command( + actions: np.ndarray, gripper_config: live.GripperConfig +) -> dict[str, np.ndarray]: + actions = np.asarray(actions, dtype=np.float32) + if actions.ndim == 1: + actions = actions.reshape(1, -1) + if actions.ndim != 2 or actions.shape[1] != G2_ACTION_DIM: + raise ValueError(f"Expected G2 action shape (T, 16), got {actions.shape}") + first = actions[0, [7, 15]].astype(np.float32) + last = actions[-1, [7, 15]].astype(np.float32) + command = live._gripper_policy_to_command(last, gripper_config) + return {"first_policy": first, "last_policy": last, "command_policy": command} + + +def _g2_gripper_policy_to_command( + policy_values: np.ndarray, gripper_config: live.GripperConfig +) -> np.ndarray: + """Pass G2 omnipicker positions through in the native SDK scale. + + G2 inference returns physical positions in [-0.785, 0], so thresholding + them into the shared client's normalized open/closed values would reverse + or discard the command. ``gripper_config`` is intentionally unused on G2. + """ + del gripper_config + pair = np.asarray(policy_values, dtype=np.float32).reshape(-1) + if pair.shape[0] != 2: + raise ValueError(f"Expected G2 gripper pair with 2 dims, got {pair.shape[0]}") + return np.clip( + pair, + G2_GRIPPER_OPEN_POSITION, + G2_GRIPPER_CLOSED_POSITION, + ).astype(np.float32) + + +def _g2_gripper_state_for_log(values: list | np.ndarray) -> np.ndarray: + """Keep G2 diagnostic state in the native range used on the wire.""" + pair = np.asarray(values, dtype=np.float32).reshape(-1) + if pair.shape[0] != 2: + raise ValueError(f"Expected G2 gripper state with 2 dims, got {pair.shape[0]}") + return np.clip( + pair, + G2_GRIPPER_OPEN_POSITION, + G2_GRIPPER_CLOSED_POSITION, + ).astype(np.float32) + + +def _clip_g2_arm_joint_limits(actions: np.ndarray) -> np.ndarray: + """Clamp G2 arm targets to the absolute limits enforced by the GDK.""" + arm_targets = np.concatenate((actions[:, 0:7], actions[:, 8:15]), axis=1) + safe_min = G2_ARM_JOINT_MIN + G2_ARM_JOINT_LIMIT_EPSILON + safe_max = G2_ARM_JOINT_MAX - G2_ARM_JOINT_LIMIT_EPSILON + clipped_targets = np.clip(arm_targets, safe_min, safe_max).astype(np.float32) + clipped_mask = clipped_targets != arm_targets + if np.any(clipped_mask): + affected = [ + ARM_JOINT_NAMES[index] + for index in range(len(ARM_JOINT_NAMES)) + if np.any(clipped_mask[:, index]) + ] + logging.warning( + "Clipped %s G2 arm target value(s) to GDK absolute joint limits; affected joints=%s", + int(np.count_nonzero(clipped_mask)), + affected, + ) + return clipped_targets + + +def _g2_to_shared_actions(actions: np.ndarray) -> np.ndarray: + """Reorder external G2 actions into the shared executor's 22-D layout.""" + actions = np.asarray(actions, dtype=np.float32) + if actions.ndim == 1: + actions = actions.reshape(1, -1) + if actions.ndim != 2 or actions.shape[1] != G2_ACTION_DIM: + raise ValueError(f"Expected G2 action shape (T, 16), got {actions.shape}") + clipped_arm_targets = _clip_g2_arm_joint_limits(actions) + shared = np.full((actions.shape[0], 22), np.nan, dtype=np.float32) + shared[:, 0:14] = clipped_arm_targets + shared[:, 14] = actions[:, 7] + shared[:, 15] = actions[:, 15] + return shared + + +def _shared_to_g2_actions(actions: np.ndarray) -> np.ndarray: + actions = np.asarray(actions, dtype=np.float32) + g2 = np.empty((actions.shape[0], G2_ACTION_DIM), dtype=np.float32) + g2[:, 0:7] = actions[:, 0:7] + g2[:, 7] = actions[:, 14] + g2[:, 8:15] = actions[:, 7:14] + g2[:, 15] = actions[:, 15] + return g2 + + +def _decode_g2_actions_for_execution( + *, + robot: G2Robot, + actions: np.ndarray, + sdk_arm_order: live.ArmOrder, +) -> np.ndarray: + """Decode checkpoint-relative actions against the latest robot state.""" + current_arm_sdk, _ = robot.arm_joint_states() + current_gripper, _ = robot.gripper_states() + current_arm_policy = live._sdk_to_policy_arm( + np.asarray(current_arm_sdk, dtype=np.float32), + sdk_arm_order, + ) + decoded = _decode_g2_relative_action( + actions, + current_arm=current_arm_policy, + current_gripper=np.asarray(current_gripper, dtype=np.float32), + ) + logging.info( + "Decoded G2 relative action against execution-boundary state | " + "delta_first_left[:3]=%s delta_first_right[:3]=%s " + "target_first_left[:3]=%s target_first_right[:3]=%s", + np.round(np.asarray(actions, dtype=np.float32).reshape(-1, G2_ACTION_DIM)[0, 0:3], 4).tolist(), + np.round(np.asarray(actions, dtype=np.float32).reshape(-1, G2_ACTION_DIM)[0, 8:11], 4).tolist(), + np.round(decoded[0, 0:3], 4).tolist(), + np.round(decoded[0, 8:11], 4).tolist(), + ) + return decoded + + +def _execute_g2_trajectory( + *, + robot: G2Robot, + actions: np.ndarray, + sdk_arm_order: live.ArmOrder, + **kwargs: Any, +) -> dict[str, Any]: + decoded = _decode_g2_actions_for_execution( + robot=robot, + actions=actions, + sdk_arm_order=sdk_arm_order, + ) + return _shared_execute_arm_trajectory( + robot=robot, + actions=_g2_to_shared_actions(decoded), + sdk_arm_order=sdk_arm_order, + **kwargs, + ) + + +def _execute_g2_direct48( + *, + robot: G2Robot, + actions: np.ndarray, + sdk_arm_order: live.ArmOrder, + **kwargs: Any, +) -> dict[str, Any]: + decoded = _decode_g2_actions_for_execution( + robot=robot, + actions=actions, + sdk_arm_order=sdk_arm_order, + ) + return _shared_execute_arm_direct48( + robot=robot, + actions=_g2_to_shared_actions(decoded), + sdk_arm_order=sdk_arm_order, + **kwargs, + ) + + +def _smooth_g2_actions( + actions: np.ndarray, config: live.ActionSmoothingConfig +) -> tuple[np.ndarray, float]: + shared, duration = _shared_maybe_smooth_actions(_g2_to_shared_actions(actions), config) + return _shared_to_g2_actions(shared), duration + + +def main() -> None: + result = agibot_gdk.gdk_init() + success = getattr(getattr(agibot_gdk, "GDKRes", object), "kSuccess", None) + if success is not None and result != success: + raise RuntimeError(f"agibot_gdk.gdk_init() failed: {result}") + + live.Camera = G2Camera + live.Robot = G2Robot + live.WebsocketClientPolicy = G2WebsocketClientPolicy + live._build_obs = _build_g2_obs + live._parse_action_row = _parse_g2_action_row + live._parse_action_first = _parse_g2_action_first + live._select_gripper_command = _select_g2_gripper_command + live._gripper_policy_to_command = _g2_gripper_policy_to_command + live._sdk_gripper_to_policy_obs = _g2_gripper_state_for_log + live._execute_arm_trajectory = _execute_g2_trajectory + live._execute_arm_direct48 = _execute_g2_direct48 + live._maybe_smooth_actions = _smooth_g2_actions + + # modality_config_g2.eval_delta_indices is [0], unlike the four-frame G1 + # evaluation history. Respect an explicit CLI override when one is given. + if "--observation-history" not in sys.argv: + sys.argv.extend(["--observation-history", "1"]) + live.main() + + +if __name__ == "__main__": + main() diff --git a/scripts/audit_g2_checkpoint.py b/scripts/audit_g2_checkpoint.py new file mode 100644 index 00000000..0904c939 --- /dev/null +++ b/scripts/audit_g2_checkpoint.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""CPU-only validation of a DreamZero G2 LoRA deployment checkpoint.""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from collections import Counter +from pathlib import Path + + +EXPECTED_ACTION_HORIZON = 24 +EXPECTED_NUM_FRAMES = 33 +EXPECTED_ACTION_DIM = 32 +EXPECTED_OUTPUT_DIM = 16 + + +def _read_safetensors_header(path: Path) -> dict: + with path.open("rb") as stream: + header_size_bytes = stream.read(8) + if len(header_size_bytes) != 8: + raise ValueError(f"{path} has an incomplete safetensors header") + header_size = int.from_bytes(header_size_bytes, "little") + if header_size <= 0 or header_size > path.stat().st_size - 8: + raise ValueError( + f"{path} has an invalid safetensors header size: {header_size}" + ) + return json.loads(stream.read(header_size)) + + +def _nested(config: dict, *keys: str): + value = config + for key in keys: + if not isinstance(value, dict) or key not in value: + return None + value = value[key] + return value + + +def audit(checkpoint: Path) -> list[str]: + errors: list[str] = [] + required = [ + checkpoint / "config.json", + checkpoint / "model.safetensors", + checkpoint / "experiment_cfg" / "conf.yaml", + checkpoint / "experiment_cfg" / "metadata.json", + ] + for path in required: + if not path.is_file(): + errors.append(f"missing required file: {path}") + if errors: + return errors + + config = json.loads((checkpoint / "config.json").read_text()) + inner = _nested(config, "action_head_cfg", "config") or {} + diffusion = inner.get("diffusion_model_cfg", {}) + checks = { + "config.action_horizon": ( + config.get("action_horizon"), + EXPECTED_ACTION_HORIZON, + ), + "action_head.action_horizon": ( + inner.get("action_horizon"), + EXPECTED_ACTION_HORIZON, + ), + "action_head.num_frames": ( + inner.get("num_frames"), + EXPECTED_NUM_FRAMES, + ), + "action_head.action_dim": ( + inner.get("action_dim"), + EXPECTED_ACTION_DIM, + ), + "diffusion.num_action_per_block": ( + diffusion.get("num_action_per_block"), + EXPECTED_ACTION_HORIZON, + ), + "diffusion.out_dim": ( + diffusion.get("out_dim"), + EXPECTED_OUTPUT_DIM, + ), + } + for name, (actual, expected) in checks.items(): + if actual != expected: + errors.append(f"{name}: expected {expected}, got {actual!r}") + + conf_text = ( + checkpoint / "experiment_cfg" / "conf.yaml" + ).read_text(errors="replace") + # Support both old LoRA-only checkpoints and new action-adapter-only + # checkpoints. The latter intentionally freezes shared DiT/LoRA and only + # stores state/action adapter parameters. + pretrained_match = re.search( + r"(?m)^pretrained_model_path:\s*(\S+)\s*$", + conf_text, + ) + if not pretrained_match: + errors.append("experiment config has no pretrained_model_path") + else: + base_path = Path(pretrained_match.group(1)) + if not base_path.is_dir(): + errors.append(f"pretrained base directory is missing: {base_path}") + elif not ( + (base_path / "model.safetensors").is_file() + or (base_path / "model.safetensors.index.json").is_file() + ): + errors.append( + f"pretrained base has no safetensors weights: {base_path}" + ) + + header = _read_safetensors_header(checkpoint / "model.safetensors") + keys = [key for key in header if key != "__metadata__"] + buckets = Counter() + for key in keys: + lowered = key.lower() + if "lora_a" in lowered: + buckets["lora_A"] += 1 + elif "lora_b" in lowered: + buckets["lora_B"] += 1 + elif "action_encoder" in lowered: + buckets["action_encoder"] += 1 + elif "action_decoder" in lowered: + buckets["action_decoder"] += 1 + elif "state_encoder" in lowered: + buckets["state_encoder"] += 1 + else: + buckets["other"] += 1 + + has_lora = buckets["lora_A"] > 0 or buckets["lora_B"] > 0 + has_action_adapter = ( + buckets["action_encoder"] > 0 + or buckets["action_decoder"] > 0 + or buckets["state_encoder"] > 0 + ) + + if has_lora: + print("checkpoint_type: LoRA") + if buckets["lora_A"] == 0: + errors.append("checkpoint contains no LoRA A weights") + if buckets["lora_A"] != buckets["lora_B"]: + errors.append( + "unbalanced LoRA weights: " + f"A={buckets['lora_A']} B={buckets['lora_B']}" + ) + elif has_action_adapter: + print("checkpoint_type: action_adapter") + else: + errors.append( + "checkpoint contains neither LoRA weights nor action adapter weights" + ) + for name, minimum in ( + ("action_encoder", 6), + ("action_decoder", 4), + ("state_encoder", 4), + ): + if buckets[name] < minimum: + errors.append( + f"incomplete {name}: expected at least {minimum}, " + f"got {buckets[name]}" + ) + + print(f"checkpoint: {checkpoint}") + print(f"pretrained_model_path: {pretrained_match.group(1) if pretrained_match else 'MISSING'}") + print(f"tensor_keys: {len(keys)}") + print(f"tensor_buckets: {dict(buckets)}") + for name, (actual, expected) in checks.items(): + print(f"{name}: {actual!r} (expected {expected})") + return errors + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint", type=Path) + args = parser.parse_args() + + errors = audit(args.checkpoint.expanduser().resolve()) + if errors: + print("G2 CHECKPOINT AUDIT FAILED:", file=sys.stderr) + for error in errors: + print(f" - {error}", file=sys.stderr) + return 1 + print("G2 CHECKPOINT AUDIT PASSED") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/data/build_g2_active_hold_windows.py b/scripts/data/build_g2_active_hold_windows.py new file mode 100644 index 00000000..0bda8320 --- /dev/null +++ b/scripts/data/build_g2_active_hold_windows.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +"""Build the active/hold sampling index for an existing policy-space G2 split.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import pandas as pd + + +ARM_DIMS = (0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14) +GRIP_DIMS = (7, 15) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("split", type=Path) + parser.add_argument("--action-horizon", type=int, default=24) + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + + output = args.split / "meta/g2_active_hold_windows.json" + if output.exists() and not args.overwrite: + raise FileExistsError(f"Refusing to overwrite {output}; pass --overwrite") + + metrics: list[tuple[int, int, float, bool]] = [] + for path in sorted(args.split.glob("data/chunk-*/*.parquet")): + frame = pd.read_parquet(path, columns=["observation.state", "action", "episode_index"]) + state = np.stack(frame["observation.state"]).astype(np.float32) + action = np.stack(frame["action"]).astype(np.float32) + if state.shape[1:] != (16,) or action.shape[1:] != (16,): + raise ValueError(f"{path}: expected 16D state/action") + if np.any(state[:, GRIP_DIMS] < 0) or np.any(state[:, GRIP_DIMS] > 1): + raise ValueError(f"{path}: state gripper is not in policy space [0,1]") + if np.any(action[:, GRIP_DIMS] < 0) or np.any(action[:, GRIP_DIMS] > 1): + raise ValueError(f"{path}: action gripper is not in policy space [0,1]") + episode = int(frame["episode_index"].iloc[0]) + for step in range(max(0, len(frame) - args.action_horizon + 1)): + future = action[step : step + args.action_horizon] + arm_motion = float( + np.linalg.norm(future[:, ARM_DIMS] - state[step, ARM_DIMS], axis=1).max() + ) + grip_transition = bool( + np.any(np.abs(future[:, GRIP_DIMS] - state[step, GRIP_DIMS]) > 0.25) + ) + metrics.append((episode, step, arm_motion, grip_transition)) + + nonzero = np.asarray([motion for _, _, motion, _ in metrics if motion > 1e-8]) + if not nonzero.size: + raise RuntimeError("No nonzero G2 arm motion found") + threshold = float(np.quantile(nonzero, 0.30)) + active: dict[str, list[int]] = {} + hold: dict[str, list[int]] = {} + for episode, step, motion, grip_transition in metrics: + target = active if motion >= threshold or grip_transition else hold + target.setdefault(str(episode), []).append(step) + + payload = { + "schema_version": 1, + "action_horizon": args.action_horizon, + "arm_motion_threshold_rule": "p30_nonzero_24_step_l2", + "arm_motion_threshold": threshold, + "gripper_transition_threshold": 0.25, + "active_count": sum(map(len, active.values())), + "hold_count": sum(map(len, hold.values())), + "active": active, + "hold": hold, + } + output.write_text(json.dumps(payload, indent=2), encoding="utf-8") + print( + f"Wrote {output}: threshold={threshold:.8f} " + f"active={payload['active_count']} hold={payload['hold_count']}" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/data/convert_lerobot_g2_to_gear.py b/scripts/data/convert_lerobot_g2_to_gear.py new file mode 100644 index 00000000..6a88abe6 --- /dev/null +++ b/scripts/data/convert_lerobot_g2_to_gear.py @@ -0,0 +1,953 @@ +#!/usr/bin/env python3 +r"""Convert a dual-arm G2 LeRobot v3 dataset into DreamZero/GEAR datasets. + +The converter is intentionally strict. It validates the 16-D joint layout, +materializes one parquet and one video per episode, forces all three camera +streams to the same CFR/resolution, writes GEAR metadata, and can physically +separate a held-out test set from the training set. + +Expected packed state/action layout (from create_g2_dataset_using_lerobot.py): + [left_joint_1..7, left_gripper, right_joint_1..7, right_gripper] + +Example (120 episodes -> 110 train + 10 held-out test): + python convert_lerobot_g2_to_gear.py \ + --source /data/.../g2_mock_light_module_joint_streaming \ + --output /data/.../g2_mock_light_module_gear \ + --test-episodes 10 --split-mode tail \ + --video-width 320 --video-height 176 \ + --video-codec libx264 --workers 8 + +Output: + /train/ # pass this directory to DreamZero training + /test/ # never read by the training job + /split_manifest.json + +For a quick smoke test, add ``--max-source-episodes 2 --test-episodes 0``. +""" + +from __future__ import annotations + +import argparse +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +import json +import logging +from pathlib import Path +import random +import shutil +import subprocess +import sys +from typing import Any, Iterable + +import numpy as np +import pyarrow as pa +import pyarrow.dataset as pads +import pyarrow.parquet as pq + + +LOG = logging.getLogger("g2-v3-to-gear") + +SOURCE_CAMERAS = { + "observation.images.head_color": "observation.images.top_head", + "observation.images.hand_left_color": "observation.images.hand_left", + "observation.images.hand_right_color": "observation.images.hand_right", +} +VIDEO_MODALITY_ORDER = ("top_head", "hand_left", "hand_right") + +JOINT_SLICES: dict[str, tuple[int, int]] = { + "left_joint_position": (0, 7), + "left_gripper_position": (7, 8), + "right_joint_position": (8, 15), + "right_gripper_position": (15, 16), +} +EXPECTED_NAMES = [ + *(f"l.joint{i}.pos" for i in range(1, 8)), + "l.gripper.pos", + *(f"r.joint{i}.pos" for i in range(1, 8)), + "r.gripper.pos", +] + + +@dataclass(frozen=True) +class SourceEpisode: + source_index: int + length: int + task: str + metadata: dict[str, Any] + + +@dataclass(frozen=True) +class OutputEpisode: + source: SourceEpisode + output_index: int + + +def nonempty_path(value: str) -> Path: + if not value.strip(): + raise argparse.ArgumentTypeError("path cannot be empty") + return Path(value) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--source", type=nonempty_path, required=True) + parser.add_argument("--output", type=nonempty_path, required=True) + parser.add_argument("--embodiment-tag", default="g2") + parser.add_argument("--workers", type=int, default=8) + parser.add_argument("--action-horizon", type=int, default=24) + parser.add_argument("--video-width", type=int, default=320) + parser.add_argument("--video-height", type=int, default=176) + parser.add_argument( + "--resize-mode", + choices=("stretch", "pad", "crop"), + default="stretch", + help="How to make all camera streams the same size.", + ) + parser.add_argument( + "--video-codec", + choices=("libx264", "h264_nvenc"), + default="libx264", + ) + parser.add_argument("--video-preset", default="veryfast") + parser.add_argument("--video-crf", type=int, default=18) + parser.add_argument( + "--test-episodes", + type=int, + default=10, + help="Number of physically isolated held-out episodes.", + ) + parser.add_argument( + "--split-mode", + choices=("tail", "random"), + default="tail", + help="tail reserves the final N source episodes; random uses --split-seed.", + ) + parser.add_argument("--split-seed", type=int, default=42) + parser.add_argument( + "--max-source-episodes", + type=int, + default=None, + help="Use only the first N source episodes (smoke tests only).", + ) + parser.add_argument( + "--min-episode-frames", + type=int, + default=None, + help="Default is action_horizon + 1; shorter episodes fail validation.", + ) + parser.add_argument("--overwrite", action="store_true") + return parser.parse_args() + + +def read_json(path: Path) -> dict[str, Any]: + with path.open(encoding="utf-8") as handle: + value = json.load(handle) + if not isinstance(value, dict): + raise ValueError(f"Expected a JSON object: {path}") + return value + + +def write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + json.dump(value, handle, indent=2, ensure_ascii=False) + handle.write("\n") + + +def write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + for row in rows: + handle.write(json.dumps(row, ensure_ascii=False) + "\n") + + +def require_tools(video_codec: str) -> None: + for executable in ("ffmpeg", "ffprobe"): + if shutil.which(executable) is None: + raise RuntimeError(f"Required executable not found: {executable}") + if video_codec == "h264_nvenc": + completed = subprocess.run( + ["ffmpeg", "-hide_banner", "-encoders"], + check=True, + capture_output=True, + text=True, + ) + if "h264_nvenc" not in completed.stdout: + raise RuntimeError("ffmpeg does not provide the h264_nvenc encoder") + + +def validate_paths(source: Path, output: Path, overwrite: bool) -> None: + source = source.resolve() + output = output.resolve() + if not source.is_dir(): + raise FileNotFoundError(source) + if output == source or output in source.parents or source in output.parents: + raise ValueError(f"Source and output must not overlap: {source} / {output}") + if output == Path(output.anchor) or output == Path.home().resolve(): + raise ValueError(f"Refusing dangerous output path: {output}") + if (output / ".git").exists(): + raise ValueError(f"Refusing to overwrite a Git repository: {output}") + if output.exists(): + if not overwrite: + raise FileExistsError(f"Output exists; pass --overwrite to rebuild: {output}") + shutil.rmtree(output) + output.mkdir(parents=True) + + +def parquet_files(path: Path) -> list[Path]: + files = sorted(path.glob("chunk-*/*.parquet")) + if not files: + raise FileNotFoundError(f"No parquet files below {path}") + return files + + +def read_many_parquets(files: list[Path]) -> pa.Table: + return pads.dataset([str(path) for path in files], format="parquet").to_table() + + +def validate_source_info(info: dict[str, Any]) -> None: + if info.get("codebase_version") != "v3.0": + raise ValueError( + f"Expected LeRobot v3.0, got {info.get('codebase_version')!r}" + ) + features = info.get("features", {}) + for key in ("observation.state", "action", *SOURCE_CAMERAS): + if key not in features: + raise ValueError(f"Missing source feature: {key}") + for key in ("observation.state", "action"): + shape = list(features[key].get("shape", [])) + if shape != [16]: + raise ValueError(f"{key} must be 16-D joint data, got shape {shape}") + names = features[key].get("names") + if names and list(names) != EXPECTED_NAMES: + raise ValueError( + f"{key} names do not match the G2 joint layout.\n" + f"Expected: {EXPECTED_NAMES}\nGot: {names}" + ) + fps = float(info.get("fps", 0)) + if fps <= 0: + raise ValueError(f"Invalid source fps: {fps}") + + +def load_source_episodes( + source: Path, info: dict[str, Any], max_source_episodes: int | None +) -> tuple[list[SourceEpisode], list[Path], pads.Dataset]: + data_files = parquet_files(source / "data") + episode_files = parquet_files(source / "meta/episodes") + tasks_path = source / "meta/tasks.parquet" + if not tasks_path.is_file(): + raise FileNotFoundError(tasks_path) + + episode_rows = read_many_parquets(episode_files).to_pylist() + task_rows = pq.read_table(tasks_path).to_pylist() + task_by_index = {int(row["task_index"]): str(row["task"]).strip() for row in task_rows} + episodes: list[SourceEpisode] = [] + for row in sorted(episode_rows, key=lambda item: int(item["episode_index"])): + source_index = int(row["episode_index"]) + declared_tasks = [str(task).strip() for task in row.get("tasks", []) if str(task).strip()] + task = declared_tasks[0] if declared_tasks else "" + if not task and "task_index" in row: + task = task_by_index.get(int(row["task_index"]), "") + if not task: + raise ValueError(f"Episode {source_index} has no language task") + if len(set(declared_tasks)) > 1: + raise ValueError(f"Episode {source_index} declares multiple tasks: {declared_tasks}") + episodes.append( + SourceEpisode( + source_index=source_index, + length=int(row["length"]), + task=task, + metadata=row, + ) + ) + + total = int(info["total_episodes"]) + if len(episodes) != total: + raise ValueError(f"Episode metadata has {len(episodes)} rows; info.json says {total}") + if max_source_episodes is not None: + if max_source_episodes < 1 or max_source_episodes > len(episodes): + raise ValueError( + f"--max-source-episodes must be in [1, {len(episodes)}]" + ) + episodes = episodes[:max_source_episodes] + + data_dataset = pads.dataset([str(path) for path in data_files], format="parquet") + required_columns = { + "observation.state", + "action", + "episode_index", + } + missing = required_columns - set(data_dataset.schema.names) + if missing: + raise ValueError(f"Source data is missing columns: {sorted(missing)}") + return episodes, data_files, data_dataset + + +def split_episodes( + episodes: list[SourceEpisode], test_count: int, mode: str, seed: int +) -> tuple[list[SourceEpisode], list[SourceEpisode]]: + if test_count < 0 or test_count >= len(episodes): + if test_count == 0: + return episodes, [] + raise ValueError(f"--test-episodes must be in [0, {len(episodes) - 1}]") + if test_count == 0: + return episodes, [] + if mode == "tail": + return episodes[:-test_count], episodes[-test_count:] + rng = random.Random(seed) + test_ids = set(rng.sample([item.source_index for item in episodes], test_count)) + train = [item for item in episodes if item.source_index not in test_ids] + test = [item for item in episodes if item.source_index in test_ids] + return train, test + + +def validate_episode_lengths( + episodes: list[SourceEpisode], min_frames: int, split_name: str +) -> None: + short = [(item.source_index, item.length) for item in episodes if item.length < min_frames] + if short: + raise ValueError( + f"{split_name} has episodes shorter than {min_frames} frames: {short[:20]}" + ) + + +def table_for_episode(data_dataset: pads.Dataset, episode: SourceEpisode) -> pa.Table: + columns = ["observation.state", "action", "episode_index"] + for optional in ("timestamp", "frame_index", "index", "task_index"): + if optional in data_dataset.schema.names: + columns.append(optional) + table = data_dataset.to_table( + filter=pads.field("episode_index") == episode.source_index, + columns=columns, + ) + if table.num_rows != episode.length: + raise ValueError( + f"Episode {episode.source_index}: data has {table.num_rows} rows; " + f"metadata says {episode.length}" + ) + return table + + +def fixed_size_vectors(column: pa.ChunkedArray, name: str) -> np.ndarray: + values = np.asarray(column.to_pylist(), dtype=np.float32) + if values.ndim != 2 or values.shape[1] != 16: + raise ValueError(f"{name} must have shape [N, 16], got {values.shape}") + if not np.isfinite(values).all(): + raise ValueError(f"{name} contains NaN or infinity") + return values + + +def output_table( + source_table: pa.Table, + episode: OutputEpisode, + fps: float, + task_index: int, + global_start: int, +) -> tuple[pa.Table, np.ndarray, np.ndarray]: + state = fixed_size_vectors(source_table["observation.state"], "observation.state") + action = fixed_size_vectors(source_table["action"], "action") + length = episode.source.length + table = pa.table( + { + "observation.state": pa.array(state.tolist(), type=pa.list_(pa.float32(), 16)), + "action": pa.array(action.tolist(), type=pa.list_(pa.float32(), 16)), + "annotation.language.action_text": pa.array( + [episode.source.task] * length, type=pa.large_string() + ), + "timestamp": pa.array( + np.arange(length, dtype=np.float32) / np.float32(fps) + ), + "frame_index": pa.array(np.arange(length, dtype=np.int64)), + "episode_index": pa.array( + np.full(length, episode.output_index, dtype=np.int64) + ), + "index": pa.array( + np.arange(global_start, global_start + length, dtype=np.int64) + ), + "task_index": pa.array(np.full(length, task_index, dtype=np.int64)), + } + ) + return table, state, action + + +class StatsAccumulator: + def __init__(self) -> None: + self.state: list[np.ndarray] = [] + self.action: list[np.ndarray] = [] + self.relative: dict[str, list[np.ndarray]] = {key: [] for key in JOINT_SLICES} + + def add(self, state: np.ndarray, action: np.ndarray, horizon: int) -> None: + self.state.append(state) + self.action.append(action) + usable = len(state) - horizon + 1 + if usable <= 0: + return + for key, (start, end) in JOINT_SLICES.items(): + reference = state[:usable, start:end] + chunks = np.stack( + [action[offset : offset + usable, start:end] - reference for offset in range(horizon)], + axis=1, + ) + self.relative[key].append(chunks.reshape(-1, end - start)) + + +def numeric_stats(values: np.ndarray) -> dict[str, list[float]]: + values64 = np.asarray(values, dtype=np.float64) + return { + "min": np.min(values64, axis=0).tolist(), + "max": np.max(values64, axis=0).tolist(), + "mean": np.mean(values64, axis=0).tolist(), + "std": np.std(values64, axis=0).tolist(), + "q01": np.quantile(values64, 0.01, axis=0).tolist(), + "q99": np.quantile(values64, 0.99, axis=0).tolist(), + } + + +def finish_stats(accumulator: StatsAccumulator) -> tuple[dict[str, Any], dict[str, Any]]: + if not accumulator.state or not accumulator.action: + raise ValueError("Cannot compute statistics for an empty split") + stats = { + "observation.state": numeric_stats(np.concatenate(accumulator.state, axis=0)), + "action": numeric_stats(np.concatenate(accumulator.action, axis=0)), + } + relative: dict[str, Any] = {} + for key, arrays in accumulator.relative.items(): + if not arrays: + raise ValueError(f"No relative-action samples for {key}") + relative[key] = numeric_stats(np.concatenate(arrays, axis=0)) + return stats, relative + + +def output_video_path(root: Path, episode_index: int, output_key: str) -> Path: + return ( + root + / f"videos/chunk-{episode_index // 1000:03d}" + / output_key + / f"episode_{episode_index:06d}.mp4" + ) + + +def probe_video(path: Path) -> tuple[int, int, int, float]: + completed = subprocess.run( + [ + "ffprobe", + "-v", + "error", + "-count_frames", + "-select_streams", + "v:0", + "-show_entries", + "stream=nb_read_frames,width,height,avg_frame_rate", + "-of", + "json", + str(path), + ], + check=True, + capture_output=True, + text=True, + ) + stream = json.loads(completed.stdout)["streams"][0] + numerator, denominator = (int(part) for part in stream["avg_frame_rate"].split("/")) + rate = numerator / denominator if denominator else 0.0 + return int(stream["nb_read_frames"]), int(stream["width"]), int(stream["height"]), rate + + +def resize_filter(width: int, height: int, mode: str) -> str: + if mode == "stretch": + return f"scale={width}:{height}:flags=lanczos,setsar=1" + if mode == "pad": + return ( + f"scale={width}:{height}:force_original_aspect_ratio=decrease:flags=lanczos," + f"pad={width}:{height}:(ow-iw)/2:(oh-ih)/2:black,setsar=1" + ) + return ( + f"scale={width}:{height}:force_original_aspect_ratio=increase:flags=lanczos," + f"crop={width}:{height},setsar=1" + ) + + +def source_video_path( + source: Path, + source_info: dict[str, Any], + metadata: dict[str, Any], + source_key: str, +) -> tuple[Path, float]: + prefix = f"videos/{source_key}" + required = ( + f"{prefix}/file_index", + f"{prefix}/chunk_index", + f"{prefix}/from_timestamp", + ) + missing = [key for key in required if key not in metadata] + if missing: + raise ValueError(f"Episode video metadata is missing: {missing}") + path = source / source_info["video_path"].format( + video_key=source_key, + chunk_index=int(metadata[f"{prefix}/chunk_index"]), + file_index=int(metadata[f"{prefix}/file_index"]), + ) + if not path.is_file(): + raise FileNotFoundError(path) + return path, float(metadata[f"{prefix}/from_timestamp"]) + + +def convert_video( + source: Path, + destination_root: Path, + source_info: dict[str, Any], + episode: OutputEpisode, + source_key: str, + output_key: str, + width: int, + height: int, + mode: str, + codec: str, + preset: str, + crf: int, +) -> dict[str, Any]: + source_path, start = source_video_path( + source, source_info, episode.source.metadata, source_key + ) + destination = output_video_path(destination_root, episode.output_index, output_key) + destination.parent.mkdir(parents=True, exist_ok=True) + fps = float(source_info["fps"]) + frame_count = episode.source.length + filters = ( + f"fps={fps:g},{resize_filter(width, height, mode)}," + f"tpad=stop_mode=clone:stop_duration=2,trim=end_frame={frame_count}," + "setpts=N/FRAME_RATE/TB" + ) + command = [ + "ffmpeg", + "-nostdin", + "-hide_banner", + "-loglevel", + "error", + "-y", + "-ss", + f"{start:.9f}", + "-i", + str(source_path), + "-an", + "-vf", + filters, + "-frames:v", + str(frame_count), + "-r", + f"{fps:g}", + "-c:v", + codec, + ] + if codec == "libx264": + command += ["-preset", preset, "-crf", str(crf)] + else: + command += ["-preset", "p4", "-cq", str(crf), "-b:v", "0"] + command += [ + "-pix_fmt", + "yuv420p", + "-g", + "2", + "-keyint_min", + "2", + "-sc_threshold", + "0", + "-movflags", + "+faststart", + str(destination), + ] + subprocess.run(command, check=True) + actual_frames, actual_width, actual_height, actual_fps = probe_video(destination) + if actual_frames != frame_count: + raise ValueError( + f"{destination}: {actual_frames} frames; expected {frame_count}" + ) + if (actual_width, actual_height) != (width, height): + raise ValueError( + f"{destination}: resolution {(actual_width, actual_height)}; " + f"expected {(width, height)}" + ) + if abs(actual_fps - fps) > 1e-3: + raise ValueError(f"{destination}: fps {actual_fps}; expected {fps}") + return { + "source_episode_index": episode.source.source_index, + "output_episode_index": episode.output_index, + "camera": output_key, + "frames": actual_frames, + "width": actual_width, + "height": actual_height, + "fps": actual_fps, + "source": str(source_path), + "output": str(destination), + } + + +def build_modality() -> dict[str, Any]: + def field(original_key: str, start: int, end: int) -> dict[str, Any]: + return { + "original_key": original_key, + "start": start, + "end": end, + "rotation_type": None, + "absolute": True, + "dtype": "float32", + "range": None, + } + + return { + "state": { + key: field("observation.state", start, end) + for key, (start, end) in JOINT_SLICES.items() + }, + "action": { + key: field("action", start, end) + for key, (start, end) in JOINT_SLICES.items() + }, + "video": { + "top_head": {"original_key": "observation.images.top_head"}, + "hand_left": {"original_key": "observation.images.hand_left"}, + "hand_right": {"original_key": "observation.images.hand_right"}, + }, + "annotation": { + "language.action_text": { + "original_key": "annotation.language.action_text" + } + }, + } + + +def build_info( + source_info: dict[str, Any], + episodes: list[OutputEpisode], + tasks: list[str], + width: int, + height: int, + split_name: str, +) -> dict[str, Any]: + fps = float(source_info["fps"]) + vector_features = { + "observation.state": { + "dtype": "float32", + "shape": [16], + "names": EXPECTED_NAMES, + }, + "action": {"dtype": "float32", "shape": [16], "names": EXPECTED_NAMES}, + } + features: dict[str, Any] = { + **vector_features, + "annotation.language.action_text": {"dtype": "string", "shape": [1]}, + "timestamp": {"dtype": "float32", "shape": [1]}, + "frame_index": {"dtype": "int64", "shape": [1]}, + "episode_index": {"dtype": "int64", "shape": [1]}, + "index": {"dtype": "int64", "shape": [1]}, + "task_index": {"dtype": "int64", "shape": [1]}, + } + for output_key in SOURCE_CAMERAS.values(): + features[output_key] = { + "dtype": "video", + "shape": [height, width, 3], + "names": ["height", "width", "channels"], + "video_info": { + "video.fps": fps, + "video.codec": "h264", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "has_audio": False, + }, + } + return { + "codebase_version": "v2.1", + "robot_type": "g2", + "total_episodes": len(episodes), + "total_frames": sum(item.source.length for item in episodes), + "total_tasks": len(tasks), + "chunks_size": 1000, + "fps": fps, + "splits": {split_name: f"0:{len(episodes)}"}, + "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", + "video_path": "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4", + "features": features, + } + + +def write_split_metadata( + root: Path, + source_info: dict[str, Any], + split_name: str, + episodes: list[OutputEpisode], + tasks: list[str], + stats: dict[str, Any], + relative_stats: dict[str, Any], + embodiment_tag: str, + width: int, + height: int, + action_horizon: int, + video_reports: list[dict[str, Any]], + stats_source: str, +) -> None: + meta = root / "meta" + task_to_index = {task: index for index, task in enumerate(tasks)} + write_json( + meta / "info.json", + build_info(source_info, episodes, tasks, width, height, split_name), + ) + write_json(meta / "modality.json", build_modality()) + write_json(meta / "embodiment.json", {"embodiment_tag": embodiment_tag}) + write_json(meta / "stats.json", stats) + write_json(meta / "relative_stats_dreamzero.json", relative_stats) + write_jsonl( + meta / "tasks.jsonl", + ({"task_index": index, "task": task} for index, task in enumerate(tasks)), + ) + write_jsonl( + meta / "episodes.jsonl", + ( + { + "episode_index": item.output_index, + "tasks": [item.source.task], + "length": item.source.length, + } + for item in episodes + ), + ) + write_json( + meta / "conversion_report.json", + { + "split": split_name, + "episode_count": len(episodes), + "frame_count": sum(item.source.length for item in episodes), + "source_episode_indices": [item.source.source_index for item in episodes], + "task_indices": { + str(item.output_index): task_to_index[item.source.task] for item in episodes + }, + "state_dim": 16, + "action_dim": 16, + "action_horizon": action_horizon, + "camera_order": list(VIDEO_MODALITY_ORDER), + "video_resolution": [width, height], + "video_count": len(video_reports), + "stats_source": stats_source, + }, + ) + + +def convert_split_data( + root: Path, + source_info: dict[str, Any], + data_dataset: pads.Dataset, + source_episodes: list[SourceEpisode], + horizon: int, +) -> tuple[list[OutputEpisode], list[str], StatsAccumulator]: + episodes = [OutputEpisode(source=item, output_index=index) for index, item in enumerate(source_episodes)] + tasks = list(dict.fromkeys(item.source.task for item in episodes)) + task_to_index = {task: index for index, task in enumerate(tasks)} + accumulator = StatsAccumulator() + global_index = 0 + for count, item in enumerate(episodes, start=1): + source_table = table_for_episode(data_dataset, item.source) + table, state, action = output_table( + source_table, + item, + float(source_info["fps"]), + task_to_index[item.source.task], + global_index, + ) + destination = ( + root + / f"data/chunk-{item.output_index // 1000:03d}" + / f"episode_{item.output_index:06d}.parquet" + ) + destination.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(table, destination, compression="zstd", row_group_size=item.source.length) + accumulator.add(state, action, horizon) + global_index += item.source.length + if count % 20 == 0 or count == len(episodes): + LOG.info("%s parquet: %d/%d", root.name, count, len(episodes)) + return episodes, tasks, accumulator + + +def convert_split_videos( + source: Path, + root: Path, + source_info: dict[str, Any], + episodes: list[OutputEpisode], + args: argparse.Namespace, +) -> list[dict[str, Any]]: + jobs = [ + (episode, source_key, output_key) + for episode in episodes + for source_key, output_key in SOURCE_CAMERAS.items() + ] + reports: list[dict[str, Any]] = [] + with ThreadPoolExecutor(max_workers=args.workers) as executor: + futures = [ + executor.submit( + convert_video, + source, + root, + source_info, + episode, + source_key, + output_key, + args.video_width, + args.video_height, + args.resize_mode, + args.video_codec, + args.video_preset, + args.video_crf, + ) + for episode, source_key, output_key in jobs + ] + for count, future in enumerate(as_completed(futures), start=1): + reports.append(future.result()) + if count % 30 == 0 or count == len(jobs): + LOG.info("%s videos: %d/%d", root.name, count, len(jobs)) + return sorted( + reports, + key=lambda item: (item["output_episode_index"], item["camera"]), + ) + + +def validate_output(root: Path, split_name: str) -> None: + meta = root / "meta" + required = ( + "info.json", + "modality.json", + "embodiment.json", + "stats.json", + "relative_stats_dreamzero.json", + "tasks.jsonl", + "episodes.jsonl", + "conversion_report.json", + ) + missing = [name for name in required if not (meta / name).is_file()] + if missing: + raise ValueError(f"{split_name}: missing metadata files: {missing}") + info = read_json(meta / "info.json") + episode_count = int(info["total_episodes"]) + data_count = len(list((root / "data").glob("chunk-*/*.parquet"))) + video_count = len(list((root / "videos").glob("chunk-*/*/*.mp4"))) + if data_count != episode_count: + raise ValueError(f"{split_name}: {data_count} parquets for {episode_count} episodes") + if video_count != episode_count * 3: + raise ValueError(f"{split_name}: {video_count} videos; expected {episode_count * 3}") + + +def main() -> int: + args = parse_args() + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + if args.workers < 1: + raise ValueError("--workers must be positive") + if args.action_horizon < 1: + raise ValueError("--action-horizon must be positive") + if min(args.video_width, args.video_height) < 2 or args.video_width % 2 or args.video_height % 2: + raise ValueError("Video width and height must be positive even integers") + + source = args.source.resolve() + output = args.output.resolve() + require_tools(args.video_codec) + validate_paths(source, output, args.overwrite) + + source_info = read_json(source / "meta/info.json") + validate_source_info(source_info) + source_episodes, data_files, data_dataset = load_source_episodes( + source, source_info, args.max_source_episodes + ) + train_source, test_source = split_episodes( + source_episodes, args.test_episodes, args.split_mode, args.split_seed + ) + min_frames = args.min_episode_frames or (args.action_horizon + 1) + validate_episode_lengths(train_source, min_frames, "train") + if test_source: + validate_episode_lengths(test_source, min_frames, "test") + + LOG.info( + "Source: %d episodes, %d parquet shards; split into %d train / %d test", + len(source_episodes), + len(data_files), + len(train_source), + len(test_source), + ) + write_json( + output / "split_manifest.json", + { + "source": str(source), + "split_mode": args.split_mode, + "split_seed": args.split_seed if args.split_mode == "random" else None, + "train_source_episode_indices": [item.source_index for item in train_source], + "test_source_episode_indices": [item.source_index for item in test_source], + }, + ) + + train_root = output / "train" + train_episodes, train_tasks, train_acc = convert_split_data( + train_root, + source_info, + data_dataset, + train_source, + args.action_horizon, + ) + train_stats, train_relative_stats = finish_stats(train_acc) + train_video_reports = convert_split_videos( + source, train_root, source_info, train_episodes, args + ) + write_split_metadata( + train_root, + source_info, + "train", + train_episodes, + train_tasks, + train_stats, + train_relative_stats, + args.embodiment_tag, + args.video_width, + args.video_height, + args.action_horizon, + train_video_reports, + "train", + ) + validate_output(train_root, "train") + + if test_source: + test_root = output / "test" + test_episodes, test_tasks, _ = convert_split_data( + test_root, + source_info, + data_dataset, + test_source, + args.action_horizon, + ) + test_video_reports = convert_split_videos( + source, test_root, source_info, test_episodes, args + ) + # Test data deliberately uses training-set normalization statistics. + write_split_metadata( + test_root, + source_info, + "test", + test_episodes, + test_tasks, + train_stats, + train_relative_stats, + args.embodiment_tag, + args.video_width, + args.video_height, + args.action_horizon, + test_video_reports, + "train", + ) + validate_output(test_root, "test") + + LOG.info("Conversion complete. DreamZero train root: %s", train_root) + if test_source: + LOG.info("Held-out test root: %s", output / "test") + return 0 + + +if __name__ == "__main__": + try: + sys.exit(main()) + except Exception: + LOG.exception("Conversion failed") + sys.exit(1) diff --git a/scripts/run_g2_server_9443_final.sh b/scripts/run_g2_server_9443_final.sh new file mode 100755 index 00000000..075dce4c --- /dev/null +++ b/scripts/run_g2_server_9443_final.sh @@ -0,0 +1,94 @@ +#!/bin/bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$ROOT_DIR" + +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY all_proxy ALL_PROXY +unset ws_proxy wss_proxy WS_PROXY WSS_PROXY + +export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}" +export PYTORCH_NVML_BASED_CUDA_CHECK="${PYTORCH_NVML_BASED_CUDA_CHECK:-1}" +export NO_ALBUMENTATIONS_UPDATE="${NO_ALBUMENTATIONS_UPDATE:-1}" +export ATTENTION_BACKEND="${ATTENTION_BACKEND:-FA2}" + +PYTHON_BIN="${PYTHON_BIN:-/data/wangk/conda/envs/dreamzero/bin/python}" + +PORT="${PORT:-9443}" + +# inference only supports 1 or 2 GPUs +NUM_GPUS=2 + +# 固定使用 GPU0 + GPU1 +CUDA_VISIBLE_DEVICES_VALUE="0,1" + + +# ============================== +# 新 action adapter checkpoint +# ============================== +MODEL_PATH="${MODEL_PATH:-/data/wangk/checkpoints/dreamzero_g2_nostate_v1_actvid_1to1_x001_3gpu/checkpoint-1500}" + + +WAN_CKPT_DIR="${WAN_CKPT_DIR:-/data/wangk/checkpoints/Wan2.1-I2V-14B-480P}" + +TOKENIZER_PATH="${TOKENIZER_PATH:-/data/wangk/checkpoints/umt5-xxl}" + +EMBODIMENT_TAG="g2" + +SERVER_MODULE="${SERVER_MODULE:-socket_optimized_AR_g2.py}" + + +VIDEO_SAVE_MODE="${VIDEO_SAVE_MODE:-full}" + +NUM_INFERENCE_TIMESTEPS="${NUM_INFERENCE_TIMESTEPS:-0}" + +OUTPUT_DIR="${OUTPUT_DIR:-/data/wangk/dreamzero/video_rollout_g2}" + + +if [ ! -x "$PYTHON_BIN" ]; then + echo "ERROR: python missing:" + echo "$PYTHON_BIN" + exit 1 +fi + + +if [ ! -d "$MODEL_PATH" ]; then + echo "ERROR checkpoint missing:" + echo "$MODEL_PATH" + exit 1 +fi + + +echo "==============================" +echo "DreamZero G2 Action Adapter" +echo "==============================" + +echo "checkpoint:" +echo "$MODEL_PATH" + +echo "GPU:" +echo "$CUDA_VISIBLE_DEVICES_VALUE" + + +echo "" +echo "Auditing checkpoint..." + +"$PYTHON_BIN" scripts/audit_g2_checkpoint.py "$MODEL_PATH" + + +echo "" +echo "Starting DreamZero server" + +CUDA_VISIBLE_DEVICES="$CUDA_VISIBLE_DEVICES_VALUE" \ +"$PYTHON_BIN" -m torch.distributed.run \ + --standalone \ + --nproc_per_node="$NUM_GPUS" \ + "$SERVER_MODULE" \ + --port "$PORT" \ + --model-path "$MODEL_PATH" \ + --wan-ckpt-dir "$WAN_CKPT_DIR" \ + --tokenizer-path "$TOKENIZER_PATH" \ + --embodiment-tag "$EMBODIMENT_TAG" \ + --video-save-mode "$VIDEO_SAVE_MODE" \ + --num-inference-timesteps "$NUM_INFERENCE_TIMESTEPS" \ + --output-dir "$OUTPUT_DIR" \ No newline at end of file diff --git a/scripts/train/train_dreamzero_g2_joint_lora.sh b/scripts/train/train_dreamzero_g2_joint_lora.sh new file mode 100644 index 00000000..d8e4ee0e --- /dev/null +++ b/scripts/train/train_dreamzero_g2_joint_lora.sh @@ -0,0 +1,244 @@ +#!/usr/bin/env bash +set -Eeuo pipefail + +# DreamZero G2 joint-space LoRA training. +# Point G2_DATA_ROOT at the GEAR train split, never at its parent/test folder. + +PROJECT_ROOT=${PROJECT_ROOT:-/home/ubuntu/projects/wangk/dreamzero} +G2_DATA_ROOT=${G2_DATA_ROOT:-/data/training_data/teleop/g2/g2_tasks_g1_g7_joint_gear_subtask_v2/train} + +OUTPUT_DIR=${OUTPUT_DIR:-/data/wangk/checkpoints/dreamzero_g2_joint_subtask_lora_v2} +WAN_CKPT_DIR=${WAN_CKPT_DIR:-/data/wangk/checkpoints/Wan2.1-I2V-14B-480P} +TOKENIZER_DIR=${TOKENIZER_DIR:-/data/wangk/checkpoints/umt5-xxl} +PRETRAINED_MODEL_PATH=${PRETRAINED_MODEL_PATH:-/data/wangk/checkpoints/DreamZero-AgiBot} + +# Only use physical GPUs 4 and above. GPU_IDS is preferred, while an existing +# CUDA_VISIBLE_DEVICES remains supported. The torchrun process count is always +# derived from the resulting list (for example, 6,7 means exactly 2 workers). +GPU_IDS=${GPU_IDS:-${CUDA_VISIBLE_DEVICES:-4,5,6,7}} +EXPECTED_EPISODES=${EXPECTED_EPISODES:-1346} +MAX_STEPS=${MAX_STEPS:-3000} +SAVE_STEPS=${SAVE_STEPS:-500} +WANDB_MODE=${WANDB_MODE:-offline} +HYDRA_FULL_ERROR=${HYDRA_FULL_ERROR:-1} + +export G2_DATA_ROOT OUTPUT_DIR WAN_CKPT_DIR TOKENIZER_DIR +export PRETRAINED_MODEL_PATH GPU_IDS +export EXPECTED_EPISODES MAX_STEPS SAVE_STEPS WANDB_MODE HYDRA_FULL_ERROR + +# Pin every Python entry point to the currently activated conda environment. +# This prevents ~/.local/bin/torchrun from launching /usr/bin/python3. +PYTHON_BIN="${PYTHON_BIN:-$(command -v python)}" +[[ -x "$PYTHON_BIN" ]] || { + echo "[ERROR] Python executable not found: $PYTHON_BIN" >&2 + exit 1 +} +export PYTHON_BIN + +fail() { + echo "[ERROR] $*" >&2 + exit 1 +} + +require_dir() { + [[ -d "$1" ]] || fail "Missing directory: $1" +} + +require_file() { + [[ -f "$1" ]] || fail "Missing file: $1" +} + +trap 'echo "[ERROR] Failed at line $LINENO" >&2' ERR + +require_dir "$PROJECT_ROOT" +require_dir "$G2_DATA_ROOT" +require_dir "$WAN_CKPT_DIR" +require_dir "$TOKENIZER_DIR" +require_dir "$PRETRAINED_MODEL_PATH" +require_file "$PROJECT_ROOT/groot/vla/experiment/experiment.py" +require_file "$PROJECT_ROOT/groot/vla/configs/data/dreamzero/g2_relative.yaml" +require_file "$PROJECT_ROOT/groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml" +require_file "$WAN_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth" +require_file "$WAN_CKPT_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" +require_file "$WAN_CKPT_DIR/Wan2.1_VAE.pth" +require_file "$G2_DATA_ROOT/meta/info.json" +require_file "$G2_DATA_ROOT/meta/modality.json" +require_file "$G2_DATA_ROOT/meta/embodiment.json" +require_file "$G2_DATA_ROOT/meta/stats.json" +require_file "$G2_DATA_ROOT/meta/relative_stats_dreamzero.json" + +"$PYTHON_BIN" - <<'PY' +import sys + +required = { + "hydra": "hydra-core", + "torch": "torch", + "omegaconf": "omegaconf", +} + +missing = [] +for module_name, package_name in required.items(): + try: + __import__(module_name) + except Exception: + missing.append(package_name) + +print("Python executable:", sys.executable) +print("Python version:", sys.version.replace("\n", " ")) + +if missing: + raise SystemExit( + "Missing packages in the active Python environment: " + + ", ".join(missing) + ) +PY + +# Keep GPU visibility and torchrun world size under one source of truth. This +# also prevents an inherited NUM_GPUS from launching more workers than visible +# devices. +IFS=',' read -r -a GPU_ARRAY <<< "$GPU_IDS" +(( ${#GPU_ARRAY[@]} > 0 )) || fail "GPU_IDS must contain at least one GPU index" + +declare -A SEEN_GPUS=() +for gpu in "${GPU_ARRAY[@]}"; do + [[ "$gpu" =~ ^[0-9]+$ ]] || fail "GPU_IDS must be comma-separated physical GPU indices; got: $GPU_IDS" + (( gpu >= 4 )) || fail "Refusing to use physical GPU $gpu: GPUs 0-3 are reserved" + [[ -z "${SEEN_GPUS[$gpu]+x}" ]] || fail "Duplicate GPU index in GPU_IDS: $gpu" + SEEN_GPUS[$gpu]=1 +done + +CUDA_VISIBLE_DEVICES=$(IFS=,; echo "${GPU_ARRAY[*]}") +NPROC_PER_NODE=${#GPU_ARRAY[@]} +export CUDA_VISIBLE_DEVICES + +# Fail before torchrun if any selected physical index does not exist. +for gpu in "${GPU_ARRAY[@]}"; do + nvidia-smi -i "$gpu" --query-gpu=index --format=csv,noheader >/dev/null \ + || fail "Physical GPU $gpu is unavailable" +done + +cd "$PROJECT_ROOT" + +echo "[1/2] Validating the G2 joint subtask training split" +"$PYTHON_BIN" - <<'PY' +import json +import os +from pathlib import Path + +root = Path(os.environ["G2_DATA_ROOT"]) +with (root / "meta/info.json").open() as handle: + info = json.load(handle) +with (root / "meta/embodiment.json").open() as handle: + embodiment = json.load(handle) + +episodes = int(info["total_episodes"]) +expected_episodes = int(os.environ["EXPECTED_EPISODES"]) +parquets = list(root.glob("data/chunk-*/*.parquet")) +videos = list(root.glob("videos/chunk-*/*/*.mp4")) + +if episodes != expected_episodes: + raise RuntimeError( + f"Expected {expected_episodes} training episodes, got {episodes}" + ) +if len(parquets) != episodes: + raise RuntimeError(f"Expected {episodes} parquets, got {len(parquets)}") +if len(videos) != episodes * 3: + raise RuntimeError(f"Expected {episodes * 3} videos, got {len(videos)}") +if info["features"]["observation.state"]["shape"] != [16]: + raise RuntimeError("observation.state must be 16-dimensional G2 joint state") +if info["features"]["action"]["shape"] != [16]: + raise RuntimeError("action must be 16-dimensional G2 joint action") +if embodiment.get("embodiment_tag") != "g2": + raise RuntimeError(f"Expected embodiment_tag=g2, got {embodiment}") + +video_features = { + key: feature + for key, feature in info["features"].items() + if feature.get("dtype") == "video" +} +expected_video_keys = { + "observation.images.top_head", + "observation.images.hand_left", + "observation.images.hand_right", +} +if set(video_features) != expected_video_keys: + raise RuntimeError(f"Unexpected video features: {sorted(video_features)}") +for key, feature in video_features.items(): + if feature["shape"] != [176, 320, 3]: + raise RuntimeError(f"Unexpected shape for {key}: {feature['shape']}") + +print( + f"Validation OK: {episodes} train episodes, " + f"{len(parquets)} parquets, {len(videos)} videos, G2 joint 16D" +) +PY + +mkdir -p "$OUTPUT_DIR" +TRAIN_LOG="$OUTPUT_DIR/train.log" + +echo "Physical GPUs selected: $CUDA_VISIBLE_DEVICES" +echo "Training processes: $NPROC_PER_NODE (one per selected GPU)" +nvidia-smi -i "$CUDA_VISIBLE_DEVICES" \ + --query-gpu=index,name,memory.total,memory.used,utilization.gpu \ + --format=csv + +echo "W&B mode: $WANDB_MODE" +echo "Checkpoint interval: every $SAVE_STEPS steps" +echo "[2/2] Starting ${MAX_STEPS}-step DreamZero G2 joint LoRA training" +CUDA_VISIBLE_DEVICES="$CUDA_VISIBLE_DEVICES" "$PYTHON_BIN" -m torch.distributed.run \ + --nproc_per_node "$NPROC_PER_NODE" \ + --standalone \ + groot/vla/experiment/experiment.py \ + report_to=wandb \ + data=dreamzero/g2_relative \ + wandb_project=dreamzero \ + train_architecture=lora \ + num_frames=33 \ + action_horizon=24 \ + num_views=3 \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frame_per_block=2 \ + num_action_per_block=24 \ + num_state_per_block=1 \ + seed=42 \ + training_args.learning_rate=1e-5 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \ + save_steps="$SAVE_STEPS" \ + training_args.warmup_ratio=0.05 \ + output_dir="$OUTPUT_DIR" \ + per_device_train_batch_size=1 \ + max_steps="$MAX_STEPS" \ + weight_decay=1e-5 \ + save_total_limit=5 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=false \ + dataloader_num_workers=1 \ + image_resolution_width=320 \ + image_resolution_height=176 \ + save_lora_only=true \ + max_chunk_size=4 \ + mixture_dataset_cls=groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec \ + single_dataset_cls=groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID \ + frame_seqlen=880 \ + save_strategy=steps \ + g2_data_root="$G2_DATA_ROOT" \ + dit_version="$WAN_CKPT_DIR" \ + text_encoder_pretrained_path="$WAN_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth" \ + image_encoder_pretrained_path="$WAN_CKPT_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" \ + vae_pretrained_path="$WAN_CKPT_DIR/Wan2.1_VAE.pth" \ + tokenizer_path="$TOKENIZER_DIR" \ + pretrained_model_path="$PRETRAINED_MODEL_PATH" \ + ++model_specific_transform.embodiment_tag_mapping.g2=26 \ + ++action_head_cfg.config.skip_component_loading=true \ + ++action_head_cfg.config.defer_lora_injection=true \ + 2>&1 | tee "$TRAIN_LOG" + +echo "Completed successfully" +echo "Dataset: $G2_DATA_ROOT" +echo "Training output: $OUTPUT_DIR" +echo "Training log: $TRAIN_LOG" \ No newline at end of file diff --git a/socket_optimized_AR_g2.py b/socket_optimized_AR_g2.py new file mode 100644 index 00000000..a476d630 --- /dev/null +++ b/socket_optimized_AR_g2.py @@ -0,0 +1,1248 @@ +import dataclasses +import logging +import socket +import asyncio +import os +import http +import logging +import time +import traceback +import torch +import tyro +from einops import rearrange +import datetime +import cv2 + +from groot.vla.model.n1_5.sim_policy import GrootSimPolicy +from groot.vla.data.schema import EmbodimentTag +import imageio +import numpy as np + +from openpi_client import base_policy as _base_policy +from openpi_client import msgpack_numpy +import websockets.asyncio.server as _server +import websockets.frames +from tianshou.data import Batch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh + +# Use roboarena policy server interface +from eval_utils.policy_server import WebsocketPolicyServer as RoboarenaServer +from eval_utils.policy_server import PolicyServerConfig + +logger = logging.getLogger(__name__) + +SIGNAL_INFER = 0 +SIGNAL_SHUTDOWN = 1 +SIGNAL_IDLE = 2 +SIGNAL_RESET_CACHE = 3 + + +def _run_policy_forward( + policy: GrootSimPolicy, + batch: Batch, +) -> tuple[object, torch.Tensor]: + """Select causal production or ordinary joint inference for A/B audits.""" + mode = os.environ.get( + "DREAMZERO_FORWARD_MODE", "causal" + ).strip().lower() + if mode == "causal": + return policy.lazy_joint_forward_causal(batch) + if mode == "joint": + return policy.lazy_joint_forward(batch) + raise ValueError( + "DREAMZERO_FORWARD_MODE must be 'causal' or 'joint', " + f"got {mode!r}" + ) + + +def _reset_policy_inference_cache(policy: object, reason: str) -> None: + trained_model = getattr(policy, "trained_model", None) + action_head = getattr(trained_model, "action_head", None) + reset_fn = getattr(action_head, "reset_inference_cache", None) + if callable(reset_fn): + reset_fn() + logger.info("Reset action-head inference cache on rank %s (%s)", dist.get_rank() if dist.is_initialized() else "?", reason) + else: + logger.warning("Policy action head does not expose reset_inference_cache(); cache reset skipped (%s)", reason) + +@dataclasses.dataclass +class Args: + port: int = 8000 + timeout_seconds: int = 50000 # 10 hours default, configurable + model_path: str = "./checkpoints/dreamzero" + wan_ckpt_dir: str | None = None # Override Wan2.1 component paths stored in config.json. + tokenizer_path: str | None = None # Override tokenizer_path stored in experiment_cfg/conf.yaml. + enable_dit_cache: bool = False # Backward-compatible alias for num_dit_steps=8. + num_dit_steps: int | None = None # Actual DiT compute steps. Supported fast masks: 5, 6, 7, 8. None keeps model default. + num_inference_timesteps: int | None = None # Positive values override diffusion steps; 0 keeps checkpoint default. + index: int = 0 + embodiment_tag: str = "oxe_droid" + max_chunk_size: int | None = None # If None, use config value. Otherwise override max_chunk_size for inference. + video_save_mode: str = "first" # one of: none, first, full. Controls generated video saved on reset/client close. + output_dir: str = "/data/wangk/dreamzero/video_rollout" + reset_cache_each_request: bool = True + + +class DistributedRoboarenaPolicyBase: + """Shared distributed inference plumbing for websocket policy wrappers.""" + + def __init__( + self, + groot_policy: GrootSimPolicy, + signal_group: dist.ProcessGroup, + output_dir: str | None = None, + video_save_mode: str = "first", + ) -> None: + self._policy = groot_policy + self._signal_group = signal_group + self._output_dir = output_dir + self._video_save_mode = video_save_mode + self._frame_buffers = self._init_frame_buffers() + self._current_session_id: str | None = None + self.video_across_time = [] + self._msg_index = 0 + + if self._output_dir: + os.makedirs(self._output_dir, exist_ok=True) + + def _init_frame_buffers(self) -> dict[str, list[np.ndarray]]: + return { + "video.top_head": [], + "video.hand_left": [], + "video.hand_right": [], + } + + def _reset_custom_state(self) -> None: + pass + + def _after_infer(self) -> None: + pass + + def _prepare_video_chunk(self, video_pred: torch.Tensor) -> torch.Tensor | None: + if self._video_save_mode == "none": + return None + return video_pred + + def _video_save_fps(self) -> int: + return 5 + + def _convert_observation(self, obs: dict) -> dict: + raise NotImplementedError + + def _convert_action(self, action_dict: dict) -> np.ndarray: + raise NotImplementedError + + def _broadcast_batch_to_workers(self, obs: dict) -> None: + import pickle + + serialized = pickle.dumps(obs) + data_size = len(serialized) + + size_tensor = torch.tensor([data_size], dtype=torch.int64, device='cuda') + dist.broadcast(size_tensor, src=0) + + data_tensor = torch.frombuffer(serialized, dtype=torch.uint8).clone().cuda() + dist.broadcast(data_tensor, src=0) + + def _extract_action_dict(self, action_chunk_dict: object) -> dict[str, object]: + action_dict: dict[str, object] = {} + for key in dir(action_chunk_dict): + if key.startswith('action.'): + action_dict[key] = getattr(action_chunk_dict, key) + return action_dict + + def _broadcast_signal_to_workers(self, signal: int) -> None: + signal_tensor = torch.tensor([signal], dtype=torch.int32, device='cpu') + dist.broadcast(signal_tensor, src=0, group=self._signal_group) + + def infer(self, obs: dict) -> np.ndarray: + session_id = obs.get('session_id') + if session_id is not None and session_id != self._current_session_id: + if self._current_session_id is not None: + logger.info("Session changed from '%s' to '%s', resetting state", self._current_session_id, session_id) + self._broadcast_signal_to_workers(SIGNAL_RESET_CACHE) + self._reset_state() + else: + logger.info("New session started: '%s'", session_id) + self._current_session_id = session_id + + if os.environ.get( + "DREAMZERO_RESET_AR_EACH_REQUEST", "true" + ).strip().lower() in {"1", "true", "yes", "on"}: + self._broadcast_signal_to_workers(SIGNAL_RESET_CACHE) + _reset_policy_inference_cache( + self._policy, + "fresh real observation before request", + ) + + self._msg_index += 1 + converted_obs = self._convert_observation(obs) + + self._broadcast_signal_to_workers(SIGNAL_INFER) + self._broadcast_batch_to_workers(converted_obs) + + batch = Batch(obs=converted_obs) + dist.barrier() + with torch.no_grad(): + result_batch, video_pred = _run_policy_forward( + self._policy, batch + ) + dist.barrier() + + video_chunk = self._prepare_video_chunk(video_pred) + if video_chunk is not None: + self.video_across_time.append(video_chunk.detach().cpu()) + action = self._convert_action(self._extract_action_dict(result_batch.act)) + self._after_infer() + return action + + def _reset_state(self, save_video: bool = True) -> None: + if save_video and len(self.video_across_time) > 0 and self._output_dir: + try: + frame_list = [] + action_head = self._policy.trained_model.action_head + device = getattr(action_head, "_device", None) + if device is None: + device = next(self._policy.trained_model.parameters()).device + video_across_time_cat = torch.cat(self.video_across_time, dim=2).to(device=device, dtype=torch.bfloat16) + frames = action_head.vae.decode( + video_across_time_cat, + tiled=action_head.tiled, + tile_size=(action_head.tile_size_height, action_head.tile_size_width), + tile_stride=(action_head.tile_stride_height, action_head.tile_stride_width), + ) + frames = rearrange(frames, 'B C T H W -> B T H W C') + frames = frames[0] + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + for frame in frames: + frame_list.append(frame) + + if frame_list: + sample_frame = frame_list[0] + if len(sample_frame.shape) == 3 and sample_frame.shape[2] in [1, 3, 4]: + save_dir = self._output_dir + os.makedirs(save_dir, exist_ok=True) + all_mp4_files = [f for f in os.listdir(save_dir) if f.endswith('.mp4')] + timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S') + num_frames = len(frame_list) + output_path = os.path.join(save_dir, f'{timestamp}_{len(all_mp4_files):06}_f{num_frames}.mp4') + imageio.mimsave(output_path, frame_list, fps=self._video_save_fps(), codec='libx264') + logger.info('Saved video on reset to: %s', output_path) + except Exception as exc: + logger.warning('Failed to save video on reset: %s', exc) + + for key in self._frame_buffers: + self._frame_buffers[key] = [] + + self.video_across_time = [] + _reset_policy_inference_cache(self._policy, "wrapper reset_state") + self._reset_custom_state() + + def reset(self, reset_info: dict) -> None: + self._broadcast_signal_to_workers(SIGNAL_RESET_CACHE) + self._reset_state(save_video=True) + + +class ARDroidRoboarenaPolicy(DistributedRoboarenaPolicyBase): + """Wrapper policy that implements roboarena.policy.BasePolicy interface for AR_droid.""" + + FRAMES_PER_CHUNK = 4 + + def __init__( + self, + groot_policy: GrootSimPolicy, + signal_group: dist.ProcessGroup, + output_dir: str | None = None, + video_save_mode: str = "first", + ) -> None: + super().__init__( + groot_policy=groot_policy, + signal_group=signal_group, + output_dir=output_dir, + video_save_mode=video_save_mode, + ) + self._reset_custom_state() + + def _init_frame_buffers(self) -> dict[str, list[np.ndarray]]: + return { + 'video.exterior_image_1_left': [], + 'video.exterior_image_2_left': [], + 'video.wrist_image_left': [], + } + + def _reset_custom_state(self) -> None: + self._is_first_call = True + + def _after_infer(self) -> None: + self._is_first_call = False + + def _convert_observation(self, obs: dict) -> dict: + converted = {} + image_key_mapping = { + 'observation/exterior_image_0_left': 'video.exterior_image_1_left', + 'observation/exterior_image_1_left': 'video.exterior_image_2_left', + 'observation/wrist_image_left': 'video.wrist_image_left', + } + + for roboarena_key, droid_key in image_key_mapping.items(): + if roboarena_key in obs: + data = obs[roboarena_key] + if isinstance(data, np.ndarray): + if data.ndim == 4: + self._frame_buffers[droid_key].extend(list(data)) + else: + self._frame_buffers[droid_key].append(data) + + num_frames = 1 if self._is_first_call else self.FRAMES_PER_CHUNK + + for droid_key, buffer in self._frame_buffers.items(): + if len(buffer) > 0: + if len(buffer) >= num_frames: + frames_to_use = buffer[-num_frames:] + else: + frames_to_use = buffer.copy() + while len(frames_to_use) < num_frames: + frames_to_use.insert(0, buffer[0]) + converted[droid_key] = np.stack(frames_to_use, axis=0) + + joint_pos = obs.get('observation/joint_position', np.zeros(7, dtype=np.float32)) + if joint_pos.ndim == 1: + joint_pos = joint_pos.reshape(1, -1) + converted['state.joint_position'] = joint_pos.astype(np.float64) + + gripper_pos = obs.get('observation/gripper_position', np.zeros(1, dtype=np.float32)) + if gripper_pos.ndim == 1: + gripper_pos = gripper_pos.reshape(1, -1) + converted['state.gripper_position'] = gripper_pos.astype(np.float64) + converted['annotation.language.action_text'] = obs.get('prompt', '') + return converted + + def _convert_action(self, action_dict: dict) -> np.ndarray: + joint_action = None + gripper_action = None + for key, value in action_dict.items(): + if 'joint_position' in key: + joint_action = value + elif 'gripper_position' in key or 'gripper' in key: + gripper_action = value + + if joint_action is None: + return np.zeros((1, 8), dtype=np.float32) + + if isinstance(joint_action, torch.Tensor): + joint_action = joint_action.cpu().numpy() + if joint_action.ndim == 1: + joint_action = joint_action.reshape(1, -1) + + num_steps = joint_action.shape[0] + if gripper_action is not None: + if isinstance(gripper_action, torch.Tensor): + gripper_action = gripper_action.cpu().numpy() + if gripper_action.ndim == 1: + gripper_action = gripper_action.reshape(-1, 1) + elif gripper_action.ndim == 0: + gripper_action = gripper_action.reshape(1, 1) + else: + gripper_action = np.zeros((num_steps, 1), dtype=np.float32) + + return np.concatenate([joint_action, gripper_action], axis=-1).astype(np.float32) + + +class AgiBotRoboarenaPolicy(DistributedRoboarenaPolicyBase): + """Adapter that converts websocket observations into AgiBot modality keys.""" + + VIDEO_KEY_MAPPING = { + 'observation/top_head': 'video.top_head', + 'observation/hand_left': 'video.hand_left', + 'observation/hand_right': 'video.hand_right', + } + STATE_KEY_MAPPING = { + 'observation/left_arm_joint_position': 'state.left_arm_joint_position', + 'observation/right_arm_joint_position': 'state.right_arm_joint_position', + 'observation/left_effector_position': 'state.left_effector_position', + 'observation/right_effector_position': 'state.right_effector_position', + 'observation/head_position': 'state.head_position', + 'observation/waist_pitch': 'state.waist_pitch', + 'observation/waist_lift': 'state.waist_lift', + } + + def __init__( + self, + groot_policy: GrootSimPolicy, + signal_group: dist.ProcessGroup, + output_dir: str | None = None, + video_save_mode: str = "first", + ) -> None: + super().__init__( + groot_policy=groot_policy, + signal_group=signal_group, + output_dir=output_dir, + video_save_mode=video_save_mode, + ) + self._action_keys = list(self._policy.modality_configs.action.modality_keys) + + def _lookup_obs_value(self, obs: dict, source_key: str, target_key: str) -> object: + if source_key in obs: + return obs[source_key] + return obs.get(target_key) + + def _normalize_video(self, value: object, target_key: str) -> np.ndarray: + if isinstance(value, dict) and value.get("__dreamzero_image_encoding__") == "jpeg_sequence": + frames = [] + expected_shape = tuple(value.get("shape", ())) + expected_dtype = np.dtype(value.get("dtype", "uint8")) + for index, frame_bytes in enumerate(value.get("frames", [])): + encoded = np.frombuffer(frame_bytes, dtype=np.uint8) + frame = cv2.imdecode(encoded, cv2.IMREAD_COLOR) + if frame is None: + raise ValueError(f"Failed to decode JPEG frame {index} for {target_key}") + frames.append(frame.astype(expected_dtype, copy=False)) + array = np.stack(frames, axis=0) + if expected_shape and tuple(array.shape) != expected_shape: + raise ValueError( + f"Decoded JPEG video for {target_key} has shape {array.shape}, expected {expected_shape}" + ) + return array + + array = np.asarray(value) + if array.ndim == 3: + return np.expand_dims(array, axis=0) + if array.ndim == 4: + return array + raise ValueError(f'AgiBot video input for {target_key} must have shape (H, W, C) or (T, H, W, C), got {array.shape}') + + def _normalize_state(self, value: object, target_key: str) -> np.ndarray: + array = np.asarray(value) + if array.ndim == 0: + return array.reshape(1, 1).astype(np.float64) + if array.ndim == 1: + return array.reshape(1, -1).astype(np.float64) + if array.ndim == 2: + return array.astype(np.float64) + raise ValueError(f'AgiBot state input for {target_key} must be 1D or 2D, got {array.shape}') + + def _prepare_video_chunk(self, video_pred: torch.Tensor) -> torch.Tensor | None: + if self._video_save_mode == "none": + return None + if video_pred.ndim != 5: + raise ValueError(f'AgiBot video prediction must be 5D (B, C, T, H, W), got {tuple(video_pred.shape)}') + if self._video_save_mode == "first": + return video_pred[:, :, :1].contiguous() + if self._video_save_mode == "full": + return video_pred.contiguous() + raise ValueError(f"Unsupported video_save_mode: {self._video_save_mode!r}; expected none, first, or full") + + def _video_save_fps(self) -> int: + return 20 + + def _convert_observation(self, obs: dict) -> dict: + converted = {} + missing_keys: list[str] = [] + + for source_key, target_key in self.VIDEO_KEY_MAPPING.items(): + value = self._lookup_obs_value(obs, source_key, target_key) + if value is None: + missing_keys.append(source_key) + continue + converted[target_key] = self._normalize_video(value, target_key) + + for source_key, target_key in self.STATE_KEY_MAPPING.items(): + value = self._lookup_obs_value(obs, source_key, target_key) + if value is None: + missing_keys.append(source_key) + continue + converted[target_key] = self._normalize_state(value, target_key) + + if missing_keys: + raise ValueError( + 'AgiBot inference requires the following observation keys: ' + + ', '.join(sorted(missing_keys)) + ) + + converted['annotation.language.action_text'] = obs.get('prompt', obs.get('annotation.language.action_text', '')) + return converted + + def _convert_action(self, action_dict: dict) -> np.ndarray: + flattened_chunks: list[np.ndarray] = [] + expected_horizon: int | None = None + missing_keys = [key for key in self._action_keys if key not in action_dict] + if missing_keys: + raise RuntimeError('Missing AgiBot action outputs: ' + ', '.join(missing_keys)) + + for action_key in self._action_keys: + value = action_dict[action_key] + if isinstance(value, torch.Tensor): + value = value.detach().cpu().numpy() + array = np.asarray(value) + if array.ndim == 0: + array = array.reshape(1, 1) + elif array.ndim == 1: + array = array.reshape(-1, 1) + else: + array = array.reshape(array.shape[0], -1) + + if expected_horizon is None: + expected_horizon = array.shape[0] + elif array.shape[0] != expected_horizon: + raise RuntimeError( + f'Inconsistent AgiBot action horizon for {action_key}: expected {expected_horizon}, got {array.shape[0]}' + ) + flattened_chunks.append(array.astype(np.float32)) + + return np.concatenate(flattened_chunks, axis=-1).astype(np.float32) + + +class G2RoboarenaPolicy(DistributedRoboarenaPolicyBase): + """Adapter for the G2 dual-arm joint-space policy.""" + FRAMES_PER_CHUNK = 4 + VIDEO_KEY_MAPPING = { + 'observation/top_head': 'video.top_head', + 'observation/hand_left': 'video.hand_left', + 'observation/hand_right': 'video.hand_right', + } + + STATE_KEY_MAPPING = { + 'observation/left_joint_position': 'state.left_joint_position', + 'observation/left_gripper_position': 'state.left_gripper_position', + 'observation/right_joint_position': 'state.right_joint_position', + 'observation/right_gripper_position': 'state.right_gripper_position', + } + + PACKED_STATE_KEYS = ( + 'observation/state', + 'observation.state', + 'state', + ) + + def __init__( + self, + groot_policy: GrootSimPolicy, + signal_group: dist.ProcessGroup, + output_dir: str | None = None, + video_save_mode: str = "first", + ) -> None: + super().__init__( + groot_policy=groot_policy, + signal_group=signal_group, + output_dir=output_dir, + video_save_mode=video_save_mode, + ) + self._action_keys = list( + self._policy.modality_configs.action.modality_keys + ) + + @staticmethod + def _lookup_obs_value( + obs: dict, + source_key: str, + target_key: str, + ) -> object: + if source_key in obs: + return obs[source_key] + return obs.get(target_key) + + @staticmethod + def _normalize_video( + value: object, + target_key: str, + ) -> np.ndarray: + if ( + isinstance(value, dict) + and value.get("__dreamzero_image_encoding__") + == "jpeg_sequence" + ): + frames = [] + expected_shape = tuple(value.get("shape", ())) + expected_dtype = np.dtype( + value.get("dtype", "uint8") + ) + for index, frame_bytes in enumerate( + value.get("frames", []) + ): + encoded = np.frombuffer( + frame_bytes, + dtype=np.uint8, + ) + frame = cv2.imdecode( + encoded, + cv2.IMREAD_COLOR, + ) + if frame is None: + raise ValueError( + f"Failed to decode JPEG frame {index} " + f"for {target_key}" + ) + # cv2.imdecode always returns BGR, while DreamZero training + # videos are decoded as RGB. Keep the model input contract RGB. + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frames.append( + frame.astype( + expected_dtype, + copy=False, + ) + ) + if not frames: + raise ValueError( + f"No JPEG frames were provided for {target_key}" + ) + array = np.stack(frames, axis=0) + if expected_shape and tuple(array.shape) != expected_shape: + raise ValueError( + f"Decoded JPEG video for {target_key} has " + f"shape {array.shape}, expected {expected_shape}" + ) + return array + + array = np.asarray(value) + if array.ndim == 3: + return np.expand_dims(array, axis=0) + if array.ndim == 4: + return array + raise ValueError( + f"G2 video input for {target_key} must have " + f"shape (H,W,C) or (T,H,W,C), got {array.shape}" + ) + + @staticmethod + def _normalize_state( + value: object, + target_key: str, + ) -> np.ndarray: + array = np.asarray(value, dtype=np.float64) + if array.ndim == 0: + array = array.reshape(1, 1) + elif array.ndim == 1: + array = array.reshape(1, -1) + elif array.ndim != 2: + raise ValueError( + f"G2 state input for {target_key} must be " + f"1D or 2D, got {array.shape}" + ) + return array + + @staticmethod + def _split_packed_state( + value: object, + ) -> dict[str, np.ndarray]: + packed = np.asarray(value, dtype=np.float64) + if packed.ndim == 1: + packed = packed.reshape(1, -1) + elif packed.ndim != 2: + raise ValueError( + "Packed G2 state must have shape (16,) or " + f"(T,16), got {packed.shape}" + ) + if packed.shape[-1] != 16: + raise ValueError( + f"Packed G2 state must contain 16 values, " + f"got {packed.shape}" + ) + return { + 'state.left_joint_position': packed[:, 0:7], + 'state.left_gripper_position': packed[:, 7:8], + 'state.right_joint_position': packed[:, 8:15], + 'state.right_gripper_position': packed[:, 15:16], + } + + def _prepare_video_chunk( + self, + video_pred: torch.Tensor, + ) -> torch.Tensor | None: + if self._video_save_mode == "none": + return None + if video_pred.ndim != 5: + raise ValueError( + "G2 video prediction must be 5D " + f"(B,C,T,H,W), got {tuple(video_pred.shape)}" + ) + if self._video_save_mode == "first": + return video_pred[:, :, :1].contiguous() + if self._video_save_mode == "full": + return video_pred.contiguous() + raise ValueError( + f"Unsupported video_save_mode: " + f"{self._video_save_mode!r}" + ) + + def _video_save_fps(self) -> int: + return 30 + + def _convert_observation(self, obs: dict) -> dict: + converted: dict[str, object] = {} + missing_video: list[str] = [] + + for source_key, target_key in self.VIDEO_KEY_MAPPING.items(): + value = self._lookup_obs_value( + obs, + source_key, + target_key, + ) + if value is None: + missing_video.append(source_key) + continue + frames = self._normalize_video( + value, + target_key, + ) + + self._frame_buffers[target_key].extend( + list(frames) + ) + + history = self._frame_buffers[target_key][-4:] + + while len(history) < 4: + history.insert(0, history[0]) + + converted[target_key] = np.stack(history, axis=0) + + if missing_video: + raise ValueError( + "G2 inference requires video keys: " + + ", ".join(sorted(missing_video)) + ) + + packed_state = None + for key in self.PACKED_STATE_KEYS: + if key in obs: + packed_state = obs[key] + break + + if packed_state is not None: + converted.update( + self._split_packed_state(packed_state) + ) + else: + missing_state: list[str] = [] + for source_key, target_key in self.STATE_KEY_MAPPING.items(): + value = self._lookup_obs_value( + obs, + source_key, + target_key, + ) + if value is None: + missing_state.append(source_key) + continue + converted[target_key] = self._normalize_state( + value, + target_key, + ) + if missing_state: + raise ValueError( + "G2 inference requires a packed 16-D state " + "under observation/state, observation.state, " + "or state; otherwise all split state keys are " + "required: " + + ", ".join(sorted(missing_state)) + ) + + expected_dims = { + 'state.left_joint_position': 7, + 'state.left_gripper_position': 1, + 'state.right_joint_position': 7, + 'state.right_gripper_position': 1, + } + for key, expected_dim in expected_dims.items(): + array = np.asarray(converted[key]) + if array.shape[-1] != expected_dim: + raise ValueError( + f"{key} must have last dimension " + f"{expected_dim}, got {array.shape}" + ) + + converted['annotation.language.action_text'] = obs.get( + 'prompt', + obs.get( + 'annotation.language.action_text', + '', + ), + ) + return converted + + def _convert_action( + self, + action_dict: dict, + ) -> np.ndarray: + missing = [ + key + for key in self._action_keys + if key not in action_dict + ] + if missing: + raise RuntimeError( + "Missing G2 action outputs: " + + ", ".join(missing) + ) + + arrays: list[np.ndarray] = [] + horizon: int | None = None + for key in self._action_keys: + value = action_dict[key] + if isinstance(value, torch.Tensor): + value = value.detach().cpu().numpy() + array = np.asarray(value) + if array.ndim == 0: + array = array.reshape(1, 1) + elif array.ndim == 1: + array = array.reshape(-1, 1) + else: + array = array.reshape(array.shape[0], -1) + + if horizon is None: + horizon = array.shape[0] + elif array.shape[0] != horizon: + raise RuntimeError( + f"Inconsistent G2 action horizon for {key}: " + f"expected {horizon}, got {array.shape[0]}" + ) + arrays.append(array.astype(np.float32)) + + action = np.concatenate(arrays, axis=-1) + if action.shape != (24, 16): + raise RuntimeError( + "G2 action must have shape (24,16), ordered as " + "[left_joint(7), left_gripper(1), " + "right_joint(7), right_gripper(1)], " + f"got {action.shape}" + ) + if not np.isfinite(action).all(): + raise RuntimeError("G2 action contains NaN or infinity") + if float(np.max(np.abs(action))) < 1e-6: + raise RuntimeError( + "G2 action is entirely zero; refusing a likely " + "checkpoint/config loading failure" + ) + return action.astype(np.float32) + + +class WebsocketPolicyServer: + """Serves a policy using the websocket protocol. See websocket_client_policy.py for a client implementation. + Currently only implements the `load` and `infer` methods. + """ + + def __init__( + self, + policy: _base_policy.BasePolicy, + host: str = "0.0.0.0", + port: int | None = None, + metadata: dict | None = None, + output_dir: str | None = None, + signal_group: dist.ProcessGroup | None = None, + ) -> None: + self._policy = policy + self._host = host + self._port = port + self._metadata = metadata or {} + self._output_dir = output_dir + logging.getLogger("websockets.server").setLevel(logging.INFO) + self.video_across_time = [] + self._msg_index = 0 + self._signal_group = signal_group + if self._output_dir: + os.makedirs(self._output_dir, exist_ok=True) + os.makedirs(os.path.join(self._output_dir, "inputs"), exist_ok=True) + + def serve_forever(self, rank: int = 0) -> None: + asyncio.run(self.run(rank)) + + async def run(self, rank: int = 0): + if rank == 0: + async with _server.serve( + self._handler, + self._host, + self._port, + compression=None, + max_size=None, + process_request=_health_check, + ping_interval=None, + ) as server: + await server.serve_forever() + else: + await self._worker_loop() + + async def _worker_loop(self): + logger.info(f"Worker loop started for rank {dist.get_rank()}") + signal_tensor = torch.zeros(1, dtype=torch.int32, device='cpu') + while True: + try: + dist.broadcast(signal_tensor, src=0, group=self._signal_group) + + signal = signal_tensor.item() + if signal == SIGNAL_SHUTDOWN: + logger.info(f"Rank {dist.get_rank()} received shutdown signal") + break + elif signal == SIGNAL_IDLE: + logger.info(f"Rank {dist.get_rank()} received idle signal. Waiting for next client.") + continue + elif signal == SIGNAL_RESET_CACHE: + logger.info(f"Rank {dist.get_rank()} received inference cache reset signal") + _reset_policy_inference_cache(self._policy, "worker signal") + continue + + batch = self._receive_batch_from_rank0() + dist.barrier() + with torch.no_grad(): + result_batch, video_pred = _run_policy_forward( + self._policy, batch + ) + dist.barrier() + + except Exception as e: + logger.error(f"Worker loop error on rank {dist.get_rank()}: {e}") + traceback.print_exc() + break + + def _receive_batch_from_rank0(self): + import pickle + + size_tensor = torch.zeros(1, dtype=torch.int64, device='cuda') + dist.broadcast(size_tensor, src=0) + data_size = size_tensor.item() + + data_tensor = torch.zeros(data_size, dtype=torch.uint8, device='cuda') + dist.broadcast(data_tensor, src=0) + + obs = pickle.loads(data_tensor.cpu().numpy().tobytes()) + return Batch(obs=obs) + + def _broadcast_batch_to_workers(self, obs): + import pickle + + serialized = pickle.dumps(obs) + data_size = len(serialized) + + size_tensor = torch.tensor([data_size], dtype=torch.int64, device='cuda') + dist.broadcast(size_tensor, src=0) + + data_tensor = torch.frombuffer(serialized, dtype=torch.uint8).clone().cuda() + dist.broadcast(data_tensor, src=0) + + async def _handler(self, websocket: _server.ServerConnection): + logger.info(f"Connection from {websocket.remote_address} opened") + packer = msgpack_numpy.Packer() + + await websocket.send(packer.pack(self._metadata)) + + signal_tensor = torch.zeros(1, dtype=torch.int32, device='cpu') + + try: + while True: + try: + data = await websocket.recv() + obs = msgpack_numpy.unpackb(data) + self._msg_index += 1 + + signal_tensor.zero_() + dist.broadcast(signal_tensor, src=0, group=self._signal_group) + + self._broadcast_batch_to_workers(obs) + batch = Batch(obs=obs) + + dist.barrier() + with torch.no_grad(): + result_batch, video_pred = _run_policy_forward( + self._policy, batch + ) + dist.barrier() + + action_chunk_dict = result_batch.act + + def batch_to_dict(batch): + out = {} + for k in dir(batch): + if not k.startswith("action."): + continue + out[k] = getattr(batch, k) + return out + + action_chunk_dict = batch_to_dict(action_chunk_dict) + await websocket.send(packer.pack(action_chunk_dict)) + + except websockets.ConnectionClosed: + logger.info(f"Connection from {websocket.remote_address} closed") + self.video_across_time = [] + break + except Exception: + await websocket.send(traceback.format_exc()) + await websocket.close( + code=websockets.frames.CloseCode.INTERNAL_ERROR, + reason="Internal server error. Traceback included in previous frame.", + ) + raise + finally: + logger.info("Rank 0: Client session ended. Sending idle signal (2) to workers.") + signal_tensor.fill_(2) + dist.broadcast(signal_tensor, src=0, group=self._signal_group) + + +def init_mesh() -> DeviceMesh: + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + + torch.cuda.set_device(local_rank) + _ = torch.cuda.is_available() + _ = torch.cuda.device_count() + + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size not in (1, 2): + raise ValueError( + f"This DreamZero inference path only supports 1 or 2 GPUs, got world_size={world_size}. " + "The action head parallelization code explicitly supports ip_size 1 or 2 only. " + "Please launch with --nproc_per_node=2 (or 1)." + ) + print(f"Rank {rank}/{world_size} (PID: {os.getpid()}) setting device to local_rank={local_rank}") + + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + + mesh = init_device_mesh( + device_type="cuda", + mesh_shape=(world_size,), + mesh_dim_names=("ip",), + ) + print(f"Rank {rank}/{world_size} (PID: {os.getpid()}) using device {device}") + + return mesh + +def _health_check(connection: _server.ServerConnection, request: _server.Request) -> _server.Response | None: + if request.path == "/healthz": + return connection.respond(http.HTTPStatus.OK, "OK\n") + return None + + +def _create_wrapper_policy( + embodiment_tag: str, + groot_policy: GrootSimPolicy, + signal_group: dist.ProcessGroup, + output_dir: str | None, + video_save_mode: str, +) -> DistributedRoboarenaPolicyBase: + if embodiment_tag == 'oxe_droid': + return ARDroidRoboarenaPolicy( + groot_policy=groot_policy, + signal_group=signal_group, + output_dir=output_dir, + video_save_mode=video_save_mode, + ) + if embodiment_tag == 'agibot': + return AgiBotRoboarenaPolicy( + groot_policy=groot_policy, + signal_group=signal_group, + output_dir=output_dir, + video_save_mode=video_save_mode, + ) + if embodiment_tag == 'g2': + return G2RoboarenaPolicy( + groot_policy=groot_policy, + signal_group=signal_group, + output_dir=output_dir, + video_save_mode=video_save_mode, + ) + raise ValueError(f'Unsupported embodiment_tag: {embodiment_tag}') + + +def _create_server_config(embodiment_tag: str) -> PolicyServerConfig: + if embodiment_tag == 'oxe_droid': + return PolicyServerConfig( + image_resolution=(180, 320), + needs_wrist_camera=True, + n_external_cameras=2, + needs_stereo_camera=False, + needs_session_id=True, + action_space='joint_position', + ) + if embodiment_tag == 'agibot': + return PolicyServerConfig( + image_resolution=(640, 480), + needs_wrist_camera=False, + n_external_cameras=3, + needs_stereo_camera=False, + needs_session_id=True, + action_space='agibot_flattened', + ) + if embodiment_tag == 'g2': + return PolicyServerConfig( + image_resolution=(176, 320), + needs_wrist_camera=False, + n_external_cameras=3, + needs_stereo_camera=False, + needs_session_id=True, + action_space='joint_position', + ) + raise ValueError(f'Unsupported embodiment_tag: {embodiment_tag}') + + +def _build_path_overrides(args: Args) -> tuple[list[str], list[str]]: + model_config_overrides: list[str] = [] + train_config_overrides: list[str] = [] + + if args.wan_ckpt_dir: + wan_ckpt_dir = os.path.abspath(args.wan_ckpt_dir) + required_files = [ + os.path.join(wan_ckpt_dir, "models_t5_umt5-xxl-enc-bf16.pth"), + os.path.join(wan_ckpt_dir, "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"), + os.path.join(wan_ckpt_dir, "Wan2.1_VAE.pth"), + ] + missing = [path for path in required_files if not os.path.exists(path)] + if missing: + raise FileNotFoundError( + "Missing Wan checkpoint component(s): " + ", ".join(missing) + ) + model_config_overrides.extend( + [ + f"action_head_cfg.config.diffusion_model_cfg.diffusion_model_pretrained_path={wan_ckpt_dir}", + f"action_head_cfg.config.text_encoder_cfg.text_encoder_pretrained_path={wan_ckpt_dir}/models_t5_umt5-xxl-enc-bf16.pth", + f"action_head_cfg.config.image_encoder_cfg.image_encoder_pretrained_path={wan_ckpt_dir}/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth", + f"action_head_cfg.config.vae_cfg.vae_pretrained_path={wan_ckpt_dir}/Wan2.1_VAE.pth", + ] + ) + + if args.tokenizer_path: + tokenizer_path = os.path.abspath(args.tokenizer_path) + if not os.path.exists(tokenizer_path): + raise FileNotFoundError(f"Tokenizer path does not exist: {tokenizer_path}") + if args.embodiment_tag.lower() == "agibot": + train_config_overrides.append( + f"transforms.agibot.transforms.10.tokenizer_path={tokenizer_path}" + ) + elif args.embodiment_tag.lower() == "oxe_droid": + train_config_overrides.append( + f"transforms.oxe_droid.transforms.10.tokenizer_path={tokenizer_path}" + ) + elif args.embodiment_tag.lower() == "g2": + train_config_overrides.append( + f"transforms.g2.transforms.10.tokenizer_path={tokenizer_path}" + ) + + return model_config_overrides, train_config_overrides + + +def main(args: Args) -> None: + os.environ["DREAMZERO_RESET_AR_EACH_REQUEST"] = ( + "true" if args.reset_cache_each_request else "false" + ) + os.environ["ENABLE_DIT_CACHE"] = "true" if args.enable_dit_cache else "false" + if args.num_dit_steps is not None: + os.environ["NUM_DIT_STEPS"] = str(args.num_dit_steps) + elif args.enable_dit_cache: + os.environ.setdefault("NUM_DIT_STEPS", "8") + os.environ.setdefault("ATTENTION_BACKEND", "FA2") + if args.video_save_mode not in {"none", "first", "full"}: + raise ValueError(f"--video-save-mode must be one of none, first, full; got {args.video_save_mode!r}") + torch._dynamo.config.recompile_limit = 800 + + embodiment_tag = args.embodiment_tag.lower() + if embodiment_tag not in {'oxe_droid', 'agibot', 'g2'}: + raise ValueError(f'Unsupported embodiment_tag: {args.embodiment_tag}') + model_path = args.model_path + model_config_overrides, train_config_overrides = _build_path_overrides(args) + policy_metadata = { + "embodiment": embodiment_tag, + "model_name": "dreamzero", + "model_path": model_path, + "wan_ckpt_dir": args.wan_ckpt_dir, + "tokenizer_path": args.tokenizer_path, + } + + device_mesh = init_mesh() + rank = dist.get_rank() + + timeout_delta = datetime.timedelta(seconds=args.timeout_seconds) + signal_group = dist.new_group(backend="gloo", timeout=timeout_delta) + logger.info(f"Rank {rank} initialized signal_group (gloo)") + + policy = GrootSimPolicy( + embodiment_tag=EmbodimentTag(embodiment_tag), + model_path=model_path, + device="cuda" if torch.cuda.is_available() else "cpu", + device_mesh=device_mesh, + model_config_overrides=model_config_overrides, + train_config_overrides=train_config_overrides, + ) + action_head = policy.trained_model.action_head + if args.num_inference_timesteps is not None: + if args.num_inference_timesteps == 0: + logging.info("Keeping checkpoint diffusion inference steps because --num-inference-timesteps=0") + elif args.num_inference_timesteps < 0: + raise ValueError( + f"--num-inference-timesteps must be non-negative, got {args.num_inference_timesteps}" + ) + else: + action_head.num_inference_steps = int(args.num_inference_timesteps) + action_head.num_inference_timesteps = int(args.num_inference_timesteps) + if hasattr(action_head, "config"): + action_head.config.num_inference_timesteps = int(args.num_inference_timesteps) + logging.info( + "Overrode action_head diffusion inference steps to %s on rank %s", + args.num_inference_timesteps, + rank, + ) + logging.info( + "[CONFIG CHECK] rank=%s action_head.num_inference_steps=%s " + "action_head.num_inference_timesteps=%s action_head.num_frame_per_block=%s " + "action_head.model.num_frame_per_block=%s NUM_DIT_STEPS=%s ENABLE_DIT_CACHE=%s", + rank, + getattr(action_head, "num_inference_steps", None), + getattr(action_head, "num_inference_timesteps", None), + getattr(action_head, "num_frame_per_block", None), + getattr(getattr(action_head, "model", None), "num_frame_per_block", None), + os.getenv("NUM_DIT_STEPS"), + os.getenv("ENABLE_DIT_CACHE"), + ) + + hostname = socket.gethostname() + local_ip = socket.gethostbyname(hostname) + + if rank == 0: + logging.info("Creating server (host: %s, ip: %s)", hostname, local_ip) + output_dir = None if args.video_save_mode == "none" else args.output_dir + if output_dir is not None: + os.makedirs(output_dir, exist_ok=True) + logging.info("Videos will be saved to: %s", output_dir) + else: + logging.info("Video saving disabled; no output directory will be created.") + else: + output_dir = None + logging.info(f"Rank {rank} starting as worker for distributed inference...") + + wrapper_policy = _create_wrapper_policy( + embodiment_tag=embodiment_tag, + groot_policy=policy, + signal_group=signal_group, + output_dir=output_dir, + video_save_mode=args.video_save_mode, + ) + + server_config = _create_server_config(embodiment_tag) + + if rank == 0: + logging.info("Using roboarena policy server interface for %s", embodiment_tag) + logging.info(f"Server config: {server_config}") + roboarena_server = RoboarenaServer( + policy=wrapper_policy, + server_config=server_config, + host="0.0.0.0", + port=args.port, + ) + roboarena_server.serve_forever() + else: + server = WebsocketPolicyServer( + policy=policy, + host="0.0.0.0", + port=args.port, + metadata=policy_metadata, + output_dir=output_dir, + signal_group=signal_group, + ) + asyncio.run(server._worker_loop()) + + +def cli() -> None: + logging.basicConfig(level=logging.INFO, force=True) + main(tyro.cli(Args)) + + +if __name__ == "__main__": + cli()