Cosmos3 action policy improvements (#6460)

Signed-off-by: Maciej Bala <mbala@nvidia.com>
Signed-off-by: MaciejBalaNV <mbala@nvidia.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/entrypoints/openai_api/test_openpi_connection.py

@@ -310,6 +310,29 @@ def test_handle_connection_closes_websocket_on_idle_timeout(monkeypatch):
     serving.infer.assert_not_called()
 
 
+def test_handle_connection_can_disable_idle_timeout(monkeypatch):
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+
+    async def fail_wait_for(*_args, **_kwargs):
+        pytest.fail("asyncio.wait_for should not be used when the idle timeout is disabled")
+
+    monkeypatch.setattr(openpi_connection.asyncio, "wait_for", fail_wait_for)
+    websocket = FakeWebSocket([{"type": "websocket.disconnect"}])
+    serving = _serving_mock()
+
+    asyncio.run(
+        openpi_connection.RobotRealtimeConnection(
+            websocket,
+            serving,
+            idle_timeout=None,
+        ).handle_connection()
+    )
+
+    assert websocket.accepted is True
+    assert websocket.closed is False
+    serving.infer.assert_not_called()
+
+
 def test_handle_connection_keeps_session_state_per_websocket(monkeypatch):
     monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
     requests = {
@@ -343,6 +366,73 @@ def test_handle_connection_keeps_session_state_per_websocket(monkeypatch):
     assert calls[2].kwargs == {"session_id": "session-b", "reset": True}
 
 
+def test_handle_connection_keeps_interleaved_parallel_session_state(monkeypatch):
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+    requests = {
+        b"a1": {"prompt": "first-a", "session_id": "session-a"},
+        b"b1": {"prompt": "first-b", "session_id": "session-b"},
+        b"a2": {"prompt": "second-a", "session_id": "session-a"},
+        b"b2": {"prompt": "second-b", "session_id": "session-b"},
+    }
+    monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data]))
+    serving = _serving_mock()
+    websocket = FakeWebSocket(
+        [{"type": "websocket.receive", "bytes": key} for key in (b"a1", b"b1", b"a2", b"b2")]
+        + [{"type": "websocket.disconnect"}]
+    )
+
+    asyncio.run(openpi_connection.RobotRealtimeConnection(websocket, serving).handle_connection())
+
+    assert [call.kwargs for call in serving.infer.await_args_list] == [
+        {"session_id": "session-a", "reset": True},
+        {"session_id": "session-b", "reset": True},
+        {"session_id": "session-a", "reset": False},
+        {"session_id": "session-b", "reset": False},
+    ]
+
+
+def test_handle_connection_bounds_tracked_sessions_without_resetting_live_ones(monkeypatch):
+    """A client may hold one socket across many episodes, minting a fresh
+    session id per episode and never sending ``reset`` (RoboLab does exactly
+    this). Session tracking must stay bounded, and evicting finished episodes
+    must not disturb the sessions that are still running.
+    """
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+    monkeypatch.setattr(openpi_connection, "MAX_TRACKED_SESSIONS", 8)
+
+    # One long-lived session interleaved with far more short episodes than the cap.
+    episodes = 40
+    keys = []
+    requests = {}
+    for episode in range(episodes):
+        for session_id in (f"episode-{episode}", "live-env"):
+            key = f"{session_id}@{episode}".encode()
+            keys.append(key)
+            requests[key] = {"prompt": "pick", "session_id": session_id}
+
+    monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data]))
+    serving = _serving_mock()
+    websocket = FakeWebSocket(
+        [{"type": "websocket.receive", "bytes": key} for key in keys] + [{"type": "websocket.disconnect"}]
+    )
+    connection = openpi_connection.RobotRealtimeConnection(websocket, serving)
+
+    asyncio.run(connection.handle_connection())
+
+    assert len(connection._seen_sessions) <= 8
+
+    resets_by_session: dict[str, list[bool]] = {}
+    for call in serving.infer.await_args_list:
+        resets_by_session.setdefault(call.kwargs["session_id"], []).append(call.kwargs["reset"])
+
+    # Each new episode resets exactly once, on its first observation.
+    for episode in range(episodes):
+        assert resets_by_session[f"episode-{episode}"] == [True]
+    # The continuously active session is never evicted, so it resets only once
+    # despite ``episodes`` other sessions churning past the cap.
+    assert resets_by_session["live-env"] == [True] + [False] * (episodes - 1)
+
+
 def test_handle_connection_reset_endpoint_resets_next_infer(monkeypatch):
     monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
     requests = {

tests/entrypoints/openai_api/test_openpi_serving.py

@@ -15,6 +15,7 @@ from starlette.testclient import TestClient
 
 from vllm_omni.entrypoints.openpi import connection as openpi_connection
 from vllm_omni.entrypoints.openpi import serving as openpi_serving
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -177,6 +178,22 @@ def test_create_policy_server_returns_none_without_policy_config():
     assert serving is None
 
 
+def test_legacy_policy_request_defaults_key_no_longer_controls_serving():
+    od_config = SimpleNamespace(
+        model_config={
+            "policy_server_config": {},
+            "policy_request_defaults": "malformed legacy value",
+        }
+    )
+
+    serving = openpi_serving.ServingRealtimeRobotOpenPI.create_policy_server(
+        engine_client=SimpleNamespace(get_diffusion_od_config=lambda: od_config),
+        model_name="generic-policy",
+    )
+
+    assert serving is not None
+
+
 def test_policy_server_config_allows_explicit_empty_config():
     serving = openpi_serving.ServingRealtimeRobotOpenPI(
         engine_client=_engine_with_policy_config(policy_config={}),
@@ -221,6 +238,44 @@ def test_build_request_uses_unique_engine_request_id_per_inference():
     assert request_a.request_id != request_b.request_id
 
 
+def test_build_request_clones_stage_defaults_before_protocol_fields():
+    default_params = OmniDiffusionSamplingParams(
+        extra_args={
+            "format_prompt_as_json": True,
+            "session_id": "configured-session-must-not-win",
+            "reset": False,
+            "nested": {"values": []},
+        }
+    )
+    od_config = SimpleNamespace(model_config={"policy_server_config": {}})
+    engine_client = SimpleNamespace(
+        get_diffusion_od_config=lambda: od_config,
+        default_sampling_params_list=[default_params],
+    )
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(
+        engine_client=engine_client,
+        model_name="custom-policy",
+    )
+
+    request_a = serving._build_request(
+        {"prompt": "pick up the object"},
+        session_id="wire-session",
+        reset=True,
+    )
+    request_b = serving._build_request(
+        {"prompt": "pick up the object"},
+        session_id="wire-session",
+        reset=False,
+    )
+
+    assert request_a.sampling_params.extra_args["format_prompt_as_json"] is True
+    assert request_a.sampling_params.extra_args["session_id"] == "wire-session"
+    assert request_a.sampling_params.extra_args["reset"] is True
+    request_a.sampling_params.extra_args["nested"]["values"].append("request-a")
+    assert request_b.sampling_params.extra_args["nested"] == {"values": []}
+    assert default_params.extra_args["nested"] == {"values": []}
+
+
 def test_build_request_forwards_seed_to_sampling_params():
     """``seed`` in the inference message is the engine-level seed; omitted, the request auto-seeds."""
     serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config())
@@ -257,6 +312,29 @@ def test_infer_keeps_session_state_but_uses_unique_engine_request_ids():
     assert sampling_params_b.extra_args["reset"] is False
 
 
+def test_infer_yields_event_loop_while_engine_is_running():
+    class DelayedEngine(RecordingEngine):
+        def generate(self, *, prompt, request_id, sampling_params_list):
+            async def _generate():
+                await asyncio.sleep(0.03)
+                yield SimpleNamespace(multimodal_output={"actions": [0.0]})
+
+            return _generate()
+
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=DelayedEngine())
+
+    async def run_test():
+        infer_task = asyncio.create_task(serving.infer({"prompt": "pick"}, session_id="session-a", reset=True))
+        ticks = 0
+        while not infer_task.done():
+            await asyncio.sleep(0.005)
+            ticks += 1
+        await infer_task
+        return ticks
+
+    assert asyncio.run(run_test()) >= 2
+
+
 def test_two_websocket_clients_without_session_id_do_not_conflict(monkeypatch):
     monkeypatch.setattr(openpi_connection, "_pack", _json_pack)
     monkeypatch.setattr(openpi_connection, "_unpack", _json_unpack)