[Naming] Fixing key names (#668)

* fixing key names

* fixing key names

Author: vmoens

test/test_transforms.py

@@ -1407,12 +1407,12 @@ def test_insert(self):
 class TestR3M:
     @pytest.mark.parametrize("tensor_pixels_key", [None, ["funny_key"]])
     def test_r3m_instantiation(self, model, tensor_pixels_key, device):
-        keys_in = ["next_pixels"]
-        keys_out = ["next_vec"]
+        in_keys = ["next_pixels"]
+        out_keys = ["next_vec"]
         r3m = R3MTransform(
             model,
-            in_keys=keys_in,
-            keys_out=keys_out,
+            in_keys=in_keys,
+            out_keys=out_keys,
             tensor_pixels_keys=tensor_pixels_key,
         )
         base_env = DiscreteActionConvMockEnvNumpy().to(device)
@@ -1438,12 +1438,12 @@ def test_r3m_instantiation(self, model, tensor_pixels_key, device):
         ],
     )
     def test_r3m_mult_images(self, model, device, stack_images, parallel):
-        keys_in = ["next_pixels", "next_pixels2"]
-        keys_out = ["next_vec"] if stack_images else ["next_vec", "next_vec2"]
+        in_keys = ["next_pixels", "next_pixels2"]
+        out_keys = ["next_vec"] if stack_images else ["next_vec", "next_vec2"]
         r3m = R3MTransform(
             model,
-            in_keys=keys_in,
-            keys_out=keys_out,
+            in_keys=in_keys,
+            out_keys=out_keys,
             stack_images=stack_images,
         )
 
@@ -1487,13 +1487,13 @@ def base_env_constructor():
         transformed_env.close()
 
     def test_r3m_parallel(self, model, device):
-        keys_in = ["next_pixels"]
-        keys_out = ["next_vec"]
+        in_keys = ["next_pixels"]
+        out_keys = ["next_vec"]
         tensor_pixels_key = None
         r3m = R3MTransform(
             model,
-            in_keys=keys_in,
-            keys_out=keys_out,
+            in_keys=in_keys,
+            out_keys=out_keys,
             tensor_pixels_keys=tensor_pixels_key,
         )
         base_env = ParallelEnv(4, lambda: DiscreteActionConvMockEnvNumpy().to(device))
@@ -1562,12 +1562,12 @@ def test_r3mnet_transform_observation_spec(
 
     @pytest.mark.parametrize("tensor_pixels_key", [None, ["funny_key"]])
     def test_r3m_spec_against_real(self, model, tensor_pixels_key, device):
-        keys_in = ["next_pixels"]
-        keys_out = ["next_vec"]
+        in_keys = ["next_pixels"]
+        out_keys = ["next_vec"]
         r3m = R3MTransform(
             model,
-            in_keys=keys_in,
-            keys_out=keys_out,
+            in_keys=in_keys,
+            out_keys=out_keys,
             tensor_pixels_keys=tensor_pixels_key,
         )
         base_env = DiscreteActionConvMockEnvNumpy().to(device)
@@ -1588,12 +1588,12 @@ def test_r3m_spec_against_real(self, model, tensor_pixels_key, device):
 class TestVIP:
     @pytest.mark.parametrize("tensor_pixels_key", [None, ["funny_key"]])
     def test_vip_instantiation(self, model, tensor_pixels_key, device):
-        keys_in = ["next_pixels"]
-        keys_out = ["next_vec"]
+        in_keys = ["next_pixels"]
+        out_keys = ["next_vec"]
         vip = VIPTransform(
             model,
-            in_keys=keys_in,
-            keys_out=keys_out,
+            in_keys=in_keys,
+            out_keys=out_keys,
             tensor_pixels_keys=tensor_pixels_key,
         )
         base_env = DiscreteActionConvMockEnvNumpy().to(device)
@@ -1613,12 +1613,12 @@ def test_vip_instantiation(self, model, tensor_pixels_key, device):
     @pytest.mark.parametrize("stack_images", [True, False])
     @pytest.mark.parametrize("parallel", [True, False])
     def test_vip_mult_images(self, model, device, stack_images, parallel):
-        keys_in = ["next_pixels", "next_pixels2"]
-        keys_out = ["next_vec"] if stack_images else ["next_vec", "next_vec2"]
+        in_keys = ["next_pixels", "next_pixels2"]
+        out_keys = ["next_vec"] if stack_images else ["next_vec", "next_vec2"]
         vip = VIPTransform(
             model,
-            in_keys=keys_in,
-            keys_out=keys_out,
+            in_keys=in_keys,
+            out_keys=out_keys,
             stack_images=stack_images,
         )
 
@@ -1662,13 +1662,13 @@ def base_env_constructor():
         transformed_env.close()
 
     def test_vip_parallel(self, model, device):
-        keys_in = ["next_pixels"]
-        keys_out = ["next_vec"]
+        in_keys = ["next_pixels"]
+        out_keys = ["next_vec"]
         tensor_pixels_key = None
         vip = VIPTransform(
             model,
-            in_keys=keys_in,
-            keys_out=keys_out,
+            in_keys=in_keys,
+            out_keys=out_keys,
             tensor_pixels_keys=tensor_pixels_key,
         )
         base_env = ParallelEnv(4, lambda: DiscreteActionConvMockEnvNumpy().to(device))
@@ -1688,13 +1688,13 @@ def test_vip_parallel(self, model, device):
         del transformed_env
 
     def test_vip_parallel_reward(self, model, device):
-        keys_in = ["next_pixels"]
-        keys_out = ["next_vec"]
+        in_keys = ["next_pixels"]
+        out_keys = ["next_vec"]
         tensor_pixels_key = None
         vip = VIPRewardTransform(
             model,
-            keys_in=keys_in,
-            keys_out=keys_out,
+            in_keys=in_keys,
+            out_keys=out_keys,
             tensor_pixels_keys=tensor_pixels_key,
         )
         base_env = ParallelEnv(4, lambda: DiscreteActionConvMockEnvNumpy().to(device))
@@ -1802,12 +1802,12 @@ def test_vipnet_transform_observation_spec(
 
     @pytest.mark.parametrize("tensor_pixels_key", [None, ["funny_key"]])
     def test_vip_spec_against_real(self, model, tensor_pixels_key, device):
-        keys_in = ["next_pixels"]
-        keys_out = ["next_vec"]
+        in_keys = ["next_pixels"]
+        out_keys = ["next_vec"]
         vip = VIPTransform(
             model,
-            in_keys=keys_in,
-            keys_out=keys_out,
+            in_keys=in_keys,
+            out_keys=out_keys,
             tensor_pixels_keys=tensor_pixels_key,
         )
         base_env = DiscreteActionConvMockEnvNumpy().to(device)

torchrl/envs/transforms/r3m.py

@@ -61,7 +61,7 @@ def __init__(self, in_keys, out_keys, model_name, del_keys: bool = True):
                 f"model {model_name} is currently not supported by R3M"
             )
         convnet.fc = Identity()
-        super().__init__(in_keys=in_keys, keys_out=out_keys)
+        super().__init__(in_keys=in_keys, out_keys=out_keys)
         self.convnet = convnet
         self.del_keys = del_keys
 
@@ -92,11 +92,11 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
 
         observation_spec = CompositeSpec(**observation_spec)
         if self.del_keys:
-            for key_in in keys:
-                del observation_spec[key_in]
+            for in_key in keys:
+                del observation_spec[in_key]
 
-        for key_out in self.keys_out:
-            observation_spec[key_out] = NdUnboundedContinuousTensorSpec(
+        for out_key in self.out_keys:
+            observation_spec[out_key] = NdUnboundedContinuousTensorSpec(
                 shape=torch.Size([self.outdim]), device=device
             )
 
@@ -161,7 +161,7 @@ class R3MTransform(Compose):
         model_name (str): one of resnet50, resnet34 or resnet18
         in_keys (list of str, optional): list of input keys. If left empty, the
             "next_pixels" key is assumed.
-        keys_out (list of str, optional): list of output keys. If left empty,
+        out_keys (list of str, optional): list of output keys. If left empty,
              "next_r3m_vec" is assumed.
         size (int, optional): Size of the image to feed to resnet.
             Defaults to 244.
@@ -187,105 +187,105 @@ def __init__(
         self,
         model_name: str,
         in_keys: List[str] = None,
-        keys_out: List[str] = None,
+        out_keys: List[str] = None,
         size: int = 244,
         stack_images: bool = True,
         download: bool = False,
         download_path: Optional[str] = None,
         tensor_pixels_keys: List[str] = None,
     ):
         super().__init__()
-        self.keys_in = in_keys
+        self.in_keys = in_keys
         self.download = download
         self.download_path = download_path
         self.model_name = model_name
-        self.keys_out = keys_out
+        self.out_keys = out_keys
         self.size = size
         self.stack_images = stack_images
         self.tensor_pixels_keys = tensor_pixels_keys
 
     def _init(self):
-        keys_in = self.keys_in
+        in_keys = self.in_keys
         model_name = self.model_name
-        keys_out = self.keys_out
+        out_keys = self.out_keys
         size = self.size
         stack_images = self.stack_images
         tensor_pixels_keys = self.tensor_pixels_keys
 
         # ToTensor
         transforms = []
         if tensor_pixels_keys:
-            for i in range(len(keys_in)):
+            for i in range(len(in_keys)):
                 transforms.append(
                     CatTensors(
-                        in_keys=[keys_in[i]],
+                        in_keys=[in_keys[i]],
                         out_key=tensor_pixels_keys[i],
                         del_keys=False,
                     )
                 )
 
         totensor = ToTensorImage(
             unsqueeze=False,
-            in_keys=keys_in,
+            in_keys=in_keys,
         )
         transforms.append(totensor)
 
         # Normalize
         mean = [0.485, 0.456, 0.406]
         std = [0.229, 0.224, 0.225]
         normalize = ObservationNorm(
-            in_keys=keys_in,
+            in_keys=in_keys,
             loc=torch.tensor(mean).view(3, 1, 1),
             scale=torch.tensor(std).view(3, 1, 1),
             standard_normal=True,
         )
         transforms.append(normalize)
 
         # Resize: note that resize is a no-op if the tensor has the desired size already
-        resize = Resize(size, size, in_keys=keys_in)
+        resize = Resize(size, size, in_keys=in_keys)
         transforms.append(resize)
 
         # R3M
-        if keys_out is None:
+        if out_keys is None:
             if stack_images:
-                keys_out = ["next_r3m_vec"]
+                out_keys = ["next_r3m_vec"]
             else:
-                keys_out = [f"next_r3m_vec_{i}" for i in range(len(keys_in))]
-        elif stack_images and len(keys_out) != 1:
+                out_keys = [f"next_r3m_vec_{i}" for i in range(len(in_keys))]
+        elif stack_images and len(out_keys) != 1:
             raise ValueError(
-                f"key_out must be of length 1 if stack_images is True. Got keys_out={keys_out}"
+                f"out_key must be of length 1 if stack_images is True. Got out_keys={out_keys}"
             )
-        elif not stack_images and len(keys_out) != len(keys_in):
+        elif not stack_images and len(out_keys) != len(in_keys):
             raise ValueError(
-                "key_out must be of length equal to in_keys if stack_images is False."
+                "out_key must be of length equal to in_keys if stack_images is False."
             )
 
-        if stack_images and len(keys_in) > 1:
+        if stack_images and len(in_keys) > 1:
             if self.is_3d:
                 unsqueeze = UnsqueezeTransform(
-                    in_keys=keys_in,
-                    keys_out=keys_in,
+                    in_keys=in_keys,
+                    out_keys=in_keys,
                     unsqueeze_dim=-4,
                 )
                 transforms.append(unsqueeze)
 
             cattensors = CatTensors(
-                keys_in,
-                keys_out[0],
+                in_keys,
+                out_keys[0],
                 dim=-4,
             )
             network = _R3MNet(
-                in_keys=keys_out,
-                out_keys=keys_out,
+                in_keys=out_keys,
+                out_keys=out_keys,
                 model_name=model_name,
                 del_keys=False,
             )
-            flatten = FlattenObservation(-2, -1, keys_out)
+            flatten = FlattenObservation(-2, -1, out_keys)
             transforms = [*transforms, cattensors, network, flatten]
         else:
             network = _R3MNet(
-                in_keys=keys_in,
-                out_keys=keys_out,
+                in_keys=in_keys,
+                out_keys=out_keys,
                 model_name=model_name,
                 del_keys=True,
             )