[Feature] Propagate interaction types through inference-server clients

test/test_inference_server.py

@@ -18,6 +18,11 @@
 from tensordict import lazy_stack, TensorDict
 from tensordict.base import TensorDictBase
 from tensordict.nn import TensorDictModule
+from tensordict.nn.probabilistic import (
+    interaction_type,
+    InteractionType,
+    set_interaction_type,
+)
 
 from torchrl.modules.inference_server import (
     InferenceClient,
@@ -1084,6 +1089,28 @@ def test_no_weight_sync(self):
             result = client(td)
             assert "action" in result.keys()
 
+    def test_policy_client_propagates_interaction_type(self):
+        class _InteractionPolicy(nn.Module):
+            def forward(self, td: TensorDictBase) -> TensorDictBase:
+                value = 1 if interaction_type() is InteractionType.RANDOM else 0
+                return TensorDict(
+                    {"action": torch.full(td.batch_size, value)},
+                    batch_size=td.batch_size,
+                )
+
+        transport = ThreadingTransport()
+        policy = _InteractionPolicy()
+        with InferenceServer(policy, transport, max_batch_size=4):
+            client = PolicyClientModule(transport, out_keys=["action"])
+            # The caller's exploration context reaches the server-side
+            # forward without any opt-in.
+            with set_interaction_type(InteractionType.RANDOM):
+                result = client(TensorDict({}, batch_size=[1]))
+            assert result["action"].item() == 1
+            # Without an active context the server runs a bare forward.
+            result = client(TensorDict({}, batch_size=[1]))
+            assert result["action"].item() == 0
+
 
 # ---------------------------------------------------------------------------
 # AsyncBatchedCollector tests

torchrl/modules/inference_server/_client.py

@@ -10,12 +10,26 @@
 from concurrent.futures import Future
 from typing import Any
 
+import torch
 from tensordict.base import TensorDictBase
 from tensordict.nn import TensorDictModuleBase
+from tensordict.nn.probabilistic import interaction_type
 from tensordict.utils import NestedKey
 
 from torchrl.modules.inference_server._transport import InferenceTransport
 
+_REMOTE_INTERACTION_TYPE_KEY = "_torchrl_inference_interaction_type"
+# Code stamped when the caller has no active interaction context; keeping the
+# key always present makes server batches homogeneous in key structure.
+_NO_INTERACTION_TYPE_CODE = -1
+_INTERACTION_TYPE_TO_CODE = {
+    "mode": 0,
+    "median": 1,
+    "mean": 2,
+    "random": 3,
+    "deterministic": 4,
+}
+
 
 class _ImmediateFuture:
     def __init__(self, result: TensorDictBase | BaseException):
@@ -135,6 +149,13 @@ class PolicyClientModule(TensorDictModuleBase):
             ``result()`` is first called; a timed-out ``result()`` keeps the
             slot. Must be at least ``1``. ``None`` means unbounded.
 
+    .. note::
+        The caller's active :func:`tensordict.nn.interaction_type` is
+        automatically attached to every transport request, and the server
+        executes the remote policy under that exploration context -- exactly
+        as a local policy would see it. In-process (plain callable) clients
+        need no propagation since the caller's context is already active.
+
     .. note::
         Version tracking is an instance of the generic *service-stamped
         metadata* pattern: a service may stamp every response with metadata
@@ -246,6 +267,27 @@ def submit(self, tensordict: TensorDictBase) -> Future | _ImmediateFuture:
         """
         release = self._acquire_inflight()
         submit = getattr(self.client, "submit", None)
+        if submit is not None:
+            # Cross-boundary request: carry the caller's exploration context
+            # so the server-side forward behaves like a local call. The key
+            # is always attached (with a sentinel when no context is active)
+            # so server batches stay homogeneous in key structure.
+            current_interaction_type = interaction_type()
+            code = (
+                _INTERACTION_TYPE_TO_CODE[current_interaction_type.value]
+                if current_interaction_type is not None
+                else _NO_INTERACTION_TYPE_CODE
+            )
+            tensordict = tensordict.clone(recurse=False)
+            tensordict.set(
+                _REMOTE_INTERACTION_TYPE_KEY,
+                torch.full(
+                    tensordict.batch_size,
+                    code,
+                    dtype=torch.int8,
+                    device=tensordict.device or torch.device("cpu"),
+                ),
+            )
         if submit is None:
             # The plain-callable path runs eagerly, so the request has
             # already completed here: free the slot immediately.

torchrl/modules/inference_server/_server.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
+import contextlib
 import multiprocessing as mp
 
 import queue
@@ -18,16 +19,29 @@
 import torch
 from tensordict import lazy_stack
 from tensordict.base import TensorDictBase
+from tensordict.nn.probabilistic import InteractionType, set_interaction_type
 from tensordict.utils import NestedKey
 from torch import nn
 
+from torchrl.modules.inference_server._client import (
+    _NO_INTERACTION_TYPE_CODE,
+    _REMOTE_INTERACTION_TYPE_KEY,
+)
 from torchrl.modules.inference_server._config import (
     InferenceDeviceConfig,
     InferenceServerConfig,
 )
 from torchrl.modules.inference_server._transport import InferenceTransport
 from torchrl.weight_update import SharedMemWeightSyncScheme, WeightSyncScheme
 
+_CODE_TO_INTERACTION_TYPE = {
+    0: InteractionType.MODE,
+    1: InteractionType.MEDIAN,
+    2: InteractionType.MEAN,
+    3: InteractionType.RANDOM,
+    4: InteractionType.DETERMINISTIC,
+}
+
 
 class InferenceServer:
     """Auto-batching inference server.
@@ -408,6 +422,31 @@ def _set_policy_version(self, result_batch: TensorDictBase) -> TensorDictBase:
         )
         return result_batch.set(self.policy_version_key, version)
 
