[Feature] ConditionalPolicySwitch transform

ghstack-source-id: eee6f67
Pull Request resolved: #2711

Author: Vincent Moens

docs/source/reference/envs.rst

@@ -816,6 +816,7 @@ to be able to create this other composition:
     CenterCrop
     ClipTransform
     Compose
+    ConditionalPolicySwitch
     Crop
     DTypeCastTransform
     DeviceCastTransform

examples/agents/ppo-chess.py

@@ -5,20 +5,24 @@
 import tensordict.nn
 import torch
 import tqdm
-from tensordict.nn import TensorDictSequential as TDSeq, TensorDictModule as TDMod, \
-    ProbabilisticTensorDictModule as TDProb, ProbabilisticTensorDictSequential as TDProbSeq
+from tensordict.nn import (
+    ProbabilisticTensorDictModule as TDProb,
+    ProbabilisticTensorDictSequential as TDProbSeq,
+    TensorDictModule as TDMod,
+    TensorDictSequential as TDSeq,
+)
 from torch import nn
 from torch.nn.utils import clip_grad_norm_
 from torch.optim import Adam
 
 from torchrl.collectors import SyncDataCollector
+from torchrl.data import LazyTensorStorage, ReplayBuffer, SamplerWithoutReplacement
 
 from torchrl.envs import ChessEnv, Tokenizer
 from torchrl.modules import MLP
 from torchrl.modules.distributions import MaskedCategorical
 from torchrl.objectives import ClipPPOLoss
 from torchrl.objectives.value import GAE
-from torchrl.data import ReplayBuffer, LazyTensorStorage, SamplerWithoutReplacement
 
 tensordict.nn.set_composite_lp_aggregate(False)
 
@@ -39,7 +43,9 @@
 embedding_moves = nn.Embedding(num_embeddings=n + 1, embedding_dim=64)
 
 # Embedding for the fen
-embedding_fen = nn.Embedding(num_embeddings=transform.tokenizer.vocab_size, embedding_dim=64)
+embedding_fen = nn.Embedding(
+    num_embeddings=transform.tokenizer.vocab_size, embedding_dim=64
+)
 
 backbone = MLP(out_features=512, num_cells=[512] * 8, activation_class=nn.ReLU)
 
@@ -49,20 +55,30 @@
 critic_head = nn.Linear(512, 1)
 critic_head.bias.data.fill_(0)
 
-prob = TDProb(in_keys=["logits", "mask"], out_keys=["action"], distribution_class=MaskedCategorical, return_log_prob=True)
+prob = TDProb(
+    in_keys=["logits", "mask"],
+    out_keys=["action"],
+    distribution_class=MaskedCategorical,
+    return_log_prob=True,
+)
+
 
 def make_mask(idx):
     mask = idx.new_zeros((*idx.shape[:-1], n + 1), dtype=torch.bool)
     return mask.scatter_(-1, idx, torch.ones_like(idx, dtype=torch.bool))[..., :-1]
 
