221 lines
6.8 KiB
Python
221 lines
6.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Action-token helpers for Cosmos3 action generation.
|
|
|
|
These helpers cover action modes that use action tokens as an auxiliary output
|
|
stream: ``policy``, ``forward_dynamics``, and ``inverse_dynamics``. The pipeline
|
|
returns predicted actions through the diffusion output payload/metadata
|
|
envelope rather than as a video frame stream.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
ACTION_MODE_POLICY = "policy"
|
|
ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics"
|
|
ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics"
|
|
ACTION_MODES = {
|
|
ACTION_MODE_POLICY,
|
|
ACTION_MODE_FORWARD_DYNAMICS,
|
|
ACTION_MODE_INVERSE_DYNAMICS,
|
|
}
|
|
|
|
|
|
EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = {
|
|
"no_action": 0,
|
|
"av": 1,
|
|
"camera_pose": 2,
|
|
"hand_pose": 3,
|
|
"pusht": 4,
|
|
"libero": 5,
|
|
"umi": 6,
|
|
"bridge_orig_lerobot": 7,
|
|
"droid_lerobot": 8,
|
|
"robomind-franka": 8,
|
|
"galbot": 9,
|
|
"robomind-franka-dual": 12,
|
|
"robomind-ur": 13,
|
|
"agibotworld": 15,
|
|
"agibot_gear_gripper": 15,
|
|
"agibot_gear_gripper_ext": 15,
|
|
"fractal": 20,
|
|
}
|
|
|
|
|
|
VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = {
|
|
"256": {
|
|
"1,1": (256, 256),
|
|
"4,3": (320, 256),
|
|
"3,4": (256, 320),
|
|
"16,9": (320, 192),
|
|
"9,16": (192, 320),
|
|
},
|
|
"480": {
|
|
"1,1": (640, 640),
|
|
"4,3": (736, 544),
|
|
"3,4": (544, 736),
|
|
"16,9": (832, 480),
|
|
"9,16": (480, 832),
|
|
},
|
|
"704": {
|
|
"1,1": (960, 960),
|
|
"4,3": (1088, 832),
|
|
"3,4": (832, 1088),
|
|
"16,9": (1280, 704),
|
|
"9,16": (704, 1280),
|
|
},
|
|
"720": {
|
|
"1,1": (960, 960),
|
|
"4,3": (1104, 832),
|
|
"3,4": (832, 1104),
|
|
"16,9": (1280, 720),
|
|
"9,16": (720, 1280),
|
|
},
|
|
}
|
|
|
|
|
|
def normalize_action_mode(mode: Any) -> str | None:
|
|
if mode is None:
|
|
return None
|
|
normalized = str(mode).strip().lower()
|
|
if not normalized:
|
|
return None
|
|
if normalized not in ACTION_MODES:
|
|
raise ValueError(f"Unsupported Cosmos3 action_mode={mode!r}; expected one of {sorted(ACTION_MODES)}.")
|
|
return normalized
|
|
|
|
|
|
def resolve_domain_id(
|
|
*,
|
|
domain_id: Any = None,
|
|
domain_name: Any = None,
|
|
require_explicit: bool = False,
|
|
) -> int:
|
|
if domain_id is not None:
|
|
resolved = int(domain_id)
|
|
if resolved < 0:
|
|
raise ValueError(f"Cosmos3 domain_id must be non-negative, got {resolved}.")
|
|
return resolved
|
|
|
|
if domain_name is None or str(domain_name).strip() == "":
|
|
if require_explicit:
|
|
raise ValueError(
|
|
"Cosmos3 action generation requires extra_args['domain_id'] or non-empty extra_args['domain_name']."
|
|
)
|
|
return 0
|
|
|
|
key = str(domain_name).strip().lower()
|
|
if key not in EMBODIMENT_TO_DOMAIN_ID:
|
|
raise ValueError(
|
|
f"Unknown Cosmos3 action domain_name={domain_name!r}; "
|
|
f"expected one of {sorted(EMBODIMENT_TO_DOMAIN_ID)} or pass domain_id directly."
|
|
)
|
|
return EMBODIMENT_TO_DOMAIN_ID[key]
|
|
|
|
|
|
def action_condition_indexes(mode: str, action_length: int) -> list[int]:
|
|
mode = normalize_action_mode(mode)
|
|
if mode == ACTION_MODE_FORWARD_DYNAMICS:
|
|
return list(range(action_length))
|
|
if mode in {ACTION_MODE_POLICY, ACTION_MODE_INVERSE_DYNAMICS}:
|
|
return []
|
|
raise AssertionError(f"Unexpected action mode: {mode!r}")
|
|
|
|
|
|
def vision_condition_indexes(mode: str, video_length: int, temporal_compression_factor: int) -> list[int]:
|
|
mode = normalize_action_mode(mode)
|
|
latent_frames = (video_length - 1) // temporal_compression_factor + 1
|
|
if mode in {ACTION_MODE_POLICY, ACTION_MODE_FORWARD_DYNAMICS}:
|
|
return [0]
|
|
if mode == ACTION_MODE_INVERSE_DYNAMICS:
|
|
return list(range(latent_frames))
|
|
raise AssertionError(f"Unexpected action mode: {mode!r}")
|
|
|
|
|
|
def action_start_frame_offset(mode: str, action_length: int, video_length: int) -> int:
|
|
del mode
|
|
if action_length == video_length - 1:
|
|
return 1
|
|
if action_length == video_length:
|
|
return 0
|
|
raise ValueError(
|
|
"Cosmos3 action_chunk_size must equal num_frames - 1 or num_frames; "
|
|
f"got action_chunk_size={action_length}, num_frames={video_length}."
|
|
)
|
|
|
|
|
|
def build_action_condition_mask(
|
|
mode: str,
|
|
action_length: int,
|
|
*,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
) -> torch.Tensor:
|
|
mask = torch.zeros(1, action_length, 1, device=device, dtype=dtype)
|
|
for idx in action_condition_indexes(mode, action_length):
|
|
mask[:, idx, :] = 1.0
|
|
return mask
|
|
|
|
|
|
def build_vision_condition_mask(
|
|
mode: str,
|
|
video_length: int,
|
|
temporal_compression_factor: int,
|
|
*,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
) -> torch.Tensor:
|
|
latent_frames = (video_length - 1) // temporal_compression_factor + 1
|
|
mask = torch.zeros(1, 1, latent_frames, 1, 1, device=device, dtype=dtype)
|
|
for idx in vision_condition_indexes(mode, video_length, temporal_compression_factor):
|
|
mask[:, :, idx, :, :] = 1.0
|
|
return mask
|
|
|
|
|
|
def pad_action_to_dim(action: torch.Tensor, action_dim: int) -> torch.Tensor:
|
|
if action.shape[-1] > action_dim:
|
|
raise ValueError(f"Cosmos3 action dimension {action.shape[-1]} exceeds model action_dim={action_dim}.")
|
|
if action.shape[-1] == action_dim:
|
|
return action
|
|
padding = torch.zeros(*action.shape[:-1], action_dim - action.shape[-1], dtype=action.dtype, device=action.device)
|
|
return torch.cat([action, padding], dim=-1)
|
|
|
|
|
|
def load_action_tensor(action: Any = None) -> torch.Tensor:
|
|
if action is None:
|
|
raise ValueError(
|
|
"Cosmos3 forward_dynamics action mode requires extra_args['action'] "
|
|
"(a JSON array / nested list / tensor of shape [T, D])."
|
|
)
|
|
if isinstance(action, torch.Tensor):
|
|
tensor = action.detach().to(dtype=torch.float32)
|
|
else:
|
|
tensor = torch.as_tensor(np.asarray(action), dtype=torch.float32)
|
|
if tensor.ndim == 3 and tensor.shape[0] == 1:
|
|
tensor = tensor.squeeze(0)
|
|
if tensor.ndim != 2:
|
|
raise ValueError(f"Cosmos3 action must have shape [T, D], got {tuple(tensor.shape)}.")
|
|
return tensor
|
|
|
|
|
|
def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]:
|
|
key = str(resolution)
|
|
if key not in VIDEO_RES_SIZE_INFO:
|
|
raise ValueError(
|
|
f"Unknown Cosmos3 action resolution={resolution!r}; expected one of {sorted(VIDEO_RES_SIZE_INFO)}."
|
|
)
|
|
input_ratio = h / w
|
|
best_size = None
|
|
best_diff = float("inf")
|
|
for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values():
|
|
diff = abs(input_ratio - cand_h / cand_w)
|
|
if diff < best_diff:
|
|
best_diff = diff
|
|
best_size = (cand_w, cand_h)
|
|
assert best_size is not None
|
|
return best_size
|