From 76d6fa4ca902a1eb94192f906343634beb555926 Mon Sep 17 00:00:00 2001 From: tactino <18781106300@163.com> Date: Thu, 1 Oct 2026 16:59:16 -0400 Subject: [PATCH] agent and rollout: a reconnect voids every chunk in flight (ConnectionReplaced); probe-v1 env for plugrl-conformance --probe The probe checker's resync scenario failed this client: after the agent reconnected inside its next call, envs mid-chunk kept executing chunks from the old connection and sent their feedback on the new one, and a failed feedback was followed by a reconnect inside the next feedback call. SPEC 7.6 forbids both. The agent now raises ConnectionReplaced instead, feedback never reconnects, and the rollout drops every chunk in flight and re-plans every env; ended episodes are still reset. The reconnect-on-server-stop resend (the SPEC 7.6 Gap) is gone with it. --- README.md | 1 + .../agent/websocket_env_client_agent.py | 162 +++++++++++------- src/plugrl_env_client/envs/probe_env.py | 106 ++++++++++++ src/plugrl_env_client/runner/rollout.py | 50 ++++-- tests/test_reconnect_drops_feedback.py | 79 ++++++++- tests/test_rollout_connection_replaced.py | 75 ++++++++ 6 files changed, 388 insertions(+), 85 deletions(-) create mode 100644 src/plugrl_env_client/envs/probe_env.py create mode 100644 tests/test_rollout_connection_replaced.py diff --git a/README.md b/README.md index b6bfef4..3f1c26f 100644 --- a/README.md +++ b/README.md @@ -123,6 +123,7 @@ uv run plugrl-run-env-client [OPTIONS] | Type | Extra required | Description | | :--- | :--- | :--- | | **dummy-v1** | — | Dummy environment for protocol and connectivity tests | +| **probe-v1** | — | The probe environment of SPEC.md section 8.1, for `plugrl-conformance --probe` | | **mujoco-v1** | `mujoco` | Gymnasium MuJoCo control, default `HalfCheetah-v5` | | **classic-v1** | `classic` | Classic control environments (e.g. CartPole) | | **atari-v1** | `atari` | Atari games via ALE | diff --git a/src/plugrl_env_client/agent/websocket_env_client_agent.py b/src/plugrl_env_client/agent/websocket_env_client_agent.py index 03f163f..c70dd9f 100644 --- a/src/plugrl_env_client/agent/websocket_env_client_agent.py +++ b/src/plugrl_env_client/agent/websocket_env_client_agent.py @@ -31,6 +31,26 @@ class ServerStopped(RuntimeError): """Raised when the server explicitly requests env clients to stop.""" +class ConnectionReplaced(RuntimeError): + """The connection the caller's actions came from is gone. + + The server keeps each environment's half of a transition - the previous + observation, the policy step state - on the connection that answered the + infer, so a new connection starts with none of it (SPEC section 7.6). Any + action chunk the caller is still executing was answered on the old one: + its feedback can never be completed, and must not be sent on the new + connection. The caller drops every chunk in flight and asks again for + all of its environments. + + It used not to be told. The agent reconnected inside its next call and + carried on, so environments in the middle of a chunk kept executing it + and then sent its feedback on the new connection, and a feedback whose + send failed was followed by a reconnect inside the next `feedback` call, + whose first message was then a stale feedback. `plugrl-conformance + --probe --scenario resync` caught both. + """ + + class WebSocketEnvClientAgent(_base_agent.BaseAgent): def __init__( self, @@ -46,6 +66,11 @@ def __init__( self._api_key = api_key self._reconnect_on_server_stop = reconnect_on_server_stop self._ws, self._server_metadata = self._wait_for_server() + # Which connection this is, and which one answered the caller's last + # infer. When they differ, the caller is holding actions from a + # connection that is gone; see ConnectionReplaced. + self._connection_id = 1 + self._actions_from: int | None = None def _close_connection(self) -> None: if self._ws is None: @@ -130,6 +155,12 @@ def _ensure_connection(self) -> None: return logger.warning("Connection closed. Attempting to re-establish connection.") self._ws, self._server_metadata = self._wait_for_server() + self._connection_id += 1 + + def _replaced(self, why: str) -> ConnectionReplaced: + """Tell the caller once; its next infer then goes out as normal.""" + self._actions_from = None + return ConnectionReplaced(why) def infer( self, @@ -143,6 +174,15 @@ def infer( ws = self._ws if ws is None: raise RuntimeError("WebSocket connection is not available.") + if ( + self._actions_from is not None + and self._actions_from != self._connection_id + ): + # Nothing is sent: the caller has to re-plan every env first. + raise self._replaced( + "The connection was replaced; the action chunks in flight " + "came from the old one and are void." + ) try: packed_data = self._packer.pack( @@ -159,6 +199,7 @@ def infer( if isinstance(response, str): raise RuntimeError(f"Error in inference server:\n{response}") + self._actions_from = self._connection_id return msgpack_numpy.unpackb(response)["data"] except ConnectionClosedOK as exc: close_code, close_reason = _get_close_details(exc) @@ -220,71 +261,66 @@ def feedback( makes the server build a transition out of an empty observation and store it, with nothing downstream able to tell. SPEC section 7.6. - So a closed connection here costs exactly one transition, and that is - the cheap outcome. The next infer reconnects and resyncs. + So a closed connection here costs the transitions in flight, and that + is the cheap outcome. This method never reconnects: a feedback is only + ever sent on the connection that answered its infer. When that + connection is gone it raises ConnectionReplaced, so the caller drops + every chunk it is still executing, and the next infer reconnects. """ - while True: - self._ensure_connection() - ws = self._ws - if ws is None: - raise RuntimeError("WebSocket connection is not available.") - - try: - packed_data = self._packer.pack( - { - "message_type": str(MessageType.FEEDBACK), - "env_indices": env_indices, - "step_ids": step_ids, - "data": { - "obs": obs, - "rewards": rewards, - "terminated": terminated, - "truncated": truncated, - "info": info, - }, - } - ) - ws.send(packed_data) - return - except ConnectionClosedOK as exc: - close_code, close_reason = _get_close_details(exc) - self._close_connection() + ws = self._ws + if ws is None or self._actions_from != self._connection_id: + raise self._replaced( + "The connection that answered this feedback's infer is gone; " + "dropping the feedback." + ) - if close_reason == SERVER_STOP_REASON: - if self._reconnect_on_server_stop: - logger.info( - "Server requested env client shutdown during FEEDBACK send, " - "but reconnect_on_server_stop is enabled. " - "Waiting for server to come back and retrying..." - ) - continue - raise ServerStopped( - "Server requested env client shutdown after algorithm stop." - ) from exc - - if close_reason == SERVER_RESYNC_REASON: - logger.info( - "Server requested session resync during FEEDBACK send. " - "Dropping stale feedback and resuming from the next infer request." - ) - return - - logger.warning( - "Connection closed during FEEDBACK send. Dropping this " - "transition and resuming from the next infer request. " - f"code={close_code}, reason={close_reason or ''}" - ) - return - except ConnectionClosedError as exc: - logger.warning( - "Connection closed during FEEDBACK send. Dropping this " - f"transition and resuming from the next infer request. {exc}" - ) - self._close_connection() - return - except Exception: - self._close_connection() - raise + try: + packed_data = self._packer.pack( + { + "message_type": str(MessageType.FEEDBACK), + "env_indices": env_indices, + "step_ids": step_ids, + "data": { + "obs": obs, + "rewards": rewards, + "terminated": terminated, + "truncated": truncated, + "info": info, + }, + } + ) + ws.send(packed_data) + except ConnectionClosedOK as exc: + close_code, close_reason = _get_close_details(exc) + self._close_connection() + + if ( + close_reason == SERVER_STOP_REASON + and not self._reconnect_on_server_stop + ): + raise ServerStopped( + "Server requested env client shutdown after algorithm stop." + ) from exc + # A resync, a stop the client waits out, or any other close: the + # feedback is dropped either way. Before, a stop with + # reconnect_on_server_stop retried, which resent it on the new + # connection - the one path that broke SPEC section 7.6. + logger.info( + "Connection closed during FEEDBACK send. Dropping the " + "transitions in flight and resuming from the next infer request. " + f"code={close_code}, reason={close_reason or ''}" + ) + raise self._replaced("dropped feedback after a close") from exc + except ConnectionClosedError as exc: + logger.warning( + "Connection closed during FEEDBACK send. Dropping the " + f"transitions in flight and resuming from the next infer request. {exc}" + ) + self._close_connection() + raise self._replaced("dropped feedback after a close") from exc + except Exception: + self._close_connection() + raise def reset(self) -> None: return diff --git a/src/plugrl_env_client/envs/probe_env.py b/src/plugrl_env_client/envs/probe_env.py new file mode 100644 index 0000000..1e4fbc8 --- /dev/null +++ b/src/plugrl_env_client/envs/probe_env.py @@ -0,0 +1,106 @@ +"""The probe environment of SPEC.md section 8.1, for `plugrl-conformance --probe`. + +Env i counts the steps of its episode in `states["t"]`, keeps the action it +applied last in `states["a"]`, pays a reward of 1 per step, and ends its +episode, terminated, when t reaches 3 + 2 * (i % 3). It never truncates. + +The conformance checker sends actions whose values encode where in the chunk +they are, so from each feedback it can tell what this client did with them: +whether it summed the chunk's reward, sent the terminal observation, and +applied the actions time-major and in order. It needs no extra: + + plugrl-conformance --probe --scenario all --action-dim 3 --client \ + "plugrl-run-env-client probe-v1 --server-host 127.0.0.1 --server-port 8000 \ + --num-envs 3 --num-episodes 1000000" + +`--env.action-dim` must match the checker's `--action-dim`. +""" + +from __future__ import annotations + +import dataclasses + +import gymnasium as gym +import numpy as np + +from plugrl_env_client.utils.registration import register_env, register_env_config + +from .base_env import ( + Action, + BaseEnv, + BaseEnvConfig, + BoolArray, + Observation, + RewardArray, +) + +UID = "Probe-v1" + + +@register_env_config(UID) +@dataclasses.dataclass +class ProbeEnvConfig(BaseEnvConfig): + action_dim: int = 3 + + +# No time limit: the probe env never truncates, and its episodes are 3 to 7 +# steps long. +@register_env(UID, max_episode_steps=None) +class ProbeEnv(BaseEnv): + def __init__( + self, + config: ProbeEnvConfig, + num_envs: int = 1, + process_id: int | None = None, + total_processes: int | None = None, + ): + super().__init__( + config=config, + num_envs=num_envs, + process_id=process_id, + total_processes=total_processes, + ) + self.action_dim = config.action_dim + # float64, so the checker's values arrive exactly whatever they are. + self.single_action_space = gym.spaces.Box( + -np.inf, np.inf, (self.action_dim,), np.float64 + ) + self.action_space = self.single_action_space + # Env indices on the wire are the positions in this vector env. + self.lengths = 3 + 2 * (np.arange(num_envs) % 3) + self.t = np.zeros(num_envs, dtype=np.int64) + self.a = np.zeros((num_envs, self.action_dim), dtype=np.float64) + + def _observation(self) -> Observation: + return Observation( + images={}, + states={ + "t": self.t[:, None].astype(np.float64), + "a": self.a.copy(), + }, + text="probe", + ) + + def reset(self, *, seed: int | None = None, options: dict | None = None) -> tuple: + self.seed_rngs(seed) + indices = np.arange(self.num_envs) + if options is not None and options.get("reset_indices") is not None: + indices = np.asarray(options["reset_indices"], dtype=np.int64) + self.t[indices] = 0 + self.a[indices] = 0.0 + return self._observation(), {} + + def step( + self, actions: Action + ) -> tuple[Observation, RewardArray, BoolArray, BoolArray, dict]: + actions = np.asarray(actions, dtype=np.float64).reshape( + self.num_envs, self.action_dim + ) + self.t += 1 + self.a = actions.copy() + reward = np.ones(self.num_envs, dtype=np.float32) + terminated = self.t == self.lengths + truncated = np.zeros(self.num_envs, dtype=np.bool_) + # No reset here: the client sends this terminal observation, then + # resets the env itself (SPEC.md section 5.4). + return self._observation(), reward, terminated, truncated, {} diff --git a/src/plugrl_env_client/runner/rollout.py b/src/plugrl_env_client/runner/rollout.py index 403e59a..f98de77 100644 --- a/src/plugrl_env_client/runner/rollout.py +++ b/src/plugrl_env_client/runner/rollout.py @@ -5,7 +5,10 @@ import numpy as np from loguru import logger -from plugrl_env_client.agent.websocket_env_client_agent import WebSocketEnvClientAgent +from plugrl_env_client.agent.websocket_env_client_agent import ( + ConnectionReplaced, + WebSocketEnvClientAgent, +) from plugrl_env_client.recorder import Recorder from plugrl_env_client.utils.rollout import _get_action_spec, _select_info, _select_obs @@ -126,11 +129,19 @@ def log_timing_summary(final: bool = False) -> None: timing.infer_obs_pack += time.perf_counter() - infer_obs_pack_started_at infer_wait_started_at = time.perf_counter() - action_chunk = agent.infer( - obs_msg, - env_indices=need_infer, - step_ids=step_id[need_infer], - )["action"] + try: + action_chunk = agent.infer( + obs_msg, + env_indices=need_infer, + step_ids=step_id[need_infer], + )["action"] + except ConnectionReplaced: + # SPEC section 7.6: every chunk still in flight came from a + # connection that is gone. Drop them all, so the next + # infer re-plans every env on the new connection. + plan_pos[:] = plan_len + chunk_reward[:] = 0.0 + continue timing.infer_wait += time.perf_counter() - infer_wait_started_at timing.infer_calls += 1 @@ -187,6 +198,7 @@ def log_timing_summary(final: bool = False) -> None: if done_indices.size: plan_pos[done_indices] = plan_len[done_indices] + replaced = False feedback_indices = np.nonzero(plan_pos >= plan_len)[0] if feedback_indices.size: feedback_started_at = time.perf_counter() @@ -200,15 +212,20 @@ def log_timing_summary(final: bool = False) -> None: timing.feedback_info_pack += ( time.perf_counter() - feedback_info_pack_started_at ) - agent.feedback( - obs=feedback_obs, - rewards=chunk_reward[feedback_indices], - terminated=terminated[feedback_indices], - truncated=truncated[feedback_indices], - info=feedback_info, - env_indices=feedback_indices, - step_ids=step_id[feedback_indices], - ) + try: + agent.feedback( + obs=feedback_obs, + rewards=chunk_reward[feedback_indices], + terminated=terminated[feedback_indices], + truncated=truncated[feedback_indices], + info=feedback_info, + env_indices=feedback_indices, + step_ids=step_id[feedback_indices], + ) + except ConnectionReplaced: + # Dropped, and so is every other chunk in flight. The + # episodes that ended still ended: the reset below runs. + replaced = True timing.feedback += time.perf_counter() - feedback_started_at timing.feedback_calls += 1 chunk_reward[feedback_indices] = 0.0 @@ -227,6 +244,9 @@ def log_timing_summary(final: bool = False) -> None: if recorder is not None: recorder.on_reset(obs, info, reset_indices=done_indices) step_id[done_indices] = 0 + if replaced: + plan_pos[:] = plan_len + chunk_reward[:] = 0.0 now = time.perf_counter() if now - last_timing_log_at >= 30.0: log_timing_summary(final=False) diff --git a/tests/test_reconnect_drops_feedback.py b/tests/test_reconnect_drops_feedback.py index 379dcca..3e17aef 100644 --- a/tests/test_reconnect_drops_feedback.py +++ b/tests/test_reconnect_drops_feedback.py @@ -13,7 +13,10 @@ The infer path is the opposite case and is covered here too: a fresh infer on a fresh connection is exactly how the two sides get back in step, so that one -must still retry. +must still retry - but only once the caller has dropped the chunks it was +still executing, which came from the old connection. The agent says so by +raising ConnectionReplaced, from a failed feedback and from the first infer +after a reconnect. """ import pytest @@ -21,9 +24,10 @@ from websockets.frames import Close from plugrl_env_client.agent.websocket_env_client_agent import ( + ConnectionReplaced, WebSocketEnvClientAgent, ) -from plugrl_protocol.websocket_protocol import SERVER_RESYNC_REASON +from plugrl_protocol.websocket_protocol import SERVER_RESYNC_REASON, SERVER_STOP_REASON class _Socket: @@ -43,21 +47,38 @@ def send(self, _data): def recv(self): raise AssertionError("no test here should get as far as a reply") + +class _AnsweringSocket(_Socket): + """Answers every infer with an empty action map.""" + + def __init__(self): + super().__init__(None) + + def recv(self): + from plugrl_protocol import msgpack_numpy + + return msgpack_numpy.Packer().pack({"message_type": "action", "data": {}}) + def close(self): self.closed = True -def _agent(monkeypatch, socket): - """An agent wired to `socket`, whose reconnect is observable, not real.""" +def _agent(monkeypatch, socket, *, fresh=_Socket, **kwargs): + """An agent wired to `socket`, whose reconnect is observable, not real. + + It is in the state a feedback is sent from: the infer it completes was + answered on this connection. + """ reconnects = [] def fake_wait(self): reconnects.append(True) - return _Socket(None), {} + return (fresh() if fresh is _AnsweringSocket else fresh(None)), {} monkeypatch.setattr(WebSocketEnvClientAgent, "_wait_for_server", fake_wait) - agent = WebSocketEnvClientAgent(host="127.0.0.1", port=1) + agent = WebSocketEnvClientAgent(host="127.0.0.1", port=1, **kwargs) agent._ws = socket + agent._actions_from = agent._connection_id reconnects.clear() # the constructor's own connect is not a reconnect return agent, reconnects @@ -77,6 +98,7 @@ def _feedback(agent): KEEPALIVE = ConnectionClosedError(Close(1011, "keepalive ping timeout"), None) PLAIN_CLOSE = ConnectionClosedOK(Close(1000, ""), None) RESYNC = ConnectionClosedOK(Close(1001, SERVER_RESYNC_REASON), None) +STOP = ConnectionClosedOK(Close(1001, SERVER_STOP_REASON), None) @pytest.mark.parametrize( @@ -91,18 +113,61 @@ def test_feedback_is_dropped_not_resent(monkeypatch, exc, label): socket = _Socket(exc) agent, reconnects = _agent(monkeypatch, socket) - _feedback(agent) # the point: this returns rather than retrying + with pytest.raises(ConnectionReplaced): # the caller drops its chunks too + _feedback(agent) assert socket.sends == 1, f"feedback was resent after {label}" assert reconnects == [], f"the feedback path reconnected after {label}" +def test_a_stop_the_client_waits_out_drops_the_feedback_too(monkeypatch): + """This path resent it, the one break of SPEC section 7.6 left.""" + socket = _Socket(STOP) + agent, reconnects = _agent(monkeypatch, socket, reconnect_on_server_stop=True) + + with pytest.raises(ConnectionReplaced): + _feedback(agent) + + assert socket.sends == 1, "feedback was resent after a stop" + assert reconnects == [], "the feedback path reconnected after a stop" + + +def test_feedback_never_reconnects_to_be_sent(monkeypatch): + """With the connection already gone, there is nothing to send it on.""" + agent, reconnects = _agent(monkeypatch, _Socket(None)) + agent._close_connection() + + with pytest.raises(ConnectionReplaced): + _feedback(agent) + + assert reconnects == [] + + def test_infer_still_retries_after_a_drop(monkeypatch): """The mirror image: a fresh infer is how the two sides resynchronise.""" socket = _Socket(KEEPALIVE) agent, reconnects = _agent(monkeypatch, socket) + agent._actions_from = None # nothing in flight yet with pytest.raises(AssertionError): # the retry reaches recv() and stops agent.infer({}, env_indices=[0], step_ids=[0]) assert reconnects == [True], "infer should have reconnected and retried" + + +def test_the_first_infer_after_a_reconnect_tells_the_caller(monkeypatch): + """Chunks answered on the old connection must not be finished on the new one. + + The caller is told before anything is sent, so it can drop them and ask + again for every env; that second infer goes out. + """ + socket = _Socket(KEEPALIVE) + agent, reconnects = _agent(monkeypatch, socket, fresh=_AnsweringSocket) + + with pytest.raises(ConnectionReplaced): + agent.infer({}, env_indices=[0], step_ids=[0]) + assert reconnects == [True] + assert agent._ws.sends == 0, "an infer went out before the caller re-planned" + + agent.infer({}, env_indices=[0, 1], step_ids=[0, 0]) + assert agent._ws.sends == 1 diff --git a/tests/test_rollout_connection_replaced.py b/tests/test_rollout_connection_replaced.py new file mode 100644 index 0000000..fd0b2e3 --- /dev/null +++ b/tests/test_rollout_connection_replaced.py @@ -0,0 +1,75 @@ +"""After a reconnect, no chunk from the old connection is finished or fed back. + +The rollout runs envs on chunks of different ages: an env whose episode just +ended gets a new chunk while the others are still executing theirs. When the +connection is replaced, every chunk still in flight was answered on the old +one, and the server's half of those transitions went with it (SPEC section +7.6). The rollout used to keep executing them and send their feedback on the +new connection; `plugrl-conformance --probe --scenario resync` caught it. +Now the agent raises ConnectionReplaced and the rollout drops them all, so +its first infer on the new connection asks for every env. +""" + +from __future__ import annotations + +import numpy as np + +from plugrl_env_client.agent.websocket_env_client_agent import ConnectionReplaced +from plugrl_env_client.envs.probe_env import ProbeEnv, ProbeEnvConfig +from plugrl_env_client.runner.rollout import rollout + +HORIZON = 4 + + +class _ReplacingAgent: + """Answers infers with zero chunks; replaces the connection on one feedback.""" + + def __init__(self, replace_on_feedback: int) -> None: + self.replace_on = replace_on_feedback + self.connection = 1 + self.chunk_from: dict[int, int] = {} + self.feedbacks = 0 + self.infers: list[tuple[int, list[int]]] = [] + self.stale: list[tuple[int, list[int]]] = [] + + def infer(self, obs, *, env_indices, step_ids): + envs = [int(e) for e in env_indices] + self.infers.append((self.connection, envs)) + for env in envs: + self.chunk_from[env] = self.connection + return {"action": np.zeros((HORIZON, len(envs), 3))} + + def feedback(self, *, env_indices, **_): + self.feedbacks += 1 + if self.feedbacks == self.replace_on: + self.connection += 1 + raise ConnectionReplaced("the server closed for a resync") + envs = [int(e) for e in env_indices] + stale = [e for e in envs if self.chunk_from.get(e) != self.connection] + if stale: + self.stale.append((self.connection, stale)) + + +def test_no_chunk_from_the_old_connection_is_fed_back_on_the_new_one(): + env = ProbeEnv(ProbeEnvConfig(), num_envs=3) # episodes of 3, 5 and 7 steps + # Feedback 1 is env 0 ending at t=3; env 0 then gets a new chunk. Feedback + # 2, at t=4, is envs 1 and 2 finishing theirs - and the connection goes. + agent = _ReplacingAgent(replace_on_feedback=2) + + rollout(env, agent, num_episodes=8, replan_steps=None, num_envs=3) + + assert agent.stale == [], "fed back a chunk answered on an earlier connection" + after = [envs for connection, envs in agent.infers if connection == 2] + assert after, "nothing was asked on the new connection" + assert after[0] == [0, 1, 2], "the first infer after a reconnect re-plans every env" + + +def test_an_episode_that_ended_is_reset_even_when_its_feedback_is_dropped(): + env = ProbeEnv(ProbeEnvConfig(), num_envs=3) + agent = _ReplacingAgent(replace_on_feedback=1) # env 0's terminal feedback + + rollout(env, agent, num_episodes=6, replan_steps=None, num_envs=3) + + # Had env 0 not been reset, its episode of 3 would never end again and + # the rollout would not finish six episodes. + assert agent.stale == []