[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"]