agent: leave out of each infer the observations the server holds - reuse-feedback-obs, SPEC 10.1

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 <ENVIRONMENT_TYPE> [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 |

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",

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,27 +29,161 @@ 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,
         host: str = "0.0.0.0",
         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:
             self._uri += f":{port}"
         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,22 +286,27 @@ 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)
 
                 response = ws.recv()
                 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 '<empty>'}"
-                )
-                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 '<empty>'}"
+            )
+            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