[Feature] Add remote policy module for inference server clients

docs/source/reference/modules_inference_server.rst

@@ -21,6 +21,7 @@ Core API
     InferenceDeviceConfig
     ProcessInferenceServer
     InferenceClient
+    PolicyClientModule
     InferenceTransport
 
 Transport Backends
@@ -109,6 +110,22 @@ drive the collector-side transfers:
         ),
     )
 
+Remote policy module
+^^^^^^^^^^^^^^^^^^^^
+
+Use :class:`PolicyClientModule` when an actor or collector expects a regular
+TensorDict policy but inference should be served by the policy server:
+
+.. code-block:: python
+
+    remote_policy = PolicyClientModule(
+        transport,
+        in_keys=["observation"],
+        out_keys=["action"],
+    )
+
+    data = remote_policy(data)
+
 Weight Synchronisation
 ^^^^^^^^^^^^^^^^^^^^^^
 

test/test_inference_server.py

@@ -7,6 +7,7 @@
 import concurrent.futures
 import importlib.util
 import multiprocessing as mp
+import pickle
 import threading
 import time
 
@@ -25,6 +26,7 @@
     InferenceServerConfig,
     InferenceTransport,
     MPTransport,
+    PolicyClientModule,
     ProcessInferenceServer,
     RayTransport,
     SlotTransport,
@@ -407,6 +409,83 @@ def test_submit_returns_future(self):
             assert "action" in result.keys()
 
 
+def _echo_client(td):
+    """Picklable stand-in client for pickling tests."""
+    return td
+
+
+class TestPolicyClientModule:
+    def test_forward_as_tensordict_module(self):
+        transport = ThreadingTransport()
+        policy = _make_policy()
+        with InferenceServer(policy, transport, max_batch_size=4):
+            remote_policy = PolicyClientModule(
+                transport,
+                in_keys=["observation"],
+                out_keys=["action"],
+            )
+            td = TensorDict({"observation": torch.randn(4)})
+            result = remote_policy(td)
+        assert result["action"].shape == (2,)
+        assert remote_policy.in_keys == ["observation"]
+        assert remote_policy.out_keys == ["action"]
+
+    def test_submit(self):
+        transport = ThreadingTransport()
+        policy = _make_policy()
+        with InferenceServer(policy, transport, max_batch_size=4):
+            remote_policy = PolicyClientModule(transport)
+            future = remote_policy.submit(TensorDict({"observation": torch.randn(4)}))
+            result = future.result(timeout=5.0)
+        assert "action" in result.keys()
+
+    def test_nested_keys(self):
+        """in_keys/out_keys accept nested keys end to end."""
+        transport = ThreadingTransport()
+        policy = TensorDictModule(
+            nn.Linear(4, 2),
+            in_keys=[("agents", "observation")],
+            out_keys=[("agents", "action")],
+        )
+        with InferenceServer(policy, transport, max_batch_size=4):
+            remote_policy = PolicyClientModule(
+                transport,
+                in_keys=[("agents", "observation")],
+                out_keys=[("agents", "action")],
+            )
+            td = TensorDict({("agents", "observation"): torch.randn(4)})
+            result = remote_policy(td)
+        assert result["agents", "action"].shape == (2,)
+        assert remote_policy.in_keys == [("agents", "observation")]
+        assert remote_policy.out_keys == [("agents", "action")]
+
+    def test_client_contract_picklable_no_lifecycle(self):
+        """Clients pickle cleanly and expose no lifecycle methods."""
+        remote_policy = PolicyClientModule(
+            _echo_client, in_keys=["observation"], out_keys=["observation"]
+        )
+        restored = pickle.loads(pickle.dumps(remote_policy))
+        td = TensorDict({"observation": torch.randn(4)})
+        result = restored(td)
+        assert "observation" in result.keys()
+        assert restored.in_keys == ["observation"]
+        # Clients carry no lifecycle rights over the service
+        for lifecycle in ("start", "shutdown", "close", "flush"):
+            assert not hasattr(restored, lifecycle)
+
+    def test_plain_callable_client_defers_errors(self):
+        """A plain-callable client defers exceptions to result()."""
+
+        def failing_client(td):
+            raise ValueError("local policy failure")
+
+        remote_policy = PolicyClientModule(failing_client)
+        future = remote_policy.submit(TensorDict({}))
+        assert future.done()
+        with pytest.raises(ValueError, match="local policy failure"):
+            future.result()
+
+
 # =============================================================================
 # Tests: ThreadingTransport (Commit 2)
 # =============================================================================

torchrl/collectors/_async_batched.py

@@ -20,6 +20,7 @@
     InferenceDeviceConfig,
     InferenceServer,
     InferenceServerConfig,
+    PolicyClientModule,
     ProcessInferenceServer,
     ThreadingTransport,
 )
@@ -432,7 +433,10 @@ def _ensure_started(self) -> None:
         # Create clients before a process server starts so response queues are
         # inherited by the child process.
         if self._clients is None:
-            self._clients = [self._transport.client() for _ in range(self._num_envs)]
+            self._clients = [
+                PolicyClientModule(self._transport.client())
+                for _ in range(self._num_envs)
+            ]
 
         # Start inference server
         if not self._server.is_alive:

