[Naming] Rename keys_in to in_keys in transforms.py and related modules (#656)

* rename keys_in to in_keys

* update format by pre-commit command

Author: sardaankita

examples/dreamer/dreamer_utils.py

@@ -93,13 +93,13 @@ def make_env_transforms(
         if cfg.grayscale:
             env.append_transform(GrayScale())
         env.append_transform(FlattenObservation())
-        env.append_transform(CatFrames(N=cfg.catframes, keys_in=["next_pixels"]))
+        env.append_transform(CatFrames(N=cfg.catframes, in_keys=["next_pixels"]))
         if stats is None:
             obs_stats = {"loc": 0.0, "scale": 1.0}
         else:
             obs_stats = stats
         obs_stats["standard_normal"] = True
-        env.append_transform(ObservationNorm(**obs_stats, keys_in=["next_pixels"]))
+        env.append_transform(ObservationNorm(**obs_stats, in_keys=["next_pixels"]))
     if norm_rewards:
         reward_scaling = 1.0
         reward_loc = 0.0
@@ -118,7 +118,7 @@ def make_env_transforms(
         ]
         float_to_double_list += ["action"]  # DMControl requires double-precision
     env.append_transform(
-        DoubleToFloat(keys_in=double_to_float_list, keys_inv_in=float_to_double_list)
+        DoubleToFloat(in_keys=double_to_float_list, in_keys_inv=float_to_double_list)
     )
 
     default_dict = {

test/test_env.py

@@ -254,7 +254,7 @@ def create_env_fn():
                     GymEnv(env_name, frame_skip=frame_skip, device=device),
                     Compose(
                         ObservationNorm(
-                            keys_in=["next_observation"], loc=0.5, scale=1.1
+                            in_keys=["next_observation"], loc=0.5, scale=1.1
                         ),
                         RewardClipping(0, 0.1),
                     ),
@@ -275,7 +275,7 @@ def t_out():
                     Compose(*[ToTensorImage(), RewardClipping(0, 0.1)])
                     if not transformed_in
                     else Compose(
-                        *[ObservationNorm(keys_in=["next_pixels"], loc=0, scale=1)]
+                        *[ObservationNorm(in_keys=["next_pixels"], loc=0, scale=1)]
                     )
                 )
 
@@ -297,14 +297,14 @@ def t_out():
                 return (
                     Compose(
                         ObservationNorm(
-                            keys_in=["next_observation"], loc=0.5, scale=1.1
+                            in_keys=["next_observation"], loc=0.5, scale=1.1
                         ),
                         RewardClipping(0, 0.1),
                     )
                     if not transformed_in
                     else Compose(
                         ObservationNorm(
-                            keys_in=["next_observation"], loc=1.0, scale=1.0
+                            in_keys=["next_observation"], loc=1.0, scale=1.0
                         )
                     )
                 )
@@ -458,8 +458,8 @@ def env1_maker():
                     CatTensors(env1_obs_keys, "next_observation_stand", del_keys=False),
                     CatTensors(env1_obs_keys, "next_observation"),
                     DoubleToFloat(
-                        keys_in=["next_observation_stand", "next_observation"],
-                        keys_inv_in=["action"],
+                        in_keys=["next_observation_stand", "next_observation"],
+                        in_keys_inv=["action"],
                     ),
                 ),
             )
@@ -471,8 +471,8 @@ def env2_maker():
                     CatTensors(env2_obs_keys, "next_observation_walk", del_keys=False),
                     CatTensors(env2_obs_keys, "next_observation"),
                     DoubleToFloat(
-                        keys_in=["next_observation_walk", "next_observation"],
-                        keys_inv_in=["action"],
+                        in_keys=["next_observation_walk", "next_observation"],
+                        in_keys_inv=["action"],
                     ),
                 ),
             )

test/test_transforms.py

@@ -406,7 +406,7 @@ class TestTransforms:
     def test_resize(self, interpolation, keys, nchannels, batch, device):
         torch.manual_seed(0)
         dont_touch = torch.randn(*batch, nchannels, 16, 16, device=device)
