[P0][RFC #3747] RL rollout serving: session management + world_model_env (#4770)

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>

tests/entrypoints/openai_api/test_api_server_guards.py

@@ -80,6 +80,11 @@ _EXPECTED_ROUTER_ROUTES = {
     ("WEBSOCKET", "/v1/realtime"),
     ("WEBSOCKET", "/v1/realtime/robot/openpi"),
     ("WEBSOCKET", "/v1/duplex"),
+    ("POST", "/v1/realtime/sessions"),
+    ("POST", "/v1/realtime/sessions/{session_id}/step"),
+    ("POST", "/v1/realtime/sessions/{session_id}/reset"),
+    ("POST", "/v1/realtime/sessions/{session_id}/close"),
+    ("GET", "/v1/realtime/sessions/{session_id}/status"),
     ("GET", "/health"),
     ("GET", "/v1/models"),
     ("POST", "/v1/images/generations"),
@@ -120,6 +125,7 @@ _DIFFUSION_APP_STATE_KEYS = {
     "openai_streaming_speech",
     "openai_streaming_video",
     "openai_serving_realtime_robot",
+    "rl_rollout_serving",
     "enable_server_load_tracking",
     "server_load_metrics",
 }
@@ -151,12 +157,14 @@ _MULTISTAGE_APP_STATE_KEYS = {
     "openai_serving_realtime",
     "openai_serving_video",
     "openai_serving_realtime_robot",
+    "rl_rollout_serving",
     "enable_server_load_tracking",
     "server_load_metrics",
 }
 _MULTISTAGE_MUST_BE_NONE = {
     "openai_serving_duplex",
     "openai_serving_realtime_robot",
+    "rl_rollout_serving",
 }
 _MULTISTAGE_MUST_BE_WIRED = _MULTISTAGE_APP_STATE_KEYS - _MULTISTAGE_MUST_BE_NONE
 

tests/entrypoints/openai_api/test_openpi_serving.py

@@ -440,3 +440,17 @@ def test_extract_actions_does_not_iterate_result_object():
     actions = serving._extract_actions(IterableResult())
 
     np.testing.assert_allclose(actions, np.array([[1.0, 2.0, 3.0]], dtype=np.float32))
+
+
+def test_drop_session_calls_dreamzero_close_hook():
+    closed: list[str] = []
+    pipeline = SimpleNamespace(close_ar_diffusion_session=lambda session_id: closed.append(session_id))
+    engine_client = SimpleNamespace(
+        model_config={"policy_server_config": TEST_POLICY_SERVER_CONFIG},
+        model_runner=SimpleNamespace(pipeline=pipeline),
+    )
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client)
+
+    serving.drop_session("rollout-1")
+
+    assert closed == ["rollout-1"]

tests/rl_rollout/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

tests/rl_rollout/test_rollout_routes.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Route-level tests for RL rollout serving endpoints."""
+
+from __future__ import annotations
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from vllm_omni.entrypoints.openai.api_server import router
+from vllm_omni.entrypoints.openai.protocol.rollout import (
+    ErrorObject,
+    RolloutStepResponse,
+    SessionMetadata,
+)
+from vllm_omni.entrypoints.openai.rollout_session import (
+    RolloutSessionClosedError,
+    RolloutSessionNotFoundError,
+)
+
+
+class _FakeRolloutServing:
+    def __init__(self, close_error: Exception | None = None, step_error_code: str | None = None) -> None:
+        self.close_error = close_error
+        self.step_error_code = step_error_code
+        self.step_request = None
+
+    async def step(self, session_id, body):
+        self.step_request = body
+        error = None
+        if self.step_error_code is not None:
+            error = ErrorObject(
+                code=self.step_error_code,
+                message=f"{self.step_error_code}: {session_id}",
+                step_id=body.step_id,
+                committed_step_id=-1,
+            )
+        return RolloutStepResponse(
+            step_id=body.step_id,
+            next_observation=None,
+            model_metadata=SessionMetadata(
+                latency_ms=0.0,
+                steps_generated=0,
+                context_length=0,
+                committed_step_id=-1,
+            ),
+            error=error,
+        )
+
+    async def close_session(self, session_id):
+        if self.close_error is not None:
+            raise self.close_error
+
+
+def _client(serving: _FakeRolloutServing | None = None) -> TestClient:
+    app = FastAPI()
+    app.include_router(router)
+    if serving is not None:
+        app.state.rl_rollout_serving = serving
+    return TestClient(app)
+
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_step_route_rejects_base64_image_payload():
+    serving = _FakeRolloutServing()
+    client = _client(serving)
+
+    response = client.post(
+        "/v1/realtime/sessions/s1/step",
+        json={
+            "step_id": 0,
+            "observation": {"images": {"front": "base64-frames"}},
+            "action": {},
+        },
+    )
+
+    assert response.status_code == 422
+    assert serving.step_request is None
+
+
+def test_step_route_accepts_nested_image_payload():
+    serving = _FakeRolloutServing()
+    client = _client(serving)
+
+    response = client.post(
+        "/v1/realtime/sessions/s1/step",
+        json={
+            "step_id": 0,
+            "observation": {"images": {"observation/exterior_image_1_left": [[[0, 0, 0]]]}},
+            "action": {},
+        },
+    )
+
+    assert response.status_code == 200
+    assert serving.step_request.observation.images == {"observation/exterior_image_1_left": [[[0, 0, 0]]]}
+
+
+def test_step_route_requires_action():
+    client = _client(_FakeRolloutServing())
+
+    response = client.post(
+        "/v1/realtime/sessions/s1/step",
+        json={"step_id": 0, "observation": {}},
+    )
+
+    assert response.status_code == 422
+
+
+def test_step_route_maps_missing_session_to_404():
+    client = _client(_FakeRolloutServing(step_error_code="session_not_found"))
+
+    response = client.post(
+        "/v1/realtime/sessions/s1/step",
+        json={"step_id": 0, "observation": {}, "action": {}},
+    )
+
+    assert response.status_code == 404
+
+
+def test_step_route_maps_closed_session_to_410():
+    client = _client(_FakeRolloutServing(step_error_code="session_closed"))
+
+    response = client.post(
+        "/v1/realtime/sessions/s1/step",
+        json={"step_id": 0, "observation": {}, "action": {}},
+    )
+
+    assert response.status_code == 410
+
+
+def test_step_route_maps_invalid_request_to_400():
+    client = _client(_FakeRolloutServing(step_error_code="invalid_request"))
+
+    response = client.post(
+        "/v1/realtime/sessions/s1/step",
+        json={"step_id": 0, "observation": {}, "action": {}},
+    )
+
+    assert response.status_code == 400
+
+
+def test_close_route_maps_missing_session_to_404():
+    client = _client(_FakeRolloutServing(close_error=RolloutSessionNotFoundError("s1")))
+
+    response = client.post("/v1/realtime/sessions/s1/close")
+
+    assert response.status_code == 404
+
+
+def test_close_route_is_idempotent_for_closed_session():
+    client = _client(_FakeRolloutServing(close_error=RolloutSessionClosedError("s1")))
+
+    response = client.post("/v1/realtime/sessions/s1/close")
+
+    assert response.status_code == 200
+    assert response.json()["closed"] is True
+
+
+def test_rollout_routes_return_501_when_serving_is_unavailable():
+    client = _client()
+
+    response = client.post("/v1/realtime/sessions/s1/close")
+
+    assert response.status_code == 501

tests/rl_rollout/test_rollout_session.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for RolloutSessionStore (RFC #3747, P0)."""
+
+import asyncio
+
+import pytest
+
+from vllm_omni.entrypoints.openai.rollout_session import (
+    RolloutSession,
+    RolloutSessionClosedError,
+    RolloutSessionNotFoundError,
+    RolloutSessionStepError,
+    RolloutSessionStore,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def store() -> RolloutSessionStore:
+    return RolloutSessionStore()
+
+
+@pytest.mark.asyncio
+async def test_create_returns_session(store):
+    s = await store.create("s1", model="dreamzero", mode="world_model_env")
+    assert isinstance(s, RolloutSession)
+    assert s.session_id == "s1"
+    assert s.committed_step_id == -1
+    assert s.context_length == 0
+    assert not s.closed
+
+
+@pytest.mark.asyncio
+async def test_get_existing(store):
+    await store.create("s1", model="dreamzero", mode="world_model_env")
+    s = await store.get("s1")
+    assert s.session_id == "s1"
+
+
+@pytest.mark.asyncio
+async def test_get_missing_raises(store):
+    with pytest.raises(RolloutSessionNotFoundError):
+        await store.get("does-not-exist")
+
+
+@pytest.mark.asyncio
+async def test_get_closed_raises(store):
+    await store.create("s1", model="dreamzero", mode="world_model_env")
+    await store.close("s1")
+    with pytest.raises(RolloutSessionClosedError):
+        await store.get("s1")
+
+
+@pytest.mark.asyncio
+async def test_close_missing_raises(store):
+    with pytest.raises(RolloutSessionNotFoundError):
+        await store.close("does-not-exist")
+
+
+@pytest.mark.asyncio
+async def test_close_is_idempotent(store):
+    await store.create("s1", model="dreamzero", mode="world_model_env")
+    await store.close("s1")
+    await store.close("s1")
+    s = await store.get("s1", include_closed=True)
+    assert s.closed
+
+
+@pytest.mark.asyncio
+async def test_advance_updates_committed_step_id(store):
+    await store.create("s1", model="dreamzero", mode="world_model_env")
+    await store.advance("s1", step_id=0)
+    s = await store.get("s1")
+    assert s.committed_step_id == 0
+    assert s.context_length == 1
+
+    await store.advance("s1", step_id=1)
+    s = await store.get("s1")
+    assert s.committed_step_id == 1
+    assert s.context_length == 2
+
+
+@pytest.mark.asyncio
+async def test_advance_rejects_duplicate_step_id(store):
+    await store.create("s1", model="dreamzero", mode="world_model_env")
+    await store.advance("s1", step_id=0)
+    with pytest.raises(RolloutSessionStepError):
+        await store.advance("s1", step_id=0)
+
+    s = await store.get("s1")
+    assert s.committed_step_id == 0
+    assert s.context_length == 1
+
+
+@pytest.mark.asyncio
+async def test_advance_rejects_out_of_order_step_id(store):
+    await store.create("s1", model="dreamzero", mode="world_model_env")
+    with pytest.raises(RolloutSessionStepError):
+        await store.advance("s1", step_id=2)
+
+    s = await store.get("s1")
+    assert s.committed_step_id == -1
+    assert s.context_length == 0
+
+
+@pytest.mark.asyncio
+async def test_reset_clears_state(store):
+    await store.create("s1", model="dreamzero", mode="world_model_env")
+    await store.advance("s1", step_id=0)
+    await store.advance("s1", step_id=1)
+    await store.advance("s1", step_id=2)
+    await store.advance("s1", step_id=3)
+    await store.reset("s1")
+    s = await store.get("s1")
+    assert s.committed_step_id == -1
+    assert s.context_length == 0
+
+
+@pytest.mark.asyncio
+async def test_reset_closed_raises(store):
+    await store.create("s1", model="dreamzero", mode="world_model_env")
+    await store.close("s1")
+    with pytest.raises(RolloutSessionClosedError):
+        await store.reset("s1")
+
+
+@pytest.mark.asyncio
+async def test_failed_step_does_not_advance(store):
+    """committed_step_id must not change when a step raises (RFC section 6.5 invariant)."""
+    await store.create("s1", model="dreamzero", mode="world_model_env")
+    await store.advance("s1", step_id=0)
+
+    # Simulate a failed step: advance is NOT called
+    s = await store.get("s1")
+    assert s.committed_step_id == 0  # unchanged
+
+
+@pytest.mark.asyncio
+async def test_concurrent_steps_are_serialised(store):
+    """Per-session lock must prevent concurrent context mutations."""
+    await store.create("s1", model="dreamzero", mode="world_model_env")
+    s = await store.get("s1")
+
+    results: list[int] = []
+
+    async def slow_advance(step_id: int, delay: float) -> None:
+        await asyncio.sleep(delay)
+        async with s.lock:
+            await asyncio.sleep(0.01)
+            await store.advance("s1", step_id=step_id)
+            results.append(step_id)
+
+    await asyncio.gather(slow_advance(0, 0.0), slow_advance(1, 0.01))
+    # Both steps ran under the per-session lock and committed in order.
+    assert sorted(results) == [0, 1]

tests/rl_rollout/test_serving_rl_rollout.py

@@ -0,0 +1,371 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for ServingRLRollout (RFC #3747, P0).
+
+The DreamZero engine is mocked. Tests verify session lifecycle, step_id
+semantics, committed_step_id invariants, and error paths without a GPU.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+from vllm_omni.entrypoints.openai.protocol.rollout import (
+    Action,
+    CreateSessionRequest,
+    Observation,
+    RolloutStepRequest,
+)
+from vllm_omni.entrypoints.openai.serving_rl_rollout import ServingRLRollout
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_fake_result(video: Any = None, *, formatter_shaped: bool = True) -> MagicMock:
+    result = MagicMock()
+    if video is None:
+        video = np.zeros((1, 3, 64, 64), dtype=np.float32)
+    if formatter_shaped:
+        # Matches output_formatter: video is peeled into images; mm only has actions.
+        result.images = video
+        result.multimodal_output = {"actions": np.zeros((1, 8), dtype=np.float32)}
+    else:
+        result.images = None
+        result.multimodal_output = {"video": video}
+    return result
+
+
+def _make_openpi(video: Any = None, raises: Exception | None = None, *, formatter_shaped: bool = True):
+    openpi = MagicMock()
+    openpi.engine_client = MagicMock()
+    openpi.model_name = None
+
+    async def _generate(*_, **__):
+        if raises:
+            raise raises
+        yield _make_fake_result(video, formatter_shaped=formatter_shaped)
+
+    openpi.engine_client.generate = _generate
+    openpi.build_request = MagicMock(
+        return_value=MagicMock(
+            prompt="",
+            prompts=[""],
+            request_id="test-req",
+            sampling_params=MagicMock(),
+        )
+    )
+    return openpi
+
+
+@pytest.fixture
+def serving() -> ServingRLRollout:
+    return ServingRLRollout(_make_openpi())
+
+
+# ---------------------------------------------------------------------------
+# Session lifecycle
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_create_session(serving):
+    resp = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    assert resp.session_id
+    assert resp.mode == "world_model_env"
+
+
+@pytest.mark.asyncio
+async def test_get_status_after_create(serving):
+    resp = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    status = await serving.get_status(resp.session_id)
+    assert status.committed_step_id == -1
+    assert status.context_length == 0
+    assert not status.closed
+
+
+@pytest.mark.asyncio
+async def test_close_session(serving):
+    resp = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    await serving.close_session(resp.session_id)
+    status = await serving.get_status(resp.session_id)
+    assert status.closed
+    assert status.committed_step_id == -1
+    serving._openpi.drop_session.assert_any_call(resp.session_id)
+    serving._openpi.drop_session.assert_any_call(f"{resp.session_id}:stateless")
+
+
+# ---------------------------------------------------------------------------
+# Step - success path
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_successful_step_advances_committed_step_id(serving):
+    sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    req = RolloutStepRequest(
+        step_id=0,
+        observation=Observation(state=[0.1, 0.2]),
+        action=Action(joint_positions=[0.0] * 7),
+    )
+    resp = await serving.step(sess.session_id, req)
+    assert resp.error is None
+    assert resp.model_metadata.committed_step_id == 0
+
+    status = await serving.get_status(sess.session_id)
+    assert status.committed_step_id == 0
+    assert status.context_length == 1
+
+
+@pytest.mark.asyncio
+async def test_repeated_steps_advance_monotonically(serving):
+    sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    for i in range(3):
+        req = RolloutStepRequest(
+            step_id=i,
+            observation=Observation(state=[float(i)]),
+            action=Action(joint_positions=[0.0] * 7),
+        )
+        resp = await serving.step(sess.session_id, req)
+        assert resp.error is None
+        assert resp.model_metadata.committed_step_id == i
+
+    status = await serving.get_status(sess.session_id)
+    assert status.committed_step_id == 2
+    assert status.context_length == 3
+
+
+@pytest.mark.asyncio
+async def test_next_observation_present_on_success(serving):
+    sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    req = RolloutStepRequest(
+        step_id=0,
+        observation=Observation(),
+        action=Action(),
+    )
+    resp = await serving.step(sess.session_id, req)
+    assert resp.next_observation is not None
+    assert "video_latent" in resp.next_observation
+    assert resp.next_observation["dtype"] == "float32"
+
+
+@pytest.mark.asyncio
+async def test_duplicate_step_does_not_rerun_or_advance(serving):
+    sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action())
+    assert (await serving.step(sess.session_id, req)).error is None
+
+    resp = await serving.step(sess.session_id, req)
+
+    assert resp.error is not None
+    assert resp.error.code == "step_already_committed"
+    assert resp.model_metadata.committed_step_id == 0
+    assert resp.model_metadata.context_length == 1
+    serving._openpi.build_request.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_out_of_order_step_does_not_rerun_or_advance(serving):
+    sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    req = RolloutStepRequest(step_id=2, observation=Observation(), action=Action())
+
+    resp = await serving.step(sess.session_id, req)
+
+    assert resp.error is not None
+    assert resp.error.code == "step_out_of_order"
+    assert resp.model_metadata.committed_step_id == -1
+    assert resp.model_metadata.context_length == 0
+    serving._openpi.build_request.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_stateless_step_does_not_commit_context(serving):
+    sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    req = RolloutStepRequest(
+        step_id=7,
+        observation=Observation(),
+        action=Action(),
+        use_session_context=False,
+    )
+
+    resp = await serving.step(sess.session_id, req)
+
+    assert resp.error is None
+    assert resp.model_metadata.committed_step_id == -1
+    assert resp.model_metadata.context_length == 0
+    serving._openpi.build_request.assert_called_once()
+    kwargs = serving._openpi.build_request.call_args.kwargs
+    assert kwargs["reset"] is True
+    assert kwargs["session_id"] == f"{sess.session_id}:stateless"
+
+
+# ---------------------------------------------------------------------------
+# Step - failure path: committed_step_id must NOT advance
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_failed_step_does_not_advance_committed_step_id():
+    serving = ServingRLRollout(_make_openpi(raises=RuntimeError("engine dead")))
+    sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action())
+    resp = await serving.step(sess.session_id, req)
+
+    assert resp.error is not None
+    assert resp.error.code == "inference_error"
+    assert resp.model_metadata.committed_step_id == -1  # must not have advanced
+
+    status = await serving.get_status(sess.session_id)
+    assert status.committed_step_id == -1
+
+
+@pytest.mark.asyncio
+async def test_output_encoding_failure_does_not_advance_committed_step_id():
+    class BadVideo:
+        def detach(self):
+            return self
+
+        def float(self):
+            return self
+
+        def cpu(self):
+            return self
+
+        def numpy(self):
+            raise TypeError("cannot encode")
+
+    serving = ServingRLRollout(_make_openpi(video=BadVideo()))
+    sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action())
+
+    resp = await serving.step(sess.session_id, req)
+
+    assert resp.error is not None
+    assert resp.error.code == "inference_error"
+    assert resp.model_metadata.committed_step_id == -1
+    status = await serving.get_status(sess.session_id)
+    assert status.committed_step_id == -1
+
+
+@pytest.mark.asyncio
+async def test_failed_step_after_success_preserves_last_committed():
+    good_serving = ServingRLRollout(_make_openpi())
+    sess = await good_serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    # Step 0 succeeds
+    await good_serving.step(sess.session_id, RolloutStepRequest(step_id=0, observation=Observation(), action=Action()))
+
+    # Swap engine to failing one - reuse the same store via internal reference
+    good_serving._openpi = _make_openpi(raises=RuntimeError("oops"))
+    resp = await good_serving.step(
+        sess.session_id, RolloutStepRequest(step_id=1, observation=Observation(), action=Action())
+    )
+    assert resp.error is not None
+    assert resp.model_metadata.committed_step_id == 0  # still at last good step
+    assert resp.model_metadata.context_length == 1
+
+
+# ---------------------------------------------------------------------------
+# Step - unknown session
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_step_unknown_session_returns_error(serving):
+    req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action())
+    resp = await serving.step("no-such-session", req)
+    assert resp.error is not None
+    assert resp.error.code == "session_not_found"
+
+
+# ---------------------------------------------------------------------------
+# Reset
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_reset_clears_committed_step_id(serving):
+    sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    await serving.step(sess.session_id, RolloutStepRequest(step_id=0, observation=Observation(), action=Action()))
+    await serving.reset_session(sess.session_id)
+    status = await serving.get_status(sess.session_id)
+    assert status.committed_step_id == -1
+    assert status.context_length == 0
+    serving._openpi.drop_session.assert_called_with(sess.session_id)
+
+
+# ---------------------------------------------------------------------------
+# action -> obs folding (module assumption)
+# ---------------------------------------------------------------------------
+
+
+def test_merge_action_into_obs_emits_dreamzero_state_keys():
+    from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs
+
+    obs = Observation(state=[9.0] * 8, prompt="test")
+    action = Action(joint_positions=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], gripper_position=0.5)
+    result = _merge_action_into_obs(obs, action)
+
+    assert result["prompt"] == "test"
+    np.testing.assert_array_almost_equal(result["observation/joint_position"], [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0])
+    np.testing.assert_array_almost_equal(result["observation/gripper_position"], [0.5])
+
+
+def test_merge_action_none_splits_obs_state_for_dreamzero():
+    from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs
+
+    obs = Observation(state=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 0.25])
+    result = _merge_action_into_obs(obs, None)
+    np.testing.assert_array_almost_equal(result["observation/joint_position"], [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0])
+    np.testing.assert_array_almost_equal(result["observation/gripper_position"], [0.25])
+
+
+def test_merge_no_state_no_action_omits_state_key():
+    from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs
+
+    obs = Observation()
+    result = _merge_action_into_obs(obs, None)
+    assert "observation/joint_position" not in result
+    assert "observation/gripper_position" not in result
+
+
+def test_merge_rejects_wrong_dreamzero_joint_dimension():
+    from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs
+
+    with pytest.raises(ValueError, match="joint_positions must have 7 values"):
+        _merge_action_into_obs(Observation(), Action(joint_positions=[1.0, 2.0]))
+
+
+@pytest.mark.asyncio
+async def test_step_reads_formatter_images_not_multimodal_video():
+    video = np.ones((1, 3, 8, 8), dtype=np.float32)
+    serving = ServingRLRollout(_make_openpi(video=video, formatter_shaped=True))
+    sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    resp = await serving.step(
+        sess.session_id,
+        RolloutStepRequest(step_id=0, observation=Observation(), action=Action()),
+    )
+    assert resp.error is None
+    assert resp.next_observation is not None
+    assert "video_latent" in resp.next_observation
+    assert resp.next_observation["shape"] == [1, 3, 8, 8]
+
+
+@pytest.mark.asyncio
+async def test_step_falls_back_to_multimodal_video():
+    video = np.ones((2, 4, 4), dtype=np.float32)
+    serving = ServingRLRollout(_make_openpi(video=video, formatter_shaped=False))
+    sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env"))
+    resp = await serving.step(
+        sess.session_id,
+        RolloutStepRequest(step_id=0, observation=Observation(), action=Action()),
+    )
+    assert resp.error is None
+    assert resp.next_observation is not None
+    assert resp.next_observation["shape"] == [2, 4, 4]

