wehub-resource-sync
130 lines
4.1 KiB
Python
130 lines
4.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""DreamZero AR client frame schedule aligned with upstream ``test_client_AR.py``."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
# Frame schedule constants (matching upstream debug_inference.py / test_client_AR.py).
|
|
ACTION_HORIZON = 24
|
|
DEFAULT_ACTION_DIM = 8
|
|
RELATIVE_OFFSETS = [-23, -16, -8, 0]
|
|
DEFAULT_NUM_AR_CHUNKS = 15
|
|
|
|
CAMERA_FILES = {
|
|
"observation/exterior_image_0_left": "exterior_image_1_left.mp4",
|
|
"observation/exterior_image_1_left": "exterior_image_2_left.mp4",
|
|
"observation/wrist_image_left": "wrist_image_left.mp4",
|
|
}
|
|
|
|
|
|
def build_frame_schedule(
|
|
total_frames: int,
|
|
num_chunks: int,
|
|
*,
|
|
logger: logging.Logger | None = None,
|
|
) -> list[list[int]]:
|
|
"""Build the frame index schedule for multi-frame chunks.
|
|
|
|
Args:
|
|
total_frames: Number of frames available in each camera video.
|
|
num_chunks: Number of 4-frame chunks to schedule after the initial frame.
|
|
|
|
Returns:
|
|
A list of frame-index lists. Each inner list has four indices. The
|
|
returned list may be shorter than ``num_chunks`` when the videos run
|
|
out of frames (upstream stops early instead of erroring).
|
|
"""
|
|
chunks: list[list[int]] = []
|
|
current_frame = 23
|
|
for _ in range(num_chunks):
|
|
indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
|
|
if indices[-1] >= total_frames:
|
|
if logger is not None:
|
|
logger.info(
|
|
"Frame %s >= %s, stopping at %s chunks",
|
|
indices[-1],
|
|
total_frames,
|
|
len(chunks),
|
|
)
|
|
break
|
|
chunks.append(indices)
|
|
current_frame += ACTION_HORIZON
|
|
return chunks
|
|
|
|
|
|
def make_obs_from_video(
|
|
camera_frames: dict[str, np.ndarray],
|
|
frame_indices: list[int],
|
|
*,
|
|
prompt: str,
|
|
session_id: str,
|
|
) -> dict[str, Any]:
|
|
"""Build an observation dict from real video frames.
|
|
|
|
For one frame each image key is ``(H, W, 3)``. For four frames each key is
|
|
``(4, H, W, 3)``.
|
|
"""
|
|
obs: dict[str, Any] = {}
|
|
for camera_key, all_frames in camera_frames.items():
|
|
selected = all_frames[frame_indices]
|
|
obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
|
|
|
|
obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
|
|
obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
|
|
obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
|
|
obs["prompt"] = prompt
|
|
obs["session_id"] = session_id
|
|
return obs
|
|
|
|
|
|
def build_ar_observations(
|
|
camera_frames: dict[str, np.ndarray],
|
|
*,
|
|
prompt: str,
|
|
session_id: str,
|
|
num_chunks: int = DEFAULT_NUM_AR_CHUNKS,
|
|
repeat_chunk_observations: bool = False,
|
|
logger: logging.Logger | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""Build the AR observation sequence used by upstream ``test_client_AR.py``.
|
|
|
|
Step 0 sends frame ``[0]``. Each subsequent step sends one 4-frame chunk.
|
|
|
|
``num_chunks`` counts only the 4-frame chunks after the initial frame, not
|
|
the total number of inferences.
|
|
"""
|
|
if num_chunks < 0:
|
|
raise ValueError("num_chunks must be non-negative")
|
|
|
|
total_frames = min(frames.shape[0] for frames in camera_frames.values())
|
|
observations = [
|
|
make_obs_from_video(
|
|
camera_frames,
|
|
[0],
|
|
prompt=prompt,
|
|
session_id=session_id,
|
|
)
|
|
]
|
|
|
|
chunk_schedule = build_frame_schedule(total_frames, num_chunks, logger=logger)
|
|
if repeat_chunk_observations and chunk_schedule and len(chunk_schedule) < num_chunks:
|
|
while len(chunk_schedule) < num_chunks:
|
|
chunk_schedule.append(chunk_schedule[-1])
|
|
|
|
for frame_indices in chunk_schedule:
|
|
observations.append(
|
|
make_obs_from_video(
|
|
camera_frames,
|
|
frame_indices,
|
|
prompt=prompt,
|
|
session_id=session_id,
|
|
)
|
|
)
|
|
return observations
|