-        resize = Resize(w=20, h=21, interpolation=interpolation, keys_in=keys)
+        resize = Resize(w=20, h=21, interpolation=interpolation, in_keys=keys)
         td = TensorDict(
             {
                 key: torch.randn(*batch, nchannels, 16, 16, device=device)
@@ -444,7 +444,7 @@ def test_resize(self, interpolation, keys, nchannels, batch, device):
     def test_centercrop(self, keys, h, nchannels, batch, device):
         torch.manual_seed(0)
         dont_touch = torch.randn(*batch, nchannels, 16, 16, device=device)
-        cc = CenterCrop(w=20, h=h, keys_in=keys)
+        cc = CenterCrop(w=20, h=h, in_keys=keys)
         if h is None:
             h = 20
         td = TensorDict(
@@ -485,7 +485,7 @@ def test_flatten(self, keys, size, nchannels, batch, device):
         torch.manual_seed(0)
         dont_touch = torch.randn(*batch, *size, nchannels, 16, 16, device=device)
         start_dim = -3 - len(size)
-        flatten = FlattenObservation(start_dim, -3, keys_in=keys)
+        flatten = FlattenObservation(start_dim, -3, in_keys=keys)
         td = TensorDict(
             {
                 key: torch.randn(*batch, *size, nchannels, 16, 16, device=device)
@@ -527,7 +527,7 @@ def test_flatten(self, keys, size, nchannels, batch, device):
     def test_unsqueeze(self, keys, size, nchannels, batch, device, unsqueeze_dim):
         torch.manual_seed(0)
         dont_touch = torch.randn(*batch, *size, nchannels, 16, 16, device=device)
-        unsqueeze = UnsqueezeTransform(unsqueeze_dim, keys_in=keys)
+        unsqueeze = UnsqueezeTransform(unsqueeze_dim, in_keys=keys)
         td = TensorDict(
             {
                 key: torch.randn(*batch, *size, nchannels, 16, 16, device=device)
@@ -586,7 +586,7 @@ def test_unsqueeze_inv(
         torch.manual_seed(0)
         keys_total = set(keys + keys_inv)
         unsqueeze = UnsqueezeTransform(
-            unsqueeze_dim, keys_in=keys, keys_inv_in=keys_inv
+            unsqueeze_dim, in_keys=keys, in_keys_inv=keys_inv
         )
         td = TensorDict(
             {
@@ -621,7 +621,7 @@ def test_unsqueeze_inv(
     def test_squeeze(self, keys, keys_inv, size, nchannels, batch, device, squeeze_dim):
         torch.manual_seed(0)
         keys_total = set(keys + keys_inv)
-        squeeze = SqueezeTransform(squeeze_dim, keys_in=keys, keys_inv_in=keys_inv)
+        squeeze = SqueezeTransform(squeeze_dim, in_keys=keys, in_keys_inv=keys_inv)
         td = TensorDict(
             {
                 key: torch.randn(*batch, *size, nchannels, 16, 16, device=device)
@@ -656,7 +656,7 @@ def test_squeeze_inv(
     ):
         torch.manual_seed(0)
         keys_total = set(keys + keys_inv)
-        squeeze = SqueezeTransform(squeeze_dim, keys_in=keys, keys_inv_in=keys_inv)
+        squeeze = SqueezeTransform(squeeze_dim, in_keys=keys, in_keys_inv=keys_inv)
         td = TensorDict(
             {
                 key: torch.randn(*batch, *size, nchannels, 16, 16, device=device)
@@ -687,7 +687,7 @@ def test_squeeze_inv(
     def test_grayscale(self, keys, device):
         torch.manual_seed(0)
         nchannels = 3
-        gs = GrayScale(keys_in=keys)
+        gs = GrayScale(in_keys=keys)
         dont_touch = torch.randn(1, nchannels, 16, 16, device=device)
         td = TensorDict(
             {key: torch.randn(1, nchannels, 16, 16, device=device) for key in keys},
@@ -720,7 +720,7 @@ def test_grayscale(self, keys, device):
     def test_totensorimage(self, keys, batch, device):
         torch.manual_seed(0)
         nchannels = 3
-        totensorimage = ToTensorImage(keys_in=keys)
+        totensorimage = ToTensorImage(in_keys=keys)
         dont_touch = torch.randn(*batch, nchannels, 16, 16, device=device)
         td = TensorDict(
             {
@@ -764,7 +764,7 @@ def test_totensorimage(self, keys, batch, device):
     @pytest.mark.parametrize("device", get_available_devices())
     def test_compose(self, keys, batch, device, nchannels=1, N=4):
         torch.manual_seed(0)
-        t1 = CatFrames(keys_in=keys, N=4)
+        t1 = CatFrames(in_keys=keys, N=4)
         t2 = FiniteTensorDictCheck()
         compose = Compose(t1, t2)
         dont_touch = torch.randn(*batch, nchannels, 16, 16, device=device)
@@ -818,8 +818,8 @@ def test_compose_inv(self, keys_inv_1, keys_inv_2, device):
         torch.manual_seed(0)
         keys_to_transform = set(keys_inv_1 + keys_inv_2)
         keys_total = set(["action_1", "action_2", "dont_touch"])
-        double2float_1 = DoubleToFloat(keys_inv_in=keys_inv_1)
-        double2float_2 = DoubleToFloat(keys_inv_in=keys_inv_2)
+        double2float_1 = DoubleToFloat(in_keys_inv=keys_inv_1)
+        double2float_2 = DoubleToFloat(in_keys_inv=keys_inv_2)
         compose = Compose(double2float_1, double2float_2)
         td = TensorDict(
             {
@@ -861,7 +861,7 @@ def test_observationnorm(
             loc = loc.to(device)
         if isinstance(scale, Tensor):
             scale = scale.to(device)
-        on = ObservationNorm(loc, scale, keys_in=keys, standard_normal=standard_normal)
+        on = ObservationNorm(loc, scale, in_keys=keys, standard_normal=standard_normal)
         dont_touch = torch.randn(1, nchannels, 16, 16, device=device)
         td = TensorDict(
             {key: torch.zeros(1, nchannels, 16, 16, device=device) for key in keys}, [1]
@@ -910,7 +910,7 @@ def test_catframes_transform_observation_spec(self):
         key1 = "first key"
         key2 = "second key"
         keys = [key1, key2]
-        cat_frames = CatFrames(N=N, keys_in=keys)
+        cat_frames = CatFrames(N=N, in_keys=keys)
         mins = [0, 0.5]
         maxes = [0.5, 1]
         observation_spec = CompositeSpec(
@@ -953,7 +953,7 @@ def test_catframes_buffer_check_latest_frame(self, device):
         key2_tensor = torch.ones(1, 1, 3, 3, device=device)
         key_tensors = [key1_tensor, key2_tensor]
         td = TensorDict(dict(zip(keys, key_tensors)), [1], device=device)
-        cat_frames = CatFrames(N=N, keys_in=keys)
+        cat_frames = CatFrames(N=N, in_keys=keys)
 
         cat_frames(td)
         latest_frame = td.get(key2)
@@ -973,7 +973,7 @@ def test_catframes_reset(self, device):
         key2_tensor = torch.ones(1, 1, 3, 3, device=device)
         key_tensors = [key1_tensor, key2_tensor]
         td = TensorDict(dict(zip(keys, key_tensors)), [1], device=device)
-        cat_frames = CatFrames(N=N, keys_in=keys)
+        cat_frames = CatFrames(N=N, in_keys=keys)
 
         cat_frames(td)
         buffer_length1 = len(cat_frames.buffer)
@@ -1014,7 +1014,7 @@ def test_finitetensordictcheck(self, device):
     def test_double2float(self, keys, keys_inv, device):
         torch.manual_seed(0)
         keys_total = set(keys + keys_inv)
-        double2float = DoubleToFloat(keys_in=keys, keys_inv_in=keys_inv)
+        double2float = DoubleToFloat(in_keys=keys, in_keys_inv=keys_inv)
         dont_touch = torch.randn(1, 3, 3, dtype=torch.double, device=device)
         td = TensorDict(
             {
@@ -1066,7 +1066,7 @@ def test_double2float(self, keys, keys_inv, device):
         ],
     )
     def test_cattensors(self, keys, device):
-        cattensors = CatTensors(keys_in=keys, out_key="observation_out", dim=-2)
+        cattensors = CatTensors(in_keys=keys, out_key="observation_out", dim=-2)
 
         dont_touch = torch.randn(1, 3, 3, dtype=torch.double, device=device)
         td = TensorDict(
@@ -1235,7 +1235,7 @@ def test_reward_scaling(self, batch, scale, loc, keys, device):
             keys_total = set([])
         else:
             keys_total = set(keys)
-        reward_scaling = RewardScaling(keys_in=keys, scale=scale, loc=loc)
+        reward_scaling = RewardScaling(in_keys=keys, scale=scale, loc=loc)
         td = TensorDict(
             {
                 **{key: torch.randn(*batch, 1, device=device) for key in keys_total},
@@ -1276,7 +1276,7 @@ def test_append(self):
         key = list(obs_spec.keys())[0]
 
         env = TransformedEnv(env)
-        env.append_transform(CatFrames(N=4, cat_dim=-1, keys_in=[key]))
+        env.append_transform(CatFrames(N=4, cat_dim=-1, in_keys=[key]))
         assert isinstance(env.transform, Compose)
         assert len(env.transform) == 1
         obs_spec = env.observation_spec
@@ -1301,7 +1301,7 @@ def test_insert(self):
         assert env._observation_spec is not None
         assert env._reward_spec is not None
 
-        env.insert_transform(0, CatFrames(N=4, cat_dim=-1, keys_in=[key]))
+        env.insert_transform(0, CatFrames(N=4, cat_dim=-1, in_keys=[key]))
 
         # transformed envs do not have spec after insert -- they need to be computed
         assert env._input_spec is None
@@ -1348,7 +1348,7 @@ def test_insert(self):
         assert env._observation_spec is None
         assert env._reward_spec is None
 
-        env.insert_transform(-5, CatFrames(N=4, cat_dim=-1, keys_in=[key]))
+        env.insert_transform(-5, CatFrames(N=4, cat_dim=-1, in_keys=[key]))
         assert isinstance(env.transform, Compose)
         assert len(env.transform) == 6
 
@@ -1411,7 +1411,7 @@ def test_r3m_instantiation(self, model, tensor_pixels_key, device):
         keys_out = ["next_vec"]
         r3m = R3MTransform(
             model,
-            keys_in=keys_in,
+            in_keys=keys_in,
             keys_out=keys_out,
             tensor_pixels_keys=tensor_pixels_key,
         )
@@ -1442,7 +1442,7 @@ def test_r3m_mult_images(self, model, device, stack_images, parallel):
         keys_out = ["next_vec"] if stack_images else ["next_vec", "next_vec2"]
         r3m = R3MTransform(
             model,
-            keys_in=keys_in,
+            in_keys=keys_in,
             keys_out=keys_out,
             stack_images=stack_images,
         )
@@ -1492,7 +1492,7 @@ def test_r3m_parallel(self, model, device):
         tensor_pixels_key = None
         r3m = R3MTransform(
             model,
-            keys_in=keys_in,
+            in_keys=keys_in,
             keys_out=keys_out,
             tensor_pixels_keys=tensor_pixels_key,
         )
@@ -1566,7 +1566,7 @@ def test_r3m_spec_against_real(self, model, tensor_pixels_key, device):
         keys_out = ["next_vec"]
         r3m = R3MTransform(
             model,
-            keys_in=keys_in,
+            in_keys=keys_in,
             keys_out=keys_out,
             tensor_pixels_keys=tensor_pixels_key,
         )
@@ -1592,7 +1592,7 @@ def test_vip_instantiation(self, model, tensor_pixels_key, device):
         keys_out = ["next_vec"]
         vip = VIPTransform(
             model,
-            keys_in=keys_in,
+            in_keys=keys_in,
             keys_out=keys_out,
             tensor_pixels_keys=tensor_pixels_key,
         )
@@ -1617,7 +1617,7 @@ def test_vip_mult_images(self, model, device, stack_images, parallel):
         keys_out = ["next_vec"] if stack_images else ["next_vec", "next_vec2"]
         vip = VIPTransform(
             model,
-            keys_in=keys_in,
+            in_keys=keys_in,
             keys_out=keys_out,
             stack_images=stack_images,
         )
@@ -1667,7 +1667,7 @@ def test_vip_parallel(self, model, device):
         tensor_pixels_key = None
         vip = VIPTransform(
             model,
-            keys_in=keys_in,
+            in_keys=keys_in,
             keys_out=keys_out,
             tensor_pixels_keys=tensor_pixels_key,
         )
@@ -1741,7 +1741,7 @@ def test_vip_spec_against_real(self, model, tensor_pixels_key, device):
         keys_out = ["next_vec"]
         vip = VIPTransform(
             model,
-            keys_in=keys_in,
+            in_keys=keys_in,
             keys_out=keys_out,
             tensor_pixels_keys=tensor_pixels_key,
         )
@@ -1762,7 +1762,7 @@ def test_batch_locked_transformed(device):
     env = TransformedEnv(
         MockBatchedLockedEnv(device),
         Compose(
-            ObservationNorm(keys_in=["next_observation"], loc=0.5, scale=1.1),
+            ObservationNorm(in_keys=["next_observation"], loc=0.5, scale=1.1),
             RewardClipping(0, 0.1),
         ),
     )
@@ -1786,7 +1786,7 @@ def test_batch_unlocked_transformed(device):
     env = TransformedEnv(
         MockBatchedUnLockedEnv(device),
         Compose(
-            ObservationNorm(keys_in=["next_observation"], loc=0.5, scale=1.1),
+            ObservationNorm(in_keys=["next_observation"], loc=0.5, scale=1.1),
             RewardClipping(0, 0.1),
         ),
     )
@@ -1806,7 +1806,7 @@ def test_batch_unlocked_with_batch_size_transformed(device):
     env = TransformedEnv(
         MockBatchedUnLockedEnv(device, batch_size=torch.Size([2])),
         Compose(
-            ObservationNorm(keys_in=["next_observation"], loc=0.5, scale=1.1),
+            ObservationNorm(in_keys=["next_observation"], loc=0.5, scale=1.1),
             RewardClipping(0, 0.1),
         ),
     )