+
 actor = TDProbSeq(
-    TDMod(
-        make_mask,
-        in_keys=["legal_moves"], out_keys=["mask"]),
+    TDMod(make_mask, in_keys=["legal_moves"], out_keys=["mask"]),
     TDMod(embedding_moves, in_keys=["legal_moves"], out_keys=["embedded_legal_moves"]),
     TDMod(embedding_fen, in_keys=["fen_tokenized"], out_keys=["embedded_fen"]),
-    TDMod(lambda *args: torch.cat([arg.view(*arg.shape[:-2], -1) for arg in args], dim=-1), in_keys=["embedded_legal_moves", "embedded_fen"],
-          out_keys=["features"]),
+    TDMod(
+        lambda *args: torch.cat(
+            [arg.view(*arg.shape[:-2], -1) for arg in args], dim=-1
+        ),
+        in_keys=["embedded_legal_moves", "embedded_fen"],
+        out_keys=["features"],
+    ),
     TDMod(backbone, in_keys=["features"], out_keys=["hidden"]),
     TDMod(actor_head, in_keys=["hidden"], out_keys=["logits"]),
     prob,
@@ -78,7 +94,9 @@ def make_mask(idx):
 
 optim = Adam(loss.parameters())
 
-gae = GAE(value_network=TDSeq(*actor[:-2], critic), gamma=0.99, lmbda=0.95, shifted=True)
+gae = GAE(
+    value_network=TDSeq(*actor[:-2], critic), gamma=0.99, lmbda=0.95, shifted=True
+)
 
 # Create a data collector
 collector = SyncDataCollector(
@@ -88,12 +106,20 @@ def make_mask(idx):
     total_frames=1_000_000,
 )
 
-replay_buffer0 = ReplayBuffer(storage=LazyTensorStorage(max_size=collector.frames_per_batch//2), batch_size=batch_size, sampler=SamplerWithoutReplacement())
-replay_buffer1 = ReplayBuffer(storage=LazyTensorStorage(max_size=collector.frames_per_batch//2), batch_size=batch_size, sampler=SamplerWithoutReplacement())
+replay_buffer0 = ReplayBuffer(
+    storage=LazyTensorStorage(max_size=collector.frames_per_batch // 2),
+    batch_size=batch_size,
+    sampler=SamplerWithoutReplacement(),
+)
+replay_buffer1 = ReplayBuffer(
+    storage=LazyTensorStorage(max_size=collector.frames_per_batch // 2),
+    batch_size=batch_size,
+    sampler=SamplerWithoutReplacement(),
+)
 
 for data in tqdm.tqdm(collector):
     data = data.filter_non_tensor_data()
-    print('data', data[0::2])
+    print("data", data[0::2])
     for i in range(num_epochs):
         replay_buffer0.empty()
         replay_buffer1.empty()
@@ -103,14 +129,24 @@ def make_mask(idx):
             # player 1
             data1 = gae(data[1::2])
             if i == 0:
-                print('win rate for 0', data0["next", "reward"].sum() / data["next", "done"].sum().clamp_min(1e-6))
-                print('win rate for 1', data1["next", "reward"].sum() / data["next", "done"].sum().clamp_min(1e-6))
+                print(
+                    "win rate for 0",
+                    data0["next", "reward"].sum()
+                    / data["next", "done"].sum().clamp_min(1e-6),
+                )
+                print(
+                    "win rate for 1",
+                    data1["next", "reward"].sum()
+                    / data["next", "done"].sum().clamp_min(1e-6),
+                )
 
             replay_buffer0.extend(data0)
             replay_buffer1.extend(data1)
 
-        n_iter = collector.frames_per_batch//(2 * batch_size)
-        for (d0, d1) in tqdm.tqdm(zip(replay_buffer0, replay_buffer1, strict=True), total=n_iter):
+        n_iter = collector.frames_per_batch // (2 * batch_size)
+        for (d0, d1) in tqdm.tqdm(
+            zip(replay_buffer0, replay_buffer1, strict=True), total=n_iter
+        ):
             loss_vals = (loss(d0) + loss(d1)) / 2
             loss_vals.sum(reduce=True).backward()
             gn = clip_grad_norm_(loss.parameters(), 100.0)

test/test_transforms.py

@@ -106,6 +106,7 @@
     CenterCrop,
     ClipTransform,
     Compose,
+    ConditionalPolicySwitch,
     Crop,
     DeviceCastTransform,
     DiscreteActionProjection,
@@ -13192,6 +13193,127 @@ def test_composite_reward_spec(self) -> None:
         assert transform.transform_reward_spec(reward_spec) == expected_reward_spec
 
 
+class TestConditionalPolicySwitch(TransformBase):
+    def test_single_trans_env_check(self):
+        base_env = CountingEnv(max_steps=15)
+        condition = lambda td: ((td.get("step_count") % 2) == 0).all()
+        # Player 0
+        policy_odd = lambda td: td.set("action", env.action_spec.zero())
+        policy_even = lambda td: td.set("action", env.action_spec.one())
+        transforms = Compose(
+            StepCounter(),
+            ConditionalPolicySwitch(condition=condition, policy=policy_even),
+        )
+        env = base_env.append_transform(transforms)
+        r = env.rollout(1000, policy_odd, break_when_all_done=True)
+        assert r.shape[0] == 15
+        assert (r["action"] == 0).all()
+        assert (r["step_count"] == torch.arange(1, r.numel() * 2, 2).unsqueeze(-1)).all()
+        assert r["next", "done"].any()
+
+        # Player 1
+        condition = lambda td: ((td.get("step_count") % 2) == 1).all()
+        transforms = Compose(
+            StepCounter(),
+            ConditionalPolicySwitch(condition=condition, policy=policy_odd),
+        )
+        env = base_env.append_transform(transforms)
+        r = env.rollout(1000, policy_even, break_when_all_done=True)
+        assert r.shape[0] == 16
+        assert (r["action"] == 1).all()
+        assert (r["step_count"] == torch.arange(0, r.numel() * 2, 2).unsqueeze(-1)).all()
+        assert r["next", "done"].any()
+
+
+    def test_trans_serial_env_check(self):
+        def make_env(max_count):
+            def make():
+                base_env = CountingEnv(max_steps=max_count)
+                transforms =
+                return base_env.append_transform(transforms)
+            return make
+
+        base_env = SerialEnv(3,
+                        [partial(CountingEnv, 6), partial(CountingEnv, 7), partial(CountingEnv, 8)])
+        condition = lambda td: ((td.get("step_count") % 2) == 0)
+        policy_odd = lambda td, base_env=base_env: td.set("action", base_env.action_spec.zero())
+        policy_even = lambda td, base_env=base_env: td.set("action", base_env.action_spec.one())
+        env = base_env.append_transform(Compose(
+                    StepCounter(),
+                    ConditionalPolicySwitch(condition=condition, policy=policy_even),
+                ))
+        r = env.rollout(100, break_when_all_done=False)
+        print(r["step_count"].squeeze())
+
+
+    def test_trans_parallel_env_check(self):
+        """tests that a transformed paprallel env (TransformedEnv(ParallelEnv(N, lambda: env()), transform)) passes the check_env_specs test."""
+        raise NotImplementedError
+
+    def test_serial_trans_env_check(self):
+        condition = lambda td: ((td.get("step_count") % 2) == 0).all()
+        # Player 0
+        policy_odd = lambda td: td.set("action", env.action_spec.zero())
+        policy_even = lambda td: td.set("action", env.action_spec.one())
+        def make_env(max_count):
+            def make():
+                base_env = CountingEnv(max_steps=max_count)
+                transforms = Compose(
+                    StepCounter(),
+                    ConditionalPolicySwitch(condition=condition, policy=policy_even),
+                )
+                return base_env.append_transform(transforms)
+            return make
+
+        env = SerialEnv(3,
+                        [make_env(6), make_env(7), make_env(8)])
+        r = env.rollout(100, break_when_all_done=False)
+        print(r["step_count"].squeeze())
+
+    def test_parallel_trans_env_check(self):
+        """tests that a parallel transformed env (ParallelEnv(N, lambda: TransformedEnv(env, transform))) passes the check_env_specs test."""
+        raise NotImplementedError
+
+    def test_transform_no_env(self):
+        """tests the transform on dummy data, without an env."""
+        raise NotImplementedError
+
+    def test_transform_compose(self):
+        """tests the transform on dummy data, without an env but inside a Compose."""
+        raise NotImplementedError
+
+    def test_transform_env(self):
+        """tests the transform on a real env.
+
+        If possible, do not use a mock env, as bugs may go unnoticed if the dynamic is too
+        simplistic. A call to reset() and step() should be tested independently, ie
+        a check that reset produces the desired output and that step() does too.
+
+        """
+        raise NotImplementedError
+
+    def test_transform_model(self):
+        """tests the transform before an nn.Module that reads the output."""
+        raise NotImplementedError
+
+    def test_transform_rb(self):
+        """tests the transform when used with a replay buffer.
+
+        If your transform is not supposed to work with a replay buffer, test that
+        an error will be raised when called or appended to a RB.
+
+        """
+        raise NotImplementedError
+
+    def test_transform_inverse(self):
+        """tests the inverse transform. If not applicable, simply skip this test.
+
+        If your transform is not supposed to work offline, test that
+        an error will be raised when called in a nn.Module.
+        """
+        raise NotImplementedError
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/envs/__init__.py

@@ -55,6 +55,7 @@
     CenterCrop,
     ClipTransform,
     Compose,
+    ConditionalPolicySwitch,
     Crop,
     DeviceCastTransform,
     DiscreteActionProjection,

torchrl/envs/transforms/__init__.py

@@ -20,6 +20,7 @@
     CenterCrop,
     ClipTransform,
     Compose,
+    ConditionalPolicySwitch,
     Crop,
     DeviceCastTransform,
     DiscreteActionProjection,

torchrl/envs/transforms/transforms.py

@@ -9974,3 +9974,53 @@ def _apply_transform(self, reward: Tensor) -> TensorDictBase:
             )
 
         return (self.weights * reward).sum(dim=-1)
+
+
+class ConditionalPolicySwitch(Transform):
+    def __init__(
+        self,
+        policy: Callable[[TensorDictBase], TensorDictBase],
+        condition: Callable[[TensorDictBase], bool],
+    ):
+        super().__init__([], [])
+        self.__dict__["policy"] = policy
+        self.condition = condition
+
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        cond = self.condition(next_tensordict)
+        if not isinstance(cond, (bool, torch.Tensor)):
+            raise RuntimeError("Calling the condition function should return a boolean or a tensor.")
+        if isinstance(cond, (torch.Tensor,)) and cond.shape not in ((1,), (), tensordict.shape):
+            raise RuntimeError("Tenspr outputs must have the shape of the tensordict, or contain a single element.")
+        if cond.any():
+            parent: TransformedEnv = self.parent
+            done = next_tensordict.get("done")
+            next_td_save = None
+            if done.any():
+                if next_tensordict.numel() == 1 or done.all():
+                    return next_tensordict
+                if parent.base_env.batch_locked:
+                    raise RuntimeError("Cannot run partial steps in a batched locked environment")
+                done = done.view(next_tensordict.shape)
+                next_td_save = next_tensordict[done]
+                next_tensordict = next_tensordict[~done]
+                tensordict = tensordict[~done]
+            td = self.policy(
+                parent.step_mdp(tensordict.copy().set("next", next_tensordict))
+            )
+            next_tensordict = parent._step(td)
+            if next_td_save is not None:
+                return torch.where(done, next_td_save, next_tensordict)
+            return next_tensordict
+        return next_tensordict
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        if self.condition(tensordict_reset):
+            parent: TransformedEnv = self.parent
+            td = self.policy(tensordict_reset)
+            return parent._step(td).exclude(*parent.reward_keys)
+        return tensordict_reset