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>
112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Pydantic schemas for RL rollout serving (RFC #3747).
|
|
|
|
P0 scope: world_model_env mode only (observation + action -> next observation).
|
|
|
|
P0 maps rollout observations/actions onto DreamZero's DROID robot_obs schema:
|
|
camera images retain their DreamZero observation/* keys, joint/gripper state is
|
|
passed through observation/joint_position and observation/gripper_position, and
|
|
the returned video payload is the DreamZero VAE latent unless a later endpoint
|
|
adds explicit latent decoding.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Literal
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class Observation(BaseModel):
|
|
"""Single-step robot observation."""
|
|
|
|
images: dict[str, list] | None = Field(
|
|
default=None,
|
|
description="Named DreamZero camera images as nested numeric lists. Base64 strings are not accepted in P0.",
|
|
)
|
|
state: list[float] | None = Field(
|
|
default=None,
|
|
description="Proprioceptive state vector (joint positions, velocities, etc.).",
|
|
)
|
|
prompt: str = Field(default="", description="Optional language conditioning.")
|
|
extra: dict[str, Any] = Field(
|
|
default_factory=dict,
|
|
description="Pass-through fields forwarded verbatim to robot_obs.",
|
|
)
|
|
|
|
|
|
class Action(BaseModel):
|
|
"""Executed action for world_model_env conditioning."""
|
|
|
|
joint_positions: list[float] | None = Field(
|
|
default=None,
|
|
description="Executed joint positions; forwarded as observation/joint_position for DreamZero.",
|
|
)
|
|
gripper_position: float | list[float] | None = Field(
|
|
default=None,
|
|
description="Executed gripper position; forwarded as observation/gripper_position for DreamZero.",
|
|
)
|
|
extra: dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class SessionMetadata(BaseModel):
|
|
latency_ms: float
|
|
steps_generated: int
|
|
context_length: int
|
|
committed_step_id: int = Field(
|
|
description="Highest step_id whose context has been atomically committed. "
|
|
"-1 means no step has been committed yet (fresh or reset session).",
|
|
)
|
|
session_memory_bytes: int | None = None
|
|
uncertainty: float | None = None
|
|
|
|
|
|
class ErrorObject(BaseModel):
|
|
code: str
|
|
message: str
|
|
step_id: int | None = None
|
|
committed_step_id: int = -1
|
|
|
|
|
|
class CreateSessionRequest(BaseModel):
|
|
model: str
|
|
mode: Literal["world_model_env"] = "world_model_env"
|
|
|
|
|
|
class CreateSessionResponse(BaseModel):
|
|
session_id: str
|
|
mode: str
|
|
created_at: float
|
|
|
|
|
|
class RolloutStepRequest(BaseModel):
|
|
step_id: int = Field(ge=0, description="Monotonically increasing per session.")
|
|
observation: Observation
|
|
action: Action = Field(
|
|
description="Required for world_model_env: executed action that produced this observation.",
|
|
)
|
|
use_session_context: bool = True
|
|
|
|
|
|
class RolloutStepResponse(BaseModel):
|
|
step_id: int
|
|
next_observation: dict[str, Any] | None = Field(
|
|
default=None,
|
|
description="Predicted next observation metadata. P0 returns DreamZero video_latent, not RGB frames.",
|
|
)
|
|
model_metadata: SessionMetadata
|
|
error: ErrorObject | None = None
|
|
|
|
|
|
class ResetSessionResponse(BaseModel):
|
|
session_id: str
|
|
committed_step_id: int
|
|
|
|
|
|
class SessionStatusResponse(BaseModel):
|
|
session_id: str
|
|
committed_step_id: int
|
|
context_length: int
|
|
closed: bool
|