vllm_omni/entrypoints/openai/api_server.py

@@ -138,6 +138,10 @@ from vllm_omni.entrypoints.openai.protocol.images import (
     ImageGenerationResponse,
     ResponseFormat,
 )
+from vllm_omni.entrypoints.openai.protocol.rollout import (
+    CreateSessionRequest,
+    RolloutStepRequest,
+)
 from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoDeleteResponse,
     VideoGenerationRequest,
@@ -146,8 +150,14 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoResponse,
 )
 from vllm_omni.entrypoints.openai.realtime_connection import RealtimeConnection
+from vllm_omni.entrypoints.openai.rollout_session import (
+    RolloutSessionCapacityError,
+    RolloutSessionClosedError,
+    RolloutSessionNotFoundError,
+)
 from vllm_omni.entrypoints.openai.serving_audio_generate import OmniOpenAIServingAudioGenerate
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+from vllm_omni.entrypoints.openai.serving_rl_rollout import ServingRLRollout
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
 from vllm_omni.entrypoints.openai.serving_video import (
@@ -604,6 +614,10 @@ async def omni_init_app_state(
             engine_client=engine_client,
             model_name=model_name,
         )
+        if state.openai_serving_realtime_robot is not None:
+            state.rl_rollout_serving = ServingRLRollout(state.openai_serving_realtime_robot)
+        else:
+            state.rl_rollout_serving = None
 
         state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False)
         state.server_load_metrics = 0
