[Feature] Track behavior policy versions in inference server

ghstack-source-id: 048a08f
Pull-Request: #3895
# Conflicts:
#	test/test_inference_server.py
#	torchrl/modules/inference_server/_server.py
# Conflicts:
#	torchrl/collectors/_async_batched.py
#	torchrl/modules/inference_server/_client.py

Author: vmoens

docs/source/reference/envs_transforms.rst

@@ -295,6 +295,7 @@ Available Transforms
     MultiAction
     NextObservationDelta
     NextStateReconstructor
+    PolicyAgeFilter
     NoopResetEnv
     ObservationNorm
     ObservationTransform

docs/source/reference/modules_inference_server.rst

@@ -121,11 +121,20 @@ TensorDict policy but inference should be served by the policy server:
     remote_policy = PolicyClientModule(
         transport,
         in_keys=["observation"],
-        out_keys=["action"],
+        out_keys=["action", "policy_version"],
     )
 
     data = remote_policy(data)
 
+The server writes ``policy_version`` by default so asynchronous collectors can
+track behavior-policy lag. This is the general *service-stamped metadata*
+pattern: any service may stamp its responses with metadata about the state it
+served them from, and the data pipeline may enforce freshness constraints on
+it. Bounded staleness is enforced by the replay buffer through
+:class:`~torchrl.envs.transforms.PolicyAgeFilter`, which drops elements whose
+stamped version lags the live version by more than ``max_policy_lag`` --
+either at extension time or dynamically at sampling time.
+
 Weight Synchronisation
 ^^^^^^^^^^^^^^^^^^^^^^
 

test/rb/test_rb_transforms.py

@@ -17,7 +17,7 @@
 from torchrl.data import ReplayBuffer, TensorDictReplayBuffer
 from torchrl.data.replay_buffers.samplers import RandomSampler, SliceSampler
 from torchrl.data.replay_buffers.storages import LazyMemmapStorage, LazyTensorStorage
-from torchrl.envs.transforms import NextStateReconstructor
+from torchrl.envs.transforms import NextStateReconstructor, PolicyAgeFilter
 from torchrl.envs.transforms.transforms import (
     BinarizeReward,
     CatFrames,
@@ -511,6 +511,99 @@ def test_bad_batch_dims_errors(self):
             NextStateReconstructor()(td)
 
 
+class TestPolicyAgeFilter:
+    @staticmethod
+    def _data(versions):
+        versions = torch.as_tensor(versions)
+        return TensorDict(
+            {
+                "observation": torch.randn(versions.shape[0], 3),
+                "policy_version": versions,
+            },
+            batch_size=[versions.shape[0]],
+        )
+
+    def test_filters_on_extend(self):
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(100),
+            transform=PolicyAgeFilter(3, max_policy_lag=1),
+        )
+        rb.extend(self._data([0, 2, 2, 3]))
+        assert len(rb) == 3
+
+    def test_filters_dynamically_on_sample(self):
+        current = {"version": 3}
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(100),
+            transform=PolicyAgeFilter(lambda: current["version"], max_policy_lag=1),
+        )
+        rb.extend(self._data([2, 2, 3, 3]))
+        assert len(rb) == 4
+        assert rb.sample(4).batch_size[0] == 4
+        # The policy moved on: version-2 data became stale since insertion.
+        # The random sampler draws with replacement, so the batch size is
+        # not deterministic; the invariant is that no stale element passes.
+        current["version"] = 4
+        sample = rb.sample(4)
+        assert (sample["policy_version"] == 3).all()
+        # Once everything is stale the batch is deterministically empty.
+        current["version"] = 10
+        sample = rb.sample(4)
+        assert sample.batch_size[0] == 0
+
+    def test_nested_version_key(self):
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(100),
+            transform=PolicyAgeFilter(
+                3, max_policy_lag=0, policy_version_key=("meta", "policy_version")
+            ),
+        )
+        data = TensorDict(
+            {
+                "observation": torch.randn(3, 2),
+                ("meta", "policy_version"): torch.tensor([1, 3, 3]),
+            },
+            batch_size=[3],
+        )
+        rb.extend(data)
+        assert len(rb) == 2
+
+    def test_row_uses_oldest_element(self):
+        # [B, T] data: a single stale element marks the whole row stale
+        filt = PolicyAgeFilter(5, max_policy_lag=2)
+        td = TensorDict(
+            {
+                "observation": torch.randn(2, 4),
+                "policy_version": torch.tensor([[3, 5, 5, 5], [4, 5, 5, 2]]),
+            },
+            batch_size=[2, 4],
+        )
+        out = filt(td)
+        assert out.batch_size[0] == 1
+
+    def test_missing_key_warns_and_passes_through(self):
+        filt = PolicyAgeFilter(3, max_policy_lag=1)
+        td = TensorDict({"observation": torch.randn(4, 3)}, batch_size=[4])
+        out = filt(td)
+        assert out.batch_size[0] == 4
+        with pytest.raises(KeyError, match="PolicyAgeFilter"):
+            PolicyAgeFilter(3, max_policy_lag=1, strict=True)(td)
+
+    def test_negative_lag_rejected(self):
+        with pytest.raises(ValueError, match="non-negative"):
+            PolicyAgeFilter(3, max_policy_lag=-1)
+
+    def test_noop_when_attached_to_env(self):
+        from torchrl.envs import TransformedEnv
+        from torchrl.testing.mocking_classes import CountingEnv
+
+        env = TransformedEnv(
+            CountingEnv(max_steps=4), PolicyAgeFilter(0, max_policy_lag=0)
+        )
+        rollout = env.rollout(3)
+        assert rollout.batch_size[0] == 3
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

