8e738ce21b
Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>
371 lines
13 KiB
Python
371 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Serving layer for RL rollout (RFC #3747, P0).
|
|
|
|
P0 scope: world_model_env mode, single-session, DreamZero backbone.
|
|
|
|
P0 maps client rollout input onto DreamZero's DROID robot_obs schema and returns
|
|
DreamZero video latents explicitly. RGB/frame decoding can be added as a later
|
|
contract without pretending the latent is already an image observation.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import binascii
|
|
import time
|
|
from collections import abc
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from vllm.logger import init_logger
|
|
from vllm.utils import random_uuid
|
|
|
|
from vllm_omni.entrypoints.openai.protocol.rollout import (
|
|
Action,
|
|
CreateSessionRequest,
|
|
CreateSessionResponse,
|
|
ErrorObject,
|
|
Observation,
|
|
ResetSessionResponse,
|
|
RolloutStepRequest,
|
|
RolloutStepResponse,
|
|
SessionMetadata,
|
|
SessionStatusResponse,
|
|
)
|
|
from vllm_omni.entrypoints.openai.rollout_session import (
|
|
RolloutSession,
|
|
RolloutSessionClosedError,
|
|
RolloutSessionNotFoundError,
|
|
RolloutSessionStore,
|
|
)
|
|
from vllm_omni.entrypoints.openpi.serving import ServingRealtimeRobotOpenPI
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _merge_action_into_obs(obs: Observation, action: Action | None) -> dict[str, Any]:
|
|
"""Build DreamZero-compatible robot_obs."""
|
|
robot_obs: dict[str, Any] = {"prompt": obs.prompt}
|
|
|
|
if obs.images:
|
|
robot_obs.update(obs.images)
|
|
|
|
if obs.extra:
|
|
for key, value in obs.extra.items():
|
|
if key.startswith("observation/") or key in {"seed", "embodiment_name"}:
|
|
robot_obs[key] = value
|
|
|
|
joint_positions = action.joint_positions if action is not None else None
|
|
gripper_position = action.gripper_position if action is not None else None
|
|
|
|
if joint_positions is None and obs.state:
|
|
joint_positions = obs.state[:7]
|
|
if len(obs.state) > 7 and gripper_position is None:
|
|
gripper_position = obs.state[7:8]
|
|
|
|
if joint_positions is not None:
|
|
joint_arr = np.asarray(joint_positions, dtype=np.float64).flatten()
|
|
if joint_arr.size != 7:
|
|
raise ValueError(f"DreamZero DROID joint_positions must have 7 values, got {joint_arr.size}.")
|
|
robot_obs["observation/joint_position"] = joint_arr
|
|
|
|
if gripper_position is not None:
|
|
gripper_arr = np.asarray(gripper_position, dtype=np.float64).flatten()
|
|
if gripper_arr.size != 1:
|
|
raise ValueError(f"DreamZero DROID gripper_position must have 1 value, got {gripper_arr.size}.")
|
|
robot_obs["observation/gripper_position"] = gripper_arr
|
|
|
|
return robot_obs
|
|
|
|
|
|
def _encode_video_output(video: Any) -> dict[str, Any]:
|
|
"""Encode DreamZero video latent tensor / ndarray to JSON metadata."""
|
|
if video is None:
|
|
return {}
|
|
if hasattr(video, "detach") and callable(video.detach):
|
|
video = video.detach()
|
|
if hasattr(video, "float") and callable(video.float):
|
|
video = video.float()
|
|
if hasattr(video, "cpu") and callable(video.cpu):
|
|
video = video.cpu()
|
|
if hasattr(video, "numpy"):
|
|
video = video.numpy()
|
|
if isinstance(video, np.ndarray):
|
|
return {
|
|
"video_latent": binascii.b2a_base64(video.tobytes(), newline=False).decode(),
|
|
"shape": list(video.shape),
|
|
"dtype": str(video.dtype),
|
|
"encoding": "base64_raw_tensor",
|
|
}
|
|
return {"video_latent": str(video), "encoding": "string"}
|
|
|
|
|
|
class ServingRLRollout:
|
|
"""HTTP serving layer for RL rollout sessions.
|
|
|
|
Wraps ServingRealtimeRobotOpenPI for world_model_env mode.
|
|
Session state and committed_step_id tracking live in RolloutSessionStore.
|
|
"""
|
|
|
|
def __init__(self, openpi_serving: ServingRealtimeRobotOpenPI) -> None:
|
|
self._openpi = openpi_serving
|
|
self._store = RolloutSessionStore()
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Session lifecycle #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
async def create_session(self, req: CreateSessionRequest) -> CreateSessionResponse:
|
|
self._validate_model(req.model)
|
|
session_id = random_uuid()
|
|
session = await self._store.create(
|
|
session_id=session_id,
|
|
model=req.model,
|
|
mode=req.mode,
|
|
)
|
|
logger.info("Created rollout session %s mode=%s", session_id, req.mode)
|
|
return CreateSessionResponse(
|
|
session_id=session_id,
|
|
mode=session.mode,
|
|
created_at=session.created_at,
|
|
)
|
|
|
|
async def reset_session(self, session_id: str) -> ResetSessionResponse:
|
|
session = await self._store.reset(session_id)
|
|
# Next infer still passes reset=True; drop GPU KV now if the engine
|
|
# session is reachable from this process.
|
|
self._drop_engine_session(session_id)
|
|
logger.info("Reset rollout session %s", session_id)
|
|
return ResetSessionResponse(
|
|
session_id=session_id,
|
|
committed_step_id=session.committed_step_id,
|
|
)
|
|
|
|
async def close_session(self, session_id: str) -> None:
|
|
await self._store.close(session_id)
|
|
self._drop_engine_session(session_id)
|
|
self._drop_engine_session(f"{session_id}:stateless")
|
|
logger.info("Closed rollout session %s", session_id)
|
|
|
|
async def get_status(self, session_id: str) -> SessionStatusResponse:
|
|
session = await self._store.get(session_id, include_closed=True)
|
|
return SessionStatusResponse(
|
|
session_id=session_id,
|
|
committed_step_id=session.committed_step_id,
|
|
context_length=session.context_length,
|
|
closed=session.closed,
|
|
)
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Step #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
async def step(
|
|
self,
|
|
session_id: str,
|
|
req: RolloutStepRequest,
|
|
) -> RolloutStepResponse:
|
|
try:
|
|
session = await self._store.get(session_id)
|
|
except RolloutSessionNotFoundError:
|
|
return self._error_response(
|
|
req.step_id,
|
|
-1,
|
|
0,
|
|
"session_not_found",
|
|
f"Session {session_id!r} does not exist.",
|
|
)
|
|
except RolloutSessionClosedError:
|
|
return self._error_response(
|
|
req.step_id,
|
|
-1,
|
|
0,
|
|
"session_closed",
|
|
f"Session {session_id!r} is closed.",
|
|
)
|
|
|
|
async with session.lock:
|
|
return await self._run_step(session, req)
|
|
|
|
async def _run_step(
|
|
self,
|
|
session: RolloutSession,
|
|
req: RolloutStepRequest,
|
|
) -> RolloutStepResponse:
|
|
t0 = time.perf_counter()
|
|
committed = session.committed_step_id
|
|
|
|
if req.use_session_context:
|
|
expected_step_id = committed + 1
|
|
if req.step_id != expected_step_id:
|
|
if req.step_id <= committed:
|
|
code = "step_already_committed"
|
|
message = f"Step {req.step_id} is already committed; highest committed step is {committed}."
|
|
else:
|
|
code = "step_out_of_order"
|
|
message = f"Expected step_id {expected_step_id}, got {req.step_id}."
|
|
return self._error_response(
|
|
req.step_id,
|
|
committed,
|
|
session.context_length,
|
|
code,
|
|
message,
|
|
)
|
|
|
|
# reset=True on first call after session create/reset
|
|
reset = committed == -1 or not req.use_session_context
|
|
engine_session_id = session.session_id
|
|
if not req.use_session_context:
|
|
engine_session_id = f"{session.session_id}:stateless"
|
|
try:
|
|
robot_obs = _merge_action_into_obs(req.observation, req.action)
|
|
video_out = await self._infer_world_model(
|
|
obs=robot_obs,
|
|
session_id=engine_session_id,
|
|
reset=reset,
|
|
)
|
|
next_observation = _encode_video_output(video_out)
|
|
except ValueError as exc:
|
|
logger.info("Invalid rollout step %d for session %s: %s", req.step_id, session.session_id, exc)
|
|
return self._error_response(
|
|
req.step_id,
|
|
committed,
|
|
session.context_length,
|
|
"invalid_request",
|
|
str(exc),
|
|
)
|
|
except Exception as exc:
|
|
logger.exception("Step %d failed for session %s", req.step_id, session.session_id)
|
|
# Do NOT advance committed_step_id on failure (RFC section 6.5).
|
|
return self._error_response(
|
|
req.step_id,
|
|
committed,
|
|
session.context_length,
|
|
"inference_error",
|
|
str(exc),
|
|
)
|
|
|
|
if req.use_session_context:
|
|
# Only advance committed_step_id on successful contextual steps.
|
|
await self._store.advance(session.session_id, req.step_id)
|
|
committed = req.step_id
|
|
|
|
latency_ms = (time.perf_counter() - t0) * 1000.0
|
|
metadata = SessionMetadata(
|
|
latency_ms=round(latency_ms, 2),
|
|
steps_generated=1,
|
|
context_length=session.context_length,
|
|
committed_step_id=committed,
|
|
)
|
|
return RolloutStepResponse(
|
|
step_id=req.step_id,
|
|
next_observation=next_observation,
|
|
model_metadata=metadata,
|
|
)
|
|
|
|
async def _infer_world_model(
|
|
self,
|
|
obs: dict[str, Any],
|
|
*,
|
|
session_id: str,
|
|
reset: bool,
|
|
) -> Any:
|
|
"""Call the engine and return the video (next-observation) output.
|
|
|
|
Reuses ServingRealtimeRobotOpenPI.build_request() so request routing,
|
|
session_id threading, and OmniDiffusionSamplingParams construction are
|
|
identical to the policy_inference path. DreamZero's formatter peels
|
|
``video`` into ``result.images`` and leaves only ``actions`` on
|
|
``multimodal_output``; read images first and fall back to the raw key.
|
|
"""
|
|
request = self._openpi.build_request(obs, session_id=session_id, reset=reset)
|
|
result = None
|
|
async for output in self._openpi.engine_client.generate(
|
|
prompt=self._request_prompt(request),
|
|
request_id=request.request_id,
|
|
sampling_params_list=[request.sampling_params],
|
|
):
|
|
result = output
|
|
|
|
if result is None:
|
|
raise RuntimeError("World model request produced no output.")
|
|
return self._extract_video_output(result)
|
|
|
|
@staticmethod
|
|
def _request_prompt(request: Any) -> Any:
|
|
prompt = getattr(request, "prompt", None)
|
|
if prompt is not None:
|
|
return prompt
|
|
return request.prompts[0]
|
|
|
|
@staticmethod
|
|
def _extract_video_output(result: Any) -> Any:
|
|
"""Prefer formatter-placed ``images``; fall back to multimodal video."""
|
|
images = getattr(result, "images", None)
|
|
if isinstance(images, (list, tuple)):
|
|
if images:
|
|
return images[0]
|
|
elif images is not None and not isinstance(images, abc.Mapping):
|
|
# Formatter assigns the peeled video tensor directly to images.
|
|
return images
|
|
|
|
multimodal_output = getattr(result, "multimodal_output", None)
|
|
if isinstance(multimodal_output, abc.Mapping):
|
|
video = multimodal_output.get("video")
|
|
if video is not None:
|
|
return video
|
|
|
|
raise RuntimeError(
|
|
"No video on result.images or multimodal_output['video']. "
|
|
"Confirm DreamZero formatter placement for world_model_env."
|
|
)
|
|
|
|
def _drop_engine_session(self, session_id: str) -> None:
|
|
drop = getattr(self._openpi, "drop_session", None)
|
|
if not callable(drop):
|
|
return
|
|
try:
|
|
drop(session_id)
|
|
except Exception:
|
|
logger.exception("Failed to drop engine session %s", session_id)
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Helpers #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
def _validate_model(self, requested_model: str) -> None:
|
|
served_model = getattr(self._openpi, "model_name", None)
|
|
if served_model is None:
|
|
return
|
|
if requested_model == served_model:
|
|
return
|
|
if requested_model.lower() == "dreamzero" and "dreamzero" in served_model.lower():
|
|
return
|
|
raise ValueError(f"Requested rollout model {requested_model!r} does not match served model {served_model!r}.")
|
|
|
|
@staticmethod
|
|
def _error_response(
|
|
step_id: int,
|
|
committed_step_id: int,
|
|
context_length: int,
|
|
code: str,
|
|
message: str,
|
|
) -> RolloutStepResponse:
|
|
metadata = SessionMetadata(
|
|
latency_ms=0.0,
|
|
steps_generated=0,
|
|
context_length=context_length,
|
|
committed_step_id=committed_step_id,
|
|
)
|
|
return RolloutStepResponse(
|
|
step_id=step_id,
|
|
next_observation=None,
|
|
model_metadata=metadata,
|
|
error=ErrorObject(
|
|
code=code,
|
|
message=message,
|
|
step_id=step_id,
|
|
committed_step_id=committed_step_id,
|
|
),
|
|
)
|