tensordict_keys dict is not longer overwritten from child classes

Author: Blonck

test/test_cost.py

@@ -223,6 +223,16 @@ class TestLossModuleBase:
             "sample_log_prob_key": "sample_log_prob",
             "state_action_value_key": "state_action_value",
         },
+        DreamerModelLoss: {
+            "reward_key": "reward",
+            "true_reward_key": "true_reward",
+            "prior_mean_key": "prior_mean",
+            "prior_std_key": "prior_std",
+            "posterior_mean_key": "posterior_mean",
+            "posterior_std_key": "posterior_std",
+            "pixels_key": "pixels",
+            "reco_pixels_key": "reco_pixels",
+        },
         DreamerActorLoss: {
             "belief_key": "belief",
             "reward_key": "reward",
@@ -437,6 +447,81 @@ def _create_value_model(self, rssm_hidden_dim, state_dim, mlp_num_units=200):
             value_model(td)
         return value_model
 
+    def _create_world_model_model(self, rssm_hidden_dim, state_dim, mlp_num_units=200):
+        mock_env = TransformedEnv(ContinuousActionConvMockEnv(pixel_shape=[3, 64, 64]))
+        default_dict = {
+            "state": UnboundedContinuousTensorSpec(state_dim),
+            "belief": UnboundedContinuousTensorSpec(rssm_hidden_dim),
+        }
+        mock_env.append_transform(
+            TensorDictPrimer(random=False, default_value=0, **default_dict)
+        )
+
+        obs_encoder = ObsEncoder()
+        obs_decoder = ObsDecoder()
+
+        rssm_prior = RSSMPrior(
+            hidden_dim=rssm_hidden_dim,
+            rnn_hidden_dim=rssm_hidden_dim,
+            state_dim=state_dim,
+            action_spec=mock_env.action_spec,
+        )
+        rssm_posterior = RSSMPosterior(hidden_dim=rssm_hidden_dim, state_dim=state_dim)
+
+        # World Model and reward model
+        rssm_rollout = RSSMRollout(
+            SafeModule(
+                rssm_prior,
+                in_keys=["state", "belief", "action"],
+                out_keys=[
+                    ("next", "prior_mean"),
+                    ("next", "prior_std"),
+                    "_",
+                    ("next", "belief"),
+                ],
+            ),
+            SafeModule(
+                rssm_posterior,
+                in_keys=[("next", "belief"), ("next", "encoded_latents")],
+                out_keys=[
+                    ("next", "posterior_mean"),
+                    ("next", "posterior_std"),
+                    ("next", "state"),
+                ],
+            ),
+        )
+        reward_module = MLP(
+            out_features=1, depth=2, num_cells=mlp_num_units, activation_class=nn.ELU
+        )
+        # World Model and reward model
+        world_modeler = SafeSequential(
+            SafeModule(
+                obs_encoder,
+                in_keys=[("next", "pixels")],
+                out_keys=[("next", "encoded_latents")],
+            ),
+            rssm_rollout,
+            SafeModule(
+                obs_decoder,
+                in_keys=[("next", "state"), ("next", "belief")],
+                out_keys=[("next", "reco_pixels")],
+            ),
+        )
+        reward_module = SafeModule(
+            reward_module,
+            in_keys=[("next", "state"), ("next", "belief")],
+            out_keys=[("next", "reward")],
+        )
+        world_model = WorldModelWrapper(world_modeler, reward_module)
+
+        with torch.no_grad():
+            td = mock_env.rollout(10)
+            td = td.unsqueeze(0).to_tensordict()
+            td["state"] = torch.zeros((1, 10, state_dim))
+            td["belief"] = torch.zeros((1, 10, rssm_hidden_dim))
+            world_model(td)
+        return world_model
+
     def _construct_loss(self, loss_module, **kwargs):
         print(f"{loss_module = }")
         if loss_module in [
@@ -466,6 +551,9 @@ def _construct_loss(self, loss_module, **kwargs):
             actor = self._create_mock_actor(action_spec_type="one_hot")
             qvalue = self._create_mock_qvalue()
             return loss_module(actor, qvalue, actor.spec["action"].space.n, **kwargs)
+        elif loss_module in [DreamerModelLoss]:
+            world_model = self._create_world_model_model(10, 5)
+            return DreamerModelLoss(world_model)
         elif loss_module in [DreamerActorLoss]:
             mb_env = self._create_mb_env(10, 5)
             actor_model = self._create_actor_model(10, 5)
@@ -497,6 +585,7 @@ def _construct_loss(self, loss_module, **kwargs):
 
     @pytest.mark.parametrize("loss_module", LOSS_MODULES)
     def test_tensordict_keys_unknown_key(self, loss_module):
+        """Test that exception is raised if an unknown key is set via .set_keys()"""
         loss_fn = self._construct_loss(loss_module)
 
         with pytest.raises(ValueError):
@@ -512,6 +601,7 @@ def test_tensordict_keys_default_values(self, loss_module):
 
     @pytest.mark.parametrize("loss_module", LOSS_MODULES)
     def test_tensordict_set_keys(self, loss_module):
+        """Test setting of tensordict keys via .set_keys()"""
         default_keys = self.DEFAULT_KEYS[loss_module]
 
         loss_fn = self._construct_loss(loss_module)
@@ -529,6 +619,7 @@ def test_tensordict_set_keys(self, loss_module):
 
     @pytest.mark.parametrize("loss_module", LOSS_MODULES)
     def test_tensordict_deprecated_ctor(self, loss_module):
+        """Test that a warning is raised if a deprecated tensordict key is set via the ctor."""
         try:
             dep_keys = self.DEPRECATED_CTOR_KEYS[loss_module]
         except KeyError:
@@ -546,6 +637,15 @@ def test_tensordict_deprecated_ctor(self, loss_module):
                     if def_key != key:
                         assert getattr(loss_fn, def_key) == def_value
 
+    @pytest.mark.parametrize("loss_module", LOSS_MODULES)
+    def test_tensordict_all_keys_tested(self, loss_module):
+        """Check that DEFAULT_KEYS contains all available tensordict keys from each loss module."""
+        tested_keys = set(self.DEFAULT_KEYS[loss_module].keys())
+
+        loss_fn = self._construct_loss(loss_module)
+        avail_keys = set(loss_fn.tensordict_keys.keys())
+        assert avail_keys.difference(tested_keys) == set()
+
 
 class TestDQN:
     seed = 0

torchrl/objectives/a2c.py

@@ -96,25 +96,16 @@ def __init__(
             policy_params = None
         self.convert_to_functional(critic, "critic", compare_against=policy_params)
 
-        self.tensordict_keys = {
+        tensordict_keys = {
             "advantage_key": "advantage",
             "value_target_key": "value_target",
             "value_key": "state_value",
             "action_key": "action",
         }
-        if advantage_key is not None:
-            warnings.warn(
-                "Setting 'advantage_key' via ctor is deprecated, use .set_keys(advantage_key='some_key') instead.",
-                category=DeprecationWarning,
-            )
-            self.tensordict_keys["advantage_key"] = advantage_key
-        if value_target_key is not None:
-            warnings.warn(
-                "Setting 'value_target_key' via ctor is deprecated, use .set_keys(value_target_key='some_key') instead.",
-                category=DeprecationWarning,
-            )
-            self.tensordict_keys["value_target_key"] = value_target_key
-        self.set_keys(**self.tensordict_keys)
+        self._set_default_tensordict_keys(tensordict_keys)
+        self._set_deprecated_ctor_keys(
+            advantage_key=advantage_key, value_target_key=value_target_key
+        )
 
         self.samples_mc_entropy = samples_mc_entropy
         self.entropy_bonus = entropy_bonus and entropy_coef

torchrl/objectives/common.py

@@ -75,7 +75,32 @@ def __init__(self):
         # self.register_forward_pre_hook(_parameters_to_tensordict)
         self.tensordict_keys = {}
 
+    def _set_default_tensordict_keys(self, tensordict_keys):
+        """Specify which tensordict keys should be used and can be configured by this loss."""
+        self.tensordict_keys = tensordict_keys
+        self.set_keys(**self.tensordict_keys)
+
+    def _set_deprecated_ctor_keys(self, **kwargs):
+        """Helper function setting a tensordict key and creating a warning for using a deprecated argument."""
+        for key, value in kwargs.items():
+            if value is not None:
+                warnings.warn(
+                    f"Setting '{key}' via ctor is deprecated, use .set_keys(advantage_key='some_key') instead.",
+                    category=DeprecationWarning,
+                )
+                self.tensordict_keys[key] = value
+        self.set_keys(**self.tensordict_keys)
+
     def set_keys(self, **kwargs):
+        """Specify tensordict key for given argument.
+
+        Examples:
+            >>> from torchrl.objectives import DQNLoss
+            >>> # initialize the DQN loss
+            >>> actor = torch.nn.Linear(3, 4)
+            >>> dqn_loss = DQNLoss(actor, action_space="one-hot")
+            >>> dqn_loss.set_keys(priority_key="td_error", action_value_key="action_value")
+        """
         for key, value in kwargs.items():
             if key not in self.tensordict_keys.keys():
                 raise ValueError(f"{key} not a valid tensordict key")

torchrl/objectives/ddpg.py

@@ -53,11 +53,11 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        self.tensordict_keys = {
+        tensordict_keys = {
             "state_action_value_key": "state_action_value",
             "priority_key": "td_error",
         }
-        self.set_keys(**self.tensordict_keys)
+        self._set_default_tensordict_keys(tensordict_keys)
 
         self.delay_actor = delay_actor
         self.delay_value = delay_value

torchrl/objectives/dqn.py

@@ -69,18 +69,13 @@ def __init__(
     ) -> None:
 
         super().__init__()
-        self.tensordict_keys = {
+        tensordict_keys = {
             "priority_key": "td_error",
             "action_value_key": "action_value",
             "action_key": "action",
         }
-        if priority_key is not None:
-            warnings.warn(
-                "Setting 'priority_key' via ctor is deprecated, use .set_keys(priotity_key='some_key') instead.",
-                category=DeprecationWarning,
-            )
-            self.tensordict_keys["priority_key"] = priority_key
-        self.set_keys(**self.tensordict_keys)
+        self._set_default_tensordict_keys(tensordict_keys)
+        self._set_deprecated_ctor_keys(priority_key=priority_key)
 
         self.delay_value = delay_value
         value_network = ensure_tensordict_compatible(
@@ -265,20 +260,15 @@ def __init__(
         priority_key: str = None,
     ):
         super().__init__()
-        self.tensordict_keys = {
+        tensordict_keys = {
             "priority_key": "td_error",
             "action_value_key": "action_value",
             "action_key": "action",
             "reward_key": "reward",
             "done_key": "done",
         }
-        if priority_key is not None:
-            warnings.warn(
-                "Setting 'priority_key' via ctor is deprecated, use .set_keys(priotity_key='some_key') instead.",
-                category=DeprecationWarning,
-            )
-            self.tensordict_keys["priority_key"] = priority_key
-        self.set_keys(**self.tensordict_keys)
+        self._set_default_tensordict_keys(tensordict_keys)
+        self._set_deprecated_ctor_keys(priority_key=priority_key)
 
         self.register_buffer("gamma", torch.tensor(gamma))
         self.delay_value = delay_value

torchrl/objectives/dreamer.py

@@ -73,40 +73,43 @@ def __init__(
         self.delayed_clamp = delayed_clamp
         self.global_average = global_average
 
-        self.tensordict_keys = {
-            "reward_key": ("next", "reward"),
-            "prior_mean_key": ("next", "prior_mean"),
-            "prior_std_key": ("next", "prior_std"),
-            "posterior_mean_key": ("next", "posterior_mean"),
-            "posterior_std_key": ("next", "posterior_std"),
-            "pixels_key": ("next", "pixels"),
-            "reco_pixels_key": ("next", "reco_pixels"),
+        tensordict_keys = {
+            "reward_key": "reward",
+            "true_reward_key": "true_reward",
+            "prior_mean_key": "prior_mean",
+            "prior_std_key": "prior_std",
+            "posterior_mean_key": "posterior_mean",
+            "posterior_std_key": "posterior_std",
+            "pixels_key": "pixels",
+            "reco_pixels_key": "reco_pixels",
         }
-        self.set_keys(**self.tensordict_keys)
+        self._set_default_tensordict_keys(tensordict_keys)
 
     def forward(self, tensordict: TensorDict) -> torch.Tensor:
         tensordict = tensordict.clone(recurse=False)
-        tensordict.rename_key_(("next", "reward"), ("next", "true_reward"))
+        tensordict.rename_key_(
+            ("next", self.reward_key), ("next", self.true_reward_key)
+        )
         tensordict = self.world_model(tensordict)
         # compute model loss
         kl_loss = self.kl_loss(
-            tensordict.get(("next", "prior_mean")),
-            tensordict.get(("next", "prior_std")),
-            tensordict.get(("next", "posterior_mean")),
-            tensordict.get(("next", "posterior_std")),
+            tensordict.get(("next", self.prior_mean_key)),
+            tensordict.get(("next", self.prior_std_key)),
+            tensordict.get(("next", self.posterior_mean_key)),
+            tensordict.get(("next", self.posterior_std_key)),
         ).unsqueeze(-1)
         reco_loss = distance_loss(
-            tensordict.get(("next", "pixels")),
-            tensordict.get(("next", "reco_pixels")),
+            tensordict.get(("next", self.pixels_key)),
+            tensordict.get(("next", self.reco_pixels_key)),
             self.reco_loss,
         )
         if not self.global_average:
             reco_loss = reco_loss.sum((-3, -2, -1))
         reco_loss = reco_loss.mean().unsqueeze(-1)
 
         reward_loss = distance_loss(
-            tensordict.get(("next", "true_reward")),
-            tensordict.get(("next", "reward")),
+            tensordict.get(("next", self.true_reward_key)),
+            tensordict.get(("next", self.reward_key)),
             self.reward_loss,
         )
         if not self.global_average:
@@ -180,13 +183,13 @@ def __init__(
         lmbda: int = None,
     ):
         super().__init__()
-        self.tensordict_keys = {
+        tensordict_keys = {
             "belief_key": "belief",
             "reward_key": "reward",
             "value_key": "state_value",
             "done_key": "done",
         }
-        self.set_keys(**self.tensordict_keys)
+        self._set_default_tensordict_keys(tensordict_keys)
 
         self.actor_model = actor_model
         self.value_model = value_model
@@ -320,10 +323,11 @@ def __init__(
         gamma: int = 0.99,
     ):
         super().__init__()
-        self.tensordict_keys = {
+        tensordict_keys = {
             "value_key": "state_value",
         }
-        self.set_keys(**self.tensordict_keys)
+        self._set_default_tensordict_keys(tensordict_keys)
+
         self.value_model = value_model
         self.value_loss = value_loss if value_loss is not None else "l2"
         self.gamma = gamma

torchrl/objectives/iql.py

@@ -75,20 +75,15 @@ def __init__(
         if not _has_functorch:
             raise ImportError("Failed to import functorch.") from FUNCTORCH_ERROR
         super().__init__()
-        self.tensordict_keys = {
+        tensordict_keys = {
             "priority_key": "td_error",
             "log_prob_key": "_log_prob",
             "action_key": "action",
             "state_action_value_key": "state_action_value",
             "value_key": "state_value",
         }
-        if priority_key is not None:
-            warnings.warn(
-                "Setting 'priority_key' via ctor is deprecated, use .set_keys(priotity_key='some_key') instead.",
-                category=DeprecationWarning,
-            )
-            self.tensordict_keys["priority_key"] = priority_key
-        self.set_keys(**self.tensordict_keys)
+        self._set_default_tensordict_keys(tensordict_keys)
+        self._set_deprecated_ctor_keys(priority_key=priority_key)
 
         # IQL parameter
         self.temperature = temperature