test/test_inference_server.py

@@ -286,6 +286,7 @@ def test_stats_accounting(self):
         assert stats["batches"] >= 1
         assert stats["avg_batch_size"] > 0
         assert stats["p95_forward_ms"] >= 0
+        assert stats["policy_version"] == 0
 
     def test_structured_config(self):
         transport = ThreadingTransport()
@@ -369,6 +370,19 @@ def recording_forward(x):
         assert all(device.type == "cuda" for device in seen_devices)
         assert next(policy.parameters()).device.type == "cuda"
 
+    def test_policy_version_is_returned(self):
+        transport = ThreadingTransport()
+        policy = _make_policy()
+        with InferenceServer(
+            policy,
+            transport,
+            policy_version=12,
+            policy_version_key=("meta", "policy_version"),
+        ):
+            client = transport.client()
+            result = client(TensorDict({"observation": torch.randn(4)}))
+        assert result["meta", "policy_version"].item() == 12
+
     @pytest.mark.gpu
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
     def test_cuda_policy_cpu_output(self):
@@ -485,6 +499,15 @@ def failing_client(td):
         with pytest.raises(ValueError, match="local policy failure"):
             future.result()
 
+    def test_update_policy_weights_cascade_bumps_version(self):
+        """The weight-sync cascade hook increments the policy version."""
+        transport = ThreadingTransport()
+        policy = _make_policy()
+        server = InferenceServer(policy, transport, policy_version=0)
+        assert server.policy_version == 0
+        server.update_policy_weights_()
+        assert server.policy_version == 1
+
 
 # =============================================================================
 # Tests: ThreadingTransport (Commit 2)
@@ -889,6 +912,8 @@ def test_weight_update_applied(self):
             result_after = client(td)
             # With zero weights the linear output should be zero (bias=0 too)
             assert torch.allclose(result_after["action"], torch.zeros(2), atol=1e-6)
+            assert result_after["policy_version"].item() == 1
+            assert server.stats()["weight_updates"] == 1
 
     def test_inference_continues_after_weight_update(self):
         """The server keeps serving after a weight update."""
@@ -1213,6 +1238,20 @@ def test_process_server_backend_smoke(self):
         collector.shutdown()
         assert total >= 20
 
