bug fix for registering backbone architecture

Author: Danfei Xu

robomimic/models/base_nets.py

@@ -476,7 +476,7 @@ def __init_subclass__(cls, **kwargs):
         This also future-proofs us for any additional encoder classes we would
         like to add ourselves.
         """
-        ObsUtils.register_encoder_core(cls)
+        ObsUtils.register_encoder_backbone(cls)
 
     # dirty hack - re-implement to pass the buck onto subclasses from ABC parent
     def output_shape(self, input_shape):

robomimic/models/obs_core.py

@@ -99,7 +99,7 @@ def __init__(
 
         # extract only relevant kwargs for this specific backbone
         backbone_kwargs = extract_class_init_kwargs_from_dict(
-                cls = ObsUtils.OBS_ENCODER_CORES[backbone_class],
+                cls = ObsUtils.OBS_ENCODER_BACKBONES[backbone_class],
                 dic=backbone_kwargs, copy=True)
 
         # visual backbone

robomimic/utils/obs_utils.py

@@ -41,8 +41,9 @@
 # in their config, without having to manually register their class internally.
 # This also future-proofs us for any additional encoder / randomizer classes we would
 # like to add ourselves.
-OBS_ENCODER_CORES = {"None": None}          # Include default None
-OBS_RANDOMIZERS = {"None": None}            # Include default None
+OBS_ENCODER_CORES = {"None": None}          # Per-modality core net as defined in obs_cores.py, e.g., "VisualCore"
+OBS_RANDOMIZERS = {"None": None}            # Obs randomizer defined in obs_cores.py, e.g., "CropRandomizer"
+OBS_ENCODER_BACKBONES = {"None": None}      # Architecture backbones for encoding obervation, e.g., "ResNet18Conv"
 
 
 def register_obs_key(target_class):
@@ -59,6 +60,10 @@ def register_randomizer(target_class):
     assert target_class not in OBS_RANDOMIZERS, f"Already registered obs randomizer {target_class}!"
     OBS_RANDOMIZERS[target_class.__name__] = target_class
 
+def register_encoder_backbone(target_class):
+    assert target_class not in OBS_ENCODER_BACKBONES, f"Already registered obs encoder backbone {target_class}!"
+    OBS_ENCODER_BACKBONES[target_class.__name__] = target_class
+
 
 class ObservationKeyToModalityDict(dict):
     """