torchrl/modules/inference_server/__init__.py

@@ -3,6 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+from torchrl.modules.inference_server._client import PolicyClientModule
 from torchrl.modules.inference_server._config import (
     InferenceDeviceConfig,
     InferenceServerConfig,
@@ -27,6 +28,7 @@
     "InferenceTransport",
     "MonarchTransport",
     "MPTransport",
+    "PolicyClientModule",
     "ProcessInferenceServer",
     "RayTransport",
     "SlotTransport",

torchrl/modules/inference_server/_client.py

@@ -0,0 +1,126 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from concurrent.futures import Future
+
+from tensordict.base import TensorDictBase
+from tensordict.nn import TensorDictModuleBase
+from tensordict.utils import NestedKey
+
+from torchrl.modules.inference_server._transport import InferenceTransport
+
+
+class _ImmediateFuture:
+    def __init__(self, result: TensorDictBase | BaseException):
+        self._result = result
+
+    def done(self) -> bool:
+        return True
+
+    def result(self, timeout: float | None = None) -> TensorDictBase:
+        if isinstance(self._result, BaseException):
+            raise self._result
+        return self._result
+
+
+class PolicyClientModule(TensorDictModuleBase):
+    """TensorDict policy wrapper for remote inference-server clients.
+
+    ``PolicyClientModule`` makes a transport client look like a TorchRL policy:
+    it accepts a :class:`~tensordict.TensorDictBase`, submits it to an
+    :class:`~torchrl.modules.inference_server.InferenceServer`, and returns the
+    TensorDict produced by the remote policy. It can be passed anywhere a
+    TensorDict policy module is expected.
+
+    This class is the reference implementation of TorchRL's service *client*
+    contract: it duck-types the domain interface (a policy client IS a
+    TensorDict policy, so consumer code cannot tell local from remote), it is
+    cheap and picklable (it can be handed to spawned workers), and it carries
+    no lifecycle rights -- clients can call the service but never start or
+    shut it down; only the owner that constructed the server can.
+
+    .. note::
+        Unlike a local :class:`~tensordict.nn.TensorDictModule`, the result
+        crosses a transport boundary, so :meth:`forward` returns a *new*
+        TensorDict rather than writing the ``out_keys`` into the input
+        TensorDict. Use the return value; do not rely on in-place updates of
+        the input.
+
+    Args:
+        client (Callable or InferenceTransport): actor-side inference client.
+            If a transport is provided, ``transport.client()`` is called.
+
+    Keyword Args:
+        in_keys (sequence of NestedKey, optional): input keys advertised by the
+            module. The full input TensorDict is still sent to the server.
+        out_keys (sequence of NestedKey, optional): output keys advertised by
+            the module.
+
+    Examples:
+        >>> import torch
+        >>> import torch.nn as nn
+        >>> from tensordict import TensorDict
+        >>> from tensordict.nn import TensorDictModule
+        >>> from torchrl.modules.inference_server import (
+        ...     InferenceServer,
+        ...     PolicyClientModule,
+        ...     ThreadingTransport,
+        ... )
+        >>> policy = TensorDictModule(
+        ...     nn.Linear(4, 2), in_keys=["observation"], out_keys=["action"]
+        ... )
+        >>> transport = ThreadingTransport()
+        >>> server = InferenceServer(policy, transport).start()
+        >>> remote_policy = PolicyClientModule(
+        ...     transport, in_keys=["observation"], out_keys=["action"]
+        ... )
+        >>> td = remote_policy(TensorDict({"observation": torch.randn(4)}))
+        >>> "action" in td.keys()
+        True
+        >>> server.shutdown()
+    """
+
+    def __init__(
+        self,
+        client: Callable[[TensorDictBase], TensorDictBase] | InferenceTransport,
+        *,
+        in_keys: Sequence[NestedKey] | None = None,
+        out_keys: Sequence[NestedKey] | None = None,
+    ) -> None:
+        super().__init__()
+        if isinstance(client, InferenceTransport):
+            client = client.client()
+        self.client = client
+        self.in_keys = list(in_keys or [])
+        self.out_keys = list(out_keys or [])
+
+    def submit(self, tensordict: TensorDictBase) -> Future | _ImmediateFuture:
+        """Submit a TensorDict request and return a future-like object.
+
+        Args:
+            tensordict (TensorDictBase): observation TensorDict to send to the
+                remote policy.
+
+        Returns:
+            Future-like object whose ``result()`` method returns a TensorDict.
+            When the wrapped client exposes ``submit`` this is the transport's
+            :class:`~concurrent.futures.Future` and submission errors raise
+            synchronously; for a plain callable client the call runs eagerly
+            and errors are deferred to ``result()`` on a reduced future that
+            only implements ``done()`` and ``result()``.
+        """
+        submit = getattr(self.client, "submit", None)
+        if submit is None:
+            try:
+                result = self.client(tensordict)
+                return _ImmediateFuture(result)
+            except Exception as exc:
+                return _ImmediateFuture(exc)
+        return submit(tensordict)
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        return self.submit(tensordict).result()