@@ -949,6 +963,7 @@ async def omni_init_app_state(
         stage_configs=state.stage_configs,
     )
     state.openai_serving_realtime_robot = None
+    state.rl_rollout_serving = None
 
     state.enable_server_load_tracking = args.enable_server_load_tracking
     state.server_load_metrics = 0
@@ -1529,6 +1544,78 @@ async def duplex_websocket(websocket: WebSocket):
     await handler.handle_session(websocket)
 
 
+# RL Rollout serving (RFC #3747, P0)
+
+
+def _rl_rollout_serving(request: Request) -> ServingRLRollout:
+    serving = getattr(request.app.state, "rl_rollout_serving", None)
+    if serving is None:
+        raise HTTPException(status_code=501, detail="RL rollout serving not available for this model.")
+    return serving
+
+
+@router.post("/v1/realtime/sessions")
+async def create_rollout_session(body: CreateSessionRequest, request: Request):
+    serving = _rl_rollout_serving(request)
+    try:
+        return (await serving.create_session(body)).model_dump()
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+    except RolloutSessionCapacityError as exc:
+        raise HTTPException(status_code=503, detail=str(exc)) from exc
+
+
+@router.post("/v1/realtime/sessions/{session_id}/step")
+async def rollout_step(session_id: str, body: RolloutStepRequest, request: Request):
+    serving = _rl_rollout_serving(request)
+    response = await serving.step(session_id, body)
+    if response.error is not None:
+        status_code = None
+        if response.error.code == "session_not_found":
+            status_code = 404
+        elif response.error.code == "session_closed":
+            status_code = 410
+        elif response.error.code in {"invalid_request", "step_already_committed", "step_out_of_order"}:
+            status_code = 400
+        if status_code is not None:
+            return JSONResponse(content=response.model_dump(), status_code=status_code)
+    return response.model_dump()
+
+
+@router.post("/v1/realtime/sessions/{session_id}/reset")
+async def reset_rollout_session(session_id: str, request: Request):
+    serving = _rl_rollout_serving(request)
+    try:
+        return (await serving.reset_session(session_id)).model_dump()
+    except RolloutSessionNotFoundError:
+        raise HTTPException(status_code=404, detail=f"Session {session_id!r} not found.")
+    except RolloutSessionClosedError:
+        raise HTTPException(status_code=410, detail=f"Session {session_id!r} is closed.")
+
+
+@router.post("/v1/realtime/sessions/{session_id}/close")
+async def close_rollout_session(session_id: str, request: Request):
+    serving = _rl_rollout_serving(request)
+    try:
+        await serving.close_session(session_id)
+        return {"session_id": session_id, "closed": True}
+    except RolloutSessionNotFoundError:
+        raise HTTPException(status_code=404, detail=f"Session {session_id!r} not found.")
+    except RolloutSessionClosedError:
+        return {"session_id": session_id, "closed": True}
+
+
+@router.get("/v1/realtime/sessions/{session_id}/status")
+async def rollout_session_status(session_id: str, request: Request):
+    serving = _rl_rollout_serving(request)
+    try:
+        return (await serving.get_status(session_id)).model_dump()
+    except RolloutSessionNotFoundError:
+        raise HTTPException(status_code=404, detail=f"Session {session_id!r} not found.")
+    except RolloutSessionClosedError:
+        raise HTTPException(status_code=410, detail=f"Session {session_id!r} is closed.")
+
+
 # Health and Model endpoints for diffusion mode
 
 

