[Feature] VIPRewardTransform (#658)

Author: vmoens

test/test_transforms.py

@@ -57,7 +57,7 @@
     TensorDictPrimer,
     UnsqueezeTransform,
 )
-from torchrl.envs.transforms.vip import _VIPNet
+from torchrl.envs.transforms.vip import _VIPNet, VIPRewardTransform
 
 if _has_gym:
     import gym
@@ -1687,6 +1687,71 @@ def test_vip_parallel(self, model, device):
         transformed_env.close()
         del transformed_env
 
+    def test_vip_parallel_reward(self, model, device):
+        keys_in = ["next_pixels"]
+        keys_out = ["next_vec"]
+        tensor_pixels_key = None
+        vip = VIPRewardTransform(
+            model,
+            keys_in=keys_in,
+            keys_out=keys_out,
+            tensor_pixels_keys=tensor_pixels_key,
+        )
+        base_env = ParallelEnv(4, lambda: DiscreteActionConvMockEnvNumpy().to(device))
+        transformed_env = TransformedEnv(base_env, vip)
+        tensordict_reset = TensorDict(
+            {"goal_image": torch.randint(0, 255, (4, 7, 7, 3), dtype=torch.uint8)},
+            [4],
+            device=device,
+        )
+        with pytest.raises(
+            KeyError,
+            match=r"VIPRewardTransform.* requires .* key to be present in the input tensordict",
+        ):
+            _ = transformed_env.reset()
+        with pytest.raises(
+            KeyError,
+            match=r"VIPRewardTransform.* requires .* key to be present in the input tensordict",
+        ):
+            _ = transformed_env.reset(tensordict_reset.select())
+
+        td = transformed_env.reset(tensordict_reset)
+        assert td.device == device
+        assert td.batch_size == torch.Size([4])
+        exp_keys = {"vec", "done", "pixels_orig", "goal_embedding", "goal_image"}
+        if tensor_pixels_key:
+            exp_keys.add(tensor_pixels_key)
+        assert set(td.keys()) == exp_keys
+
+        td = transformed_env.rand_step(td)
+        exp_keys = exp_keys.union({"next_vec", "next_pixels_orig", "action", "reward"})
+        assert set(td.keys()) == exp_keys, td
+
+        tensordict_reset = TensorDict(
+            {"goal_image": torch.randint(0, 255, (4, 7, 7, 3), dtype=torch.uint8)},
+            [4],
+            device=device,
+        )
+        td = transformed_env.rollout(
+            3, auto_reset=False, tensordict=transformed_env.reset(tensordict_reset)
+        )
+        assert set(td.keys()) == exp_keys, td
+        # test that we do compute the reward we want
+        cur_embedding = td["next_vec"]
+        goal_embedding = td["goal_embedding"]
+        last_embedding = td["vec"]
+        explicit_reward = -torch.norm(cur_embedding - goal_embedding, dim=-1) - (
+            -torch.norm(last_embedding - goal_embedding, dim=-1)
+        )
+        torch.testing.assert_close(explicit_reward, td["reward"].squeeze())
+        # test that there is only one goal embedding
+        goal = td["goal_embedding"]
+        goal_expand = td["goal_embedding"][:, :1].expand_as(td["goal_embedding"])
+        torch.testing.assert_close(goal, goal_expand)
+
+        transformed_env.close()
+        del transformed_env
+
     @pytest.mark.parametrize("del_keys", [True, False])
     @pytest.mark.parametrize(
         "in_keys",

torchrl/envs/transforms/transforms.py

@@ -431,7 +431,11 @@ def set_seed(self, seed: int, static_seed: bool = False) -> int:
         return self.base_env.set_seed(seed, static_seed=static_seed)
 
     def _reset(self, tensordict: Optional[TensorDictBase] = None, **kwargs):
-        out_tensordict = self.base_env.reset(execute_step=False, **kwargs)
+        if tensordict is not None:
+            tensordict = tensordict.clone(recurse=False)
+        out_tensordict = self.base_env.reset(
+            tensordict=tensordict, execute_step=False, **kwargs
+        )
         out_tensordict = self.transform.reset(out_tensordict)
         out_tensordict = self.transform(out_tensordict)
         return out_tensordict

torchrl/envs/transforms/vip.py

@@ -14,6 +14,7 @@
     CompositeSpec,
     NdUnboundedContinuousTensorSpec,
 )
+from torchrl.data.tensordict.tensordict import TensorDictBase
 from torchrl.envs.transforms import (
     ToTensorImage,
     Compose,
@@ -306,3 +307,45 @@ def dtype(self):
     transform_reward_spec = _init_first(Compose.transform_reward_spec)
     reset = _init_first(Compose.reset)
     init = _init_first(Compose.init)
+
+
+class VIPRewardTransform(VIPTransform):
+    """A VIP transform to compute rewards based on embedded similarity.
+
+    This class will update the reward computation
+    """
+
+    def reset(self, tensordict: TensorDictBase) -> TensorDictBase:
+        if "goal_embedding" not in tensordict.keys():
+            tensordict = self._embed_goal(tensordict)
+        return super().reset(tensordict)
+
+    def _embed_goal(self, tensordict):
+        if "goal_image" not in tensordict.keys():
+            raise KeyError(
+                f"{self.__class__.__name__}.reset() requires a `'goal_image'` key to be "
+                f"present in the input tensordict."
+            )
+        tensordict_in = tensordict.select("goal_image").rename_key(
+            "goal_image", self.keys_in[0]
+        )
+        tensordict_in = super(VIPRewardTransform, self).forward(tensordict_in)
+        tensordict = tensordict.update(
+            tensordict_in.rename_key(self.keys_out[0], "goal_embedding")
+        )
+        return tensordict
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        if "goal_embedding" not in tensordict.keys():
+            tensordict = self._embed_goal(tensordict)
+        tensordict = super().forward(tensordict)
+        cur_embedding = tensordict.get(self.keys_out[0])
+        last_embedding_key = self.keys_out[0].split("next_")[1]
+        last_embedding = tensordict.get(last_embedding_key, None)
+        if last_embedding is not None:
+            goal_embedding = tensordict["goal_embedding"]
+            reward = -torch.norm(cur_embedding - goal_embedding, dim=-1) - (
+                -torch.norm(last_embedding - goal_embedding, dim=-1)
+            )
+            tensordict.set("reward", reward)
+        return tensordict