+    def _interaction_type_context(self, batch: TensorDictBase):
+        code = batch.get(_REMOTE_INTERACTION_TYPE_KEY, default=None)
+        if code is None:
+            return contextlib.nullcontext(), batch
+        if not isinstance(code, torch.Tensor):
+            interaction_code = int(code)
+        else:
+            flat_code = code.reshape(-1)
+            if flat_code.numel() == 0:
+                return contextlib.nullcontext(), batch.exclude(
+                    _REMOTE_INTERACTION_TYPE_KEY, inplace=False
+                )
+            interaction_code = int(flat_code[0].item())
+            if not flat_code.eq(interaction_code).all():
+                raise RuntimeError(
+                    "InferenceServer received a mixed interaction-type batch. "
+                    "Use homogeneous server requests or a smaller max_batch_size."
+                )
+        batch = batch.exclude(_REMOTE_INTERACTION_TYPE_KEY, inplace=False)
+        if interaction_code == _NO_INTERACTION_TYPE_CODE:
+            # Sentinel: the caller had no active interaction context.
+            return contextlib.nullcontext(), batch
+        interaction_type_value = _CODE_TO_INTERACTION_TYPE[interaction_code]
+        return set_interaction_type(interaction_type_value), batch
+
     @torch.no_grad()
     def _run(self) -> None:
         self._init_weight_sync()
@@ -462,10 +501,14 @@ def _run(self) -> None:
                         batch = batch.to(self.policy_device)
                     forward_start = time.monotonic()
                     with self._model_lock:
-                        result_batch = self.model(batch)
-                    if self.output_device is not None:
-                        result_batch = result_batch.to(self.output_device)
-                    result_batch = self._set_policy_version(result_batch)
+                        interaction_context, batch = self._interaction_type_context(
+                            batch
+                        )
+                        with interaction_context:
+                            result_batch = self.model(batch)
+                        if self.output_device is not None:
+                            result_batch = result_batch.to(self.output_device)
+                        result_batch = self._set_policy_version(result_batch)
                     forward_ms = (time.monotonic() - forward_start) * 1000.0
                     self._record_batch_stats(
                         batch_size=len(callbacks),