amend

test/test_transforms.py

@@ -9,6 +9,7 @@
 
 import itertools
 import pickle
+import re
 import sys
 from copy import copy
 from functools import partial
@@ -4878,6 +4879,27 @@ def test_sum_reward(self, keys, device):
  def test_transform_inverse(self):
  raise pytest.skip("No inverse for RewardSum")
 
+ @pytest.mark.parametrize("in_keys", [["reward"], ["reward_1", "reward_2"]])
+ @pytest.mark.parametrize("reset_keys", [["_reset"], ["_reset1", "_reset2"]])
+ def test_keys_length_error(self, in_keys, reset_keys, batch=10):
+ t = RewardSum(in_keys=in_keys, reset_keys=reset_keys)
+ reset_dict = {
+ reset_key: torch.zeros(batch, dtype=torch.bool) for reset_key in reset_keys
+ }
+ reward_sum_dict = {out_key: torch.randn(batch) for out_key in t.out_keys}
+ reset_dict.update(reward_sum_dict)
+ td = TensorDict(reset_dict, [])
+ if len(in_keys) != len(reset_keys):
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ f"Could not match the env reset_keys {reset_keys} with the in_keys {in_keys}"
+ ),
+ ):
+ t.reset(td)
+ else:
+ t.reset(td)
+
 
 class TestReward2Go(TransformBase):
  @pytest.mark.parametrize("device", get_default_devices())

torchrl/envs/transforms/transforms.py

@@ -4692,6 +4692,7 @@ def __init__(
  """Initialises the transform. Filters out non-reward input keys and defines output keys."""
  super().__init__(in_keys=in_keys, out_keys=out_keys)
  self._reset_keys = reset_keys
+ self._keys_checked = False
 
  @property
  def in_keys(self):
@@ -4770,9 +4771,7 @@ def _check_match(reset_keys, in_keys):
  return False
  return True
 
- if len(reset_keys) != len(self.in_keys) or not _check_match(
- reset_keys, self.in_keys
- ):
+ if not _check_match(reset_keys, self.in_keys):
  raise ValueError(
  f"Could not match the env reset_keys {reset_keys} with the {type(self)} in_keys {self.in_keys}. "
  f"Please provide the reset_keys manually. Reset entries can be "
@@ -4781,6 +4780,14 @@ def _check_match(reset_keys, in_keys):
  )
  reset_keys = copy(reset_keys)
  self._reset_keys = reset_keys
+
+ if not self._keys_checked and len(reset_keys) != len(self.in_keys):
+ raise ValueError(
+ f"Could not match the env reset_keys {reset_keys} with the in_keys {self.in_keys}. "
+ "Please make sure that these have the same length."
+ )
+ self._keys_checked = True
+
  return reset_keys
 
  @reset_keys.setter