+    def test_policy_version_key_none_disables_annotations(self):
+        collector = AsyncBatchedCollector(
+            create_env_fn=[_counting_env_factory] * 2,
+            policy=_make_counting_policy(),
+            frames_per_batch=10,
+            total_frames=10,
+            max_batch_size=2,
+            env_backend="threading",
+            policy_version_key=None,
+        )
+        for batch in collector:
+            assert "policy_version" not in batch.keys()
+        collector.shutdown()
+
     def test_invalid_server_backend_raises(self):
         with pytest.raises(ValueError, match="backend"):
             InferenceServerConfig(backend="not-a-backend")

torchrl/collectors/_async_batched.py

@@ -11,7 +11,7 @@
 from typing import Literal
 
 import torch
-from tensordict import lazy_stack, TensorDictBase
+from tensordict import lazy_stack, NestedKey, TensorDictBase
 
 from torchrl._utils import _maybe_record_function_decorator, logger as torchrl_logger
 from torchrl.collectors._base import BaseCollector
@@ -184,6 +184,11 @@ class AsyncBatchedCollector(BaseCollector):
             placement (``policy_device``, ``output_device``, ``env_device``,
             ``storing_device``) for the whole collection pipeline. Mutually
             exclusive with ``device``.
+        policy_version (int, optional): initial behavior-policy version
+            attached to server outputs. Defaults to ``0``.
+        policy_version_key (NestedKey or None, optional): TensorDict key used
+            for behavior-policy version annotations. ``None`` disables
+            annotations. Defaults to ``"policy_version"``.
         backend (str, optional): global default backend for both
             environments and policy inference.  Specific overrides
             ``env_backend`` and ``policy_backend`` take precedence when set.
@@ -271,6 +276,8 @@ def __init__(
         create_env_kwargs: dict | list[dict] | None = None,
         server_config: InferenceServerConfig | None = None,
         device_config: InferenceDeviceConfig | None = None,
+        policy_version: int = 0,
+        policy_version_key: NestedKey | None = "policy_version",
     ):
         if policy is not None and policy_factory is not None:
             raise TypeError("policy and policy_factory are mutually exclusive.")
@@ -372,6 +379,8 @@ def __init__(
                 weight_sync_model_id=weight_sync_model_id,
                 collect_stats=_server_defaults.collect_stats,
                 stats_window_size=_server_defaults.stats_window_size,
+                policy_version=policy_version,
+                policy_version_key=policy_version_key,
             )
         else:
             self._server = InferenceServer(
@@ -386,7 +395,10 @@ def __init__(
                 weight_sync_model_id=weight_sync_model_id,
                 collect_stats=_server_defaults.collect_stats,
                 stats_window_size=_server_defaults.stats_window_size,
+                policy_version=policy_version,
+                policy_version_key=policy_version_key,
             )
+        self._policy_version_key = policy_version_key
 
         # ---- collector settings -----------------------------------------------
         self.requested_frames_per_batch = frames_per_batch
@@ -490,6 +502,11 @@ def server_stats(self, *, reset: bool = False) -> dict[str, float | int]:
             return {}
         return stats(reset=reset)
 
+    @property
+    def policy_version(self) -> int:
+        """The live behavior-policy version of the inference server."""
+        return self._server.policy_version
+
     # ------------------------------------------------------------------
     # Rollout: drain the result queue
     # ------------------------------------------------------------------

torchrl/envs/transforms/__init__.py

@@ -16,7 +16,7 @@
 from .module import ModuleTransform
 from .r3m import R3MTransform
 from .ray_service import RayTransform
-from .rb_transforms import MultiStepTransform, NextStateReconstructor
+from .rb_transforms import MultiStepTransform, NextStateReconstructor, PolicyAgeFilter
 from .rnd import RNDTransform, RunningMeanStd
 from .transforms import (
     ActionChunkTransform,
@@ -179,6 +179,7 @@ def __getattr__(name: str):
     "MultiStepTransform",
     "NextObservationDelta",
     "NextStateReconstructor",
+    "PolicyAgeFilter",
     "NoopResetEnv",
     "ObservationNorm",
     "ObservationTransform",