diff --git a/README.md b/README.md index b6bfef4..429e0af 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,13 @@ `video`. E45 measured them at 45 MB and 77 MB of the 222 MB install (`plugrl-server/experiments/e45-env-side-footprint/`). * **Distributed communication**: `websockets` plus `msgpack` for asynchronous transfer between the server and any number of env clients, across machines. + * **Each observation sent once**: a feedback carries the observation the + next chunk starts from, and the next infer used to send it again. When the + server offers `reuse-feedback-obs` (plugrl-protocol `SPEC.md` section + 10.1), the infer leaves out every observation the server already has, so + only the observation after a reset crosses the link twice. Nothing changes + with a server that does not offer it, and `--no-reuse-feedback-obs` turns + it off. * **Modular design**: Separates the environment-side runtime (`plugrl-env-client`) from the shared protocol layer (`plugrl-protocol`). * **Command-line interface**: A `tyro`-powered CLI for starting and managing env clients. * **Gymnasium integration**: Works with standard `Gymnasium` environments. @@ -123,6 +130,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/pyproject.toml b/pyproject.toml index ad53af4..7f753b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ "websockets>=15.0.1,<16.0.0", "msgpack>=1.1.1,<2.0.0", "tyro>=0.9.31,<0.10.0", - "plugrl-protocol @ git+https://github.com/PlugRL/plugrl-protocol.git@840296ecbe39dddc6fb9f3ac92479decac0797bd", + "plugrl-protocol @ git+https://github.com/PlugRL/plugrl-protocol.git@4c877bc9929fadbf49fbce8a5cd7d141bb3566a3", # The recorder's PNG snapshots. Its mp4 videos need ffmpeg: the `video` # extra. "imageio>=2.37.2,<3.0.0", 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..43f9938 100644 --- a/src/plugrl_env_client/agent/websocket_env_client_agent.py +++ b/src/plugrl_env_client/agent/websocket_env_client_agent.py @@ -2,6 +2,7 @@ from typing import Any, Dict, Optional, Tuple from loguru import logger +import numpy as np import websockets.sync.client from websockets.exceptions import ( ConnectionClosed, @@ -10,6 +11,7 @@ ) from plugrl_protocol import msgpack_numpy +from plugrl_protocol.reuse import REUSE_FEEDBACK_OBS from plugrl_protocol.websocket_protocol import ( MessageType, SERVER_RESYNC_REASON, @@ -27,10 +29,81 @@ def _get_close_details(exc: ConnectionClosed) -> tuple[int | None, str]: return None, "" +def _row(obs: Dict, i: int) -> Dict: + """Row `i` of a batched observation, copied, so later writes cannot reach it.""" + row: Dict = {} + for group, value in obs.items(): + if isinstance(value, dict): + row[group] = { + name: np.array(array[i : i + 1]) for name, array in value.items() + } + elif isinstance(value, np.ndarray) and value.ndim >= 1: + row[group] = np.array(value[i : i + 1]) + elif isinstance(value, (list, tuple)): + row[group] = list(value[i : i + 1]) + else: + row[group] = value + return row + + +def _same(a: Any, b: Any) -> bool: + if isinstance(a, dict) or isinstance(b, dict): + return ( + isinstance(a, dict) + and isinstance(b, dict) + and a.keys() == b.keys() + and all(_same(a[k], b[k]) for k in a) + ) + if isinstance(a, np.ndarray) or isinstance(b, np.ndarray): + return ( + isinstance(a, np.ndarray) + and isinstance(b, np.ndarray) + and a.shape == b.shape + and a.dtype == b.dtype + and bool(np.array_equal(a, b)) + ) + return type(a) is type(b) and a == b + + +def _take(obs: Dict, rows: np.ndarray) -> Dict: + """Rows `rows` of a batched observation, in order.""" + out: Dict = {} + for group, value in obs.items(): + if isinstance(value, dict): + out[group] = {name: array[rows] for name, array in value.items()} + elif isinstance(value, np.ndarray) and value.ndim >= 1: + out[group] = value[rows] + elif isinstance(value, (list, tuple)): + out[group] = [value[i] for i in rows.tolist()] + else: + out[group] = value + return out + + 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, @@ -38,6 +111,7 @@ def __init__( port: Optional[int] = None, api_key: Optional[str] = None, reconnect_on_server_stop: bool = False, + reuse_feedback_obs: bool = True, ) -> None: self._uri = f"ws://{host}" if port is not None: @@ -45,9 +119,71 @@ def __init__( self._packer = msgpack_numpy.Packer() self._api_key = api_key self._reconnect_on_server_stop = reconnect_on_server_stop + self._reuse_feedback_obs = reuse_feedback_obs + # SPEC section 10.1: per env, the observation its last feedback on + # this connection carried, when that feedback did not end the episode. + # Only this connection's server holds them, so a new one starts empty. + self._held: Dict[int, Dict] = {} 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 _reuses(self) -> bool: + """Whether this connection's infers leave out what the server holds.""" + features = (self._server_metadata or {}).get("features") + return ( + self._reuse_feedback_obs + and isinstance(features, list) + and REUSE_FEEDBACK_OBS in features + ) + + def _infer_message(self, obs: Dict, env_indices: Any, step_ids: Any) -> Dict: + """An infer that leaves out each row the server already holds. + + A row is left out only when the server holds that env's observation + from its last feedback here and the row is identical to it, so a + caller never has to know the feature exists. + """ + message = { + "message_type": str(MessageType.INFER), + "data": obs, + "env_indices": env_indices, + "step_ids": step_ids, + } + if not (self._reuses() and self._held): + return message + envs = np.asarray(env_indices).reshape(-1).tolist() + reuse = np.asarray( + [ + env in self._held and _same(_row(obs, i), self._held[env]) + for i, env in enumerate(envs) + ], + dtype=np.bool_, + ) + if reuse.any(): + message["data"] = _take(obs, np.nonzero(~reuse)[0]) + message["reuse"] = reuse + return message + + def _remember( + self, obs: Dict, terminated: Any, truncated: Any, env_indices: Any + ) -> None: + if not self._reuses(): + return + ended = np.logical_or(np.asarray(terminated), np.asarray(truncated)) + ended = ended.reshape(-1) + for i, env in enumerate(np.asarray(env_indices).reshape(-1).tolist()): + if ended[i]: + # It carried the terminal observation, not the reset one. + self._held.pop(env, None) + else: + self._held[env] = _row(obs, i) def _close_connection(self) -> None: + self._held.clear() if self._ws is None: return try: @@ -129,7 +265,14 @@ def _ensure_connection(self) -> None: if self._ws is not None: return logger.warning("Connection closed. Attempting to re-establish connection.") + self._held.clear() 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,15 +286,19 @@ 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( - { - "message_type": str(MessageType.INFER), - "data": obs, - "env_indices": env_indices, - "step_ids": step_ids, - } + self._infer_message(obs, env_indices, step_ids) ) ws.send(packed_data) @@ -159,6 +306,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 +368,67 @@ 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() - - 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 + 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_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) + self._remember(obs, terminated, truncated, env_indices) + 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/cli.py b/src/plugrl_env_client/cli.py index 366b1f8..c57813b 100644 --- a/src/plugrl_env_client/cli.py +++ b/src/plugrl_env_client/cli.py @@ -30,6 +30,9 @@ class Args: server_host: str = "0.0.0.0" server_port: int = 8000 reconnect_on_server_stop: bool = False + # Leave out of each infer the observations the server already holds from + # the last feedback, when it offers that (SPEC section 10.1). + reuse_feedback_obs: bool = True _CONFIGS_DICT = { @@ -54,6 +57,7 @@ def make_agent_factory(args: Args): host=args.server_host, port=args.server_port, reconnect_on_server_stop=args.reconnect_on_server_stop, + reuse_feedback_obs=args.reuse_feedback_obs, ) 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_reuse_feedback_obs.py b/tests/test_reuse_feedback_obs.py new file mode 100644 index 0000000..0abce4f --- /dev/null +++ b/tests/test_reuse_feedback_obs.py @@ -0,0 +1,169 @@ +"""SPEC section 10.1: an infer leaves out the observations the server holds. + +Every observation used to cross the link twice: in the feedback that ends a +chunk, and again in the next infer. When the server lists +`reuse-feedback-obs`, the agent marks in `reuse` each row whose observation +the server already has, and sends only the others. The server holds an env's +observation from its last feedback on this connection, unless that feedback +ended the episode, so these tests pin down when a row may be left out and +when it must not be. +""" + +import numpy as np + +from plugrl_env_client.agent.websocket_env_client_agent import ( + ConnectionReplaced, + WebSocketEnvClientAgent, +) +from plugrl_protocol import msgpack_numpy +from plugrl_protocol.reuse import REUSE_FEEDBACK_OBS + +OFFERED = {"features": [REUSE_FEEDBACK_OBS]} + + +class _Socket: + """Records every message sent; answers each infer with an empty action.""" + + def __init__(self): + self.sent: list[dict] = [] + + def send(self, data): + self.sent.append(msgpack_numpy.unpackb(data)) + + def recv(self): + return msgpack_numpy.Packer().pack({"message_type": "action", "data": {}}) + + def close(self): + pass + + def infers(self) -> list[dict]: + return [m for m in self.sent if m["message_type"] == "infer"] + + +def _agent(monkeypatch, metadata=OFFERED, **kwargs): + sockets: list[_Socket] = [] + + def fake_wait(self): + sockets.append(_Socket()) + return sockets[-1], dict(metadata) + + monkeypatch.setattr(WebSocketEnvClientAgent, "_wait_for_server", fake_wait) + agent = WebSocketEnvClientAgent(host="127.0.0.1", port=1, **kwargs) + return agent, sockets + + +def _obs(*values): + values = np.asarray(values, dtype=np.float32) + return { + "images": { + "cam": np.broadcast_to(values[:, None, None], (len(values), 2, 2)).copy() + }, + "states": {"obs": values[:, None]}, + "text": np.asarray(["go"] * len(values)), + } + + +def _infer(agent, obs, envs): + agent.infer(obs, env_indices=np.asarray(envs), step_ids=np.zeros(len(envs))) + + +def _feedback(agent, obs, envs, done=None): + done = np.asarray(done or [False] * len(envs)) + agent.feedback( + obs, + rewards=np.ones(len(envs), np.float32), + terminated=done, + truncated=np.zeros(len(envs), bool), + info={}, + env_indices=np.asarray(envs), + step_ids=np.zeros(len(envs)), + ) + + +def test_the_rows_the_server_holds_are_left_out(monkeypatch): + agent, sockets = _agent(monkeypatch) + _infer(agent, _obs(1, 2), [0, 1]) + _feedback(agent, _obs(10, 11), [0, 1]) + + _infer(agent, _obs(10, 11), [0, 1]) + + first, second = sockets[0].infers() + assert "reuse" not in first # nothing was held yet + assert second["reuse"].tolist() == [True, True] + assert second["data"]["states"]["obs"].shape == (0, 1) + assert second["data"]["images"]["cam"].shape == (0, 2, 2) + + +def test_a_row_unlike_the_one_fed_back_is_sent(monkeypatch): + agent, sockets = _agent(monkeypatch) + _infer(agent, _obs(1, 2), [0, 1]) + _feedback(agent, _obs(10, 11), [0, 1]) + + _infer(agent, _obs(10, 99), [0, 1]) + + second = sockets[0].infers()[1] + assert second["reuse"].tolist() == [True, False] + assert second["data"]["states"]["obs"].tolist() == [[99.0]] + assert second["data"]["text"].tolist() == ["go"] + + +def test_an_env_whose_feedback_ended_its_episode_sends_its_reset(monkeypatch): + """That feedback carried the terminal observation, even if the reset looks the same.""" + agent, sockets = _agent(monkeypatch) + _infer(agent, _obs(1, 2), [0, 1]) + _feedback(agent, _obs(10, 11), [0, 1], done=[True, False]) + + _infer(agent, _obs(10, 11), [0, 1]) # env 0's reset equals its terminal obs + + second = sockets[0].infers()[1] + assert second["reuse"].tolist() == [False, True] + assert second["data"]["states"]["obs"].tolist() == [[10.0]] + + +def test_nothing_is_left_out_unless_the_server_offers_it(monkeypatch): + agent, sockets = _agent(monkeypatch, metadata={"protocol_version": 1}) + _infer(agent, _obs(1), [0]) + _feedback(agent, _obs(10), [0]) + _infer(agent, _obs(10), [0]) + + assert all("reuse" not in m for m in sockets[0].infers()) + + +def test_it_can_be_turned_off(monkeypatch): + agent, sockets = _agent(monkeypatch, reuse_feedback_obs=False) + _infer(agent, _obs(1), [0]) + _feedback(agent, _obs(10), [0]) + _infer(agent, _obs(10), [0]) + + assert all("reuse" not in m for m in sockets[0].infers()) + + +def test_a_new_connection_holds_nothing(monkeypatch): + """SPEC section 10.1: a client MUST NOT reuse just after a reconnect.""" + agent, sockets = _agent(monkeypatch) + _infer(agent, _obs(1), [0]) + _feedback(agent, _obs(10), [0]) + agent._close_connection() # the server went away + try: + _infer(agent, _obs(10), [0]) + except ConnectionReplaced: + pass # the caller drops its chunks and asks again + _infer(agent, _obs(10), [0]) + + assert len(sockets) == 2 + assert all("reuse" not in m for m in sockets[1].infers()) + + +def test_an_observation_changed_in_place_after_its_feedback_is_sent(monkeypatch): + """The agent keeps its own copy, so a caller reusing buffers cannot fool it.""" + agent, sockets = _agent(monkeypatch) + obs = _obs(10) + _infer(agent, _obs(1), [0]) + _feedback(agent, obs, [0]) + obs["states"]["obs"][0, 0] = 42.0 # the same buffer, a new step + + _infer(agent, obs, [0]) + + second = sockets[0].infers()[1] + assert "reuse" not in second + assert second["data"]["states"]["obs"].tolist() == [[42.0]] 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 == [] diff --git a/uv.lock b/uv.lock index 64a445e..f8225b3 100644 --- a/uv.lock +++ b/uv.lock @@ -2388,7 +2388,7 @@ requires-dist = [ { name = "mujoco", marker = "extra == 'libero'", specifier = "==2.3.7" }, { name = "mujoco", marker = "extra == 'robomimic'", specifier = "==2.3.7" }, { name = "patchelf", marker = "extra == 'd4rl'", specifier = ">=0.17.2.4,<0.18.0" }, - { name = "plugrl-protocol", git = "https://github.com/PlugRL/plugrl-protocol.git?rev=840296ecbe39dddc6fb9f3ac92479decac0797bd" }, + { name = "plugrl-protocol", git = "https://github.com/PlugRL/plugrl-protocol.git?rev=4c877bc9929fadbf49fbce8a5cd7d141bb3566a3" }, { name = "pyopengl", marker = "extra == 'robomimic'", specifier = "==3.1.4" }, { name = "robomimic", marker = "extra == 'robomimic'", git = "https://github.com/ARISE-Initiative/robomimic.git?rev=v0.4.0" }, { name = "robosuite", marker = "extra == 'robomimic'", specifier = "==1.4.1" }, @@ -2407,7 +2407,7 @@ dev = [ [[package]] name = "plugrl-protocol" version = "0.1.0" -source = { git = "https://github.com/PlugRL/plugrl-protocol.git?rev=840296ecbe39dddc6fb9f3ac92479decac0797bd#840296ecbe39dddc6fb9f3ac92479decac0797bd" } +source = { git = "https://github.com/PlugRL/plugrl-protocol.git?rev=4c877bc9929fadbf49fbce8a5cd7d141bb3566a3#4c877bc9929fadbf49fbce8a5cd7d141bb3566a3" } dependencies = [ { name = "msgpack" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-17-plugrl-env-client-libero' and extra == 'extra-17-plugrl-env-client-mujoco') or (extra == 'extra-17-plugrl-env-client-mujoco' and extra == 'extra-17-plugrl-env-client-robomimic')" },