vllm_omni/entrypoints/openai/protocol/rollout.py

@@ -0,0 +1,111 @@
+# 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

vllm_omni/entrypoints/openai/rollout_session.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""In-memory session store for RL rollout serving (RFC #3747, P0).
+
+Each RolloutSession holds a per-session asyncio.Lock so that concurrent step
+requests for the same session are serialised. committed_step_id advances only
+on successful step completion; failed or timed-out steps leave it unchanged,
+preserving the atomic-context-commit guarantee from section 6.5 of the RFC.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections import OrderedDict
+from dataclasses import dataclass, field
+from typing import Literal
+
+
+class RolloutSessionNotFoundError(KeyError):
+    pass
+
+
+class RolloutSessionClosedError(RuntimeError):
+    pass
+
+
+class RolloutSessionStepError(ValueError):
+    pass
+
+
+class RolloutSessionCapacityError(RuntimeError):
+    pass
+
+
+@dataclass
+class RolloutSession:
+    session_id: str
+    model: str
+    mode: Literal["world_model_env"]
+    created_at: float = field(default_factory=time.time)
+    committed_step_id: int = -1
+    context_length: int = 0
+    closed: bool = False
+    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+
+
+class RolloutSessionStore:
+    """Thread-safe (asyncio) store for active rollout sessions."""
+
+    def __init__(self, max_sessions: int = 64) -> None:
+        self._sessions: OrderedDict[str, RolloutSession] = OrderedDict()
+        self._max_sessions = max_sessions
+        self._store_lock = asyncio.Lock()
+
+    async def create(self, session_id: str, model: str, mode: Literal["world_model_env"]) -> RolloutSession:
+        async with self._store_lock:
+            self._evict_closed_sessions_for_create()
+            if len(self._sessions) >= self._max_sessions:
+                raise RolloutSessionCapacityError(f"Maximum active rollout sessions reached: {self._max_sessions}.")
+            session = RolloutSession(
+                session_id=session_id,
+                model=model,
+                mode=mode,
+            )
+            self._sessions[session_id] = session
+            self._sessions.move_to_end(session_id)
+            return session
+
+    async def get(self, session_id: str, *, include_closed: bool = False) -> RolloutSession:
+        session = self._sessions.get(session_id)
+        if session is None:
+            raise RolloutSessionNotFoundError(session_id)
+        self._sessions.move_to_end(session_id)
+        if session.closed and not include_closed:
+            raise RolloutSessionClosedError(session_id)
+        return session
+
+    async def reset(self, session_id: str) -> RolloutSession:
+        session = await self.get(session_id)
+        async with session.lock:
+            session.committed_step_id = -1
+            session.context_length = 0
+        return session
+
+    async def close(self, session_id: str) -> None:
+        session = await self.get(session_id, include_closed=True)
+        async with session.lock:
+            session.closed = True
+
+    async def advance(self, session_id: str, step_id: int) -> None:
+        """Commit a successfully completed step. Called only on success."""
+        session = await self.get(session_id)
+        expected_step_id = session.committed_step_id + 1
+        if step_id != expected_step_id:
+            raise RolloutSessionStepError(f"Expected step_id {expected_step_id}, got {step_id}.")
+        session.committed_step_id = step_id
+        session.context_length += 1
+
+    def _evict_closed_sessions_for_create(self) -> None:
+        while len(self._sessions) >= self._max_sessions:
+            for session_id, session in self._sessions.items():
+                if session.closed:
+                    del self._sessions[session_id]
+                    break
+            else:
+                break

vllm_omni/entrypoints/openai/serving_rl_rollout.py

@@ -0,0 +1,370 @@
+# 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,
+            ),
+        )

