Merge pull request #1 from ARISE-Initiative/master

fix scan modality (ARISE-Initiative#119)

robomimic/models/base_nets.py

@@ -851,6 +851,11 @@ def __init__(
 
         # Get activation requested
         activation = CONV_ACTIVATIONS[activation]
+
+        # Add layer kwargs
+        conv_kwargs["out_channels"] = out_channels
+        conv_kwargs["kernel_size"] = kernel_size
+        conv_kwargs["stride"] = stride
 
         # Generate network
         self.n_layers = len(out_channels)

robomimic/models/obs_core.py

@@ -231,8 +231,9 @@ def __init__(
             conv_kwargs = dict()
 
         # Generate backbone network
+        # N input channels is assumed to be the first dimension
         self.backbone = BaseNets.Conv1dBase(
-            input_channel=1,
+            input_channel=self.input_shape[0],
             activation=conv_activation,
             **conv_kwargs,
         )

robomimic/utils/obs_utils.py

@@ -944,10 +944,34 @@ class ScanModality(Modality):
 
     @classmethod
     def _default_obs_processor(cls, obs):
+        # Channel swaps ([...,] L, C) --> ([...,] C, L)
+
+        # First, add extra dimension at 2nd to last index to treat this as a frame
+        shape = obs.shape
+        new_shape = [*shape[:-2], 1, *shape[-2:]]
+        obs = obs.reshape(new_shape)
+
+        # Convert shape
+        obs = batch_image_hwc_to_chw(obs)
+
+        # Remove extra dimension (it's the second from last dimension)
+        obs = obs.squeeze(-2)
         return obs
 
     @classmethod
     def _default_obs_unprocessor(cls, obs):
+        # Channel swaps ([B,] C, L) --> ([B,] L, C)
+
+        # First, add extra dimension at 1st index to treat this as a frame
+        shape = obs.shape
+        new_shape = [*shape[:-2], 1, *shape[-2:]]
+        obs = obs.reshape(new_shape)
+
+        # Convert shape
+        obs = batch_image_chw_to_hwc(obs)
+
+        # Remove extra dimension (it's the second from last dimension)
+        obs = obs.squeeze(-2)
         return obs