vllm_omni/entrypoints/openpi/serving.py

@@ -119,10 +119,32 @@ class ServingRealtimeRobotOpenPI:
     def reset(self, obs: dict) -> None:
         """Compatibility hook; per-connection state lives in RobotRealtimeConnection."""
 
+    def drop_session(self, session_id: str) -> None:
+        """Best-effort release of model-side session state for a closed rollout."""
+        drop = getattr(self.engine_client, "drop_session", None)
+        if callable(drop):
+            drop(session_id)
+            return
+        pipeline = self._pipeline()
+        for name in ("close_ar_diffusion_session", "drop_session_state"):
+            close = getattr(pipeline, name, None)
+            if callable(close):
+                close(session_id)
+                return
+
+    def _pipeline(self) -> Any:
+        engine = self.engine_client
+        for attr in ("model_runner", "runner", "diffusion_model_runner"):
+            runner = getattr(engine, attr, None)
+            pipeline = getattr(runner, "pipeline", None) if runner is not None else None
+            if pipeline is not None:
+                return pipeline
+        return getattr(engine, "pipeline", None)
+
     async def infer(self, obs: dict, *, session_id: str, reset: bool) -> ActionOutput:
         """raw obs → engine → actions."""
         # Build request, run inference through AsyncOmni
-        request = self._build_request(obs, session_id=session_id, reset=reset)
+        request = self.build_request(obs, session_id=session_id, reset=reset)
         result = None
         # OpenPI policy serving is one request -> one action reply. AsyncOmni
         # exposes an async iterator, so consume it to completion and use the
@@ -141,6 +163,10 @@ class ServingRealtimeRobotOpenPI:
     def _next_request_id(self, session_id: str) -> str:
         return f"robot-{session_id}-{next(self._request_counter)}"
 
+    def build_request(self, obs: dict, *, session_id: str, reset: bool) -> Any:
+        """Build an engine request from raw robot obs."""
+        return self._build_request(obs, session_id=session_id, reset=reset)
+
     def _build_request(self, obs: dict, *, session_id: str, reset: bool) -> Any:
         """Build engine request from raw robot obs.