revert dp changes, make act and tdmpc batch friendly

lerobot/common/policies/abstract.py

@@ -1,15 +1,20 @@
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from collections import deque
 
 import torch
 from torch import Tensor, nn
 
 
-class AbstractPolicy(nn.Module):
+class AbstractPolicy(nn.Module, ABC):
+    """Base policy which all policies should be derived from.
+
+    The forward method should generally not be overriden as it plays the role of handling multi-step policies. See its
+    documentation for more information.
+    """
+
     @abstractmethod
     def update(self, replay_buffer, step):
         """One step of the policy's learning algorithm."""
-        pass
 
     def save(self, fp):
         torch.save(self.state_dict(), fp)
@@ -24,7 +29,6 @@ class AbstractPolicy(nn.Module):
 
         Should return a (batch_size, n_action_steps, *) tensor of actions.
         """
-        pass
 
     def forward(self, *args, **kwargs):
         """Inference step that makes multi-step policies compatible with their single-step environments.

lerobot/common/policies/act/policy.py

@@ -153,10 +153,6 @@ class ActionChunkingTransformerPolicy(AbstractPolicy):
 
         self.eval()
 
-        # TODO(rcadene): remove unsqueeze hack to add bsize=1
-        observation["image", "top"] = observation["image", "top"].unsqueeze(0)
-        # observation["state"] = observation["state"].unsqueeze(0)
-
         # TODO(rcadene): remove hack
         # add 1 camera dimension
         observation["image", "top"] = observation["image", "top"].unsqueeze(1)
@@ -180,11 +176,8 @@ class ActionChunkingTransformerPolicy(AbstractPolicy):
             # exp_weights = torch.from_numpy(exp_weights).cuda().unsqueeze(dim=1)
             # raw_action = (actions_for_curr_step * exp_weights).sum(dim=0, keepdim=True)
 
-        # remove bsize=1
-        action = action.squeeze(0)
-
         # take first predicted action or n first actions
-        action = action[0] if self.n_action_steps == 1 else action[: self.n_action_steps]
+        action = action[: self.n_action_steps]
         return action
 
     def _forward(self, qpos, image, actions=None, is_pad=None):

lerobot/common/policies/diffusion/model/multi_image_obs_encoder.py

@@ -1,37 +1,15 @@
 import copy
-from typing import Dict, Optional, Tuple, Union
+from typing import Dict, Tuple, Union
 
-import timm
 import torch
 import torch.nn as nn
 import torchvision
-from robomimic.models.base_nets import SpatialSoftmax
 
 from lerobot.common.policies.diffusion.model.crop_randomizer import CropRandomizer
 from lerobot.common.policies.diffusion.model.module_attr_mixin import ModuleAttrMixin
 from lerobot.common.policies.diffusion.pytorch_utils import replace_submodules
 
 
-class RgbEncoder(nn.Module):
-    """Following `VisualCore` from Robomimic 0.2.0."""
-
-    def __init__(self, input_shape, model_name="resnet18", pretrained=False, num_keypoints=32):
-        """
-        resnet_name: a timm model name.
-        pretrained: whether to use timm pretrained weights.
-        num_keypoints: Number of keypoints for SpatialSoftmax (default value of 32 matches PushT Image).
-        """
-        super().__init__()
-        self.backbone = timm.create_model(model_name, pretrained, num_classes=0, global_pool="")
-        # Figure out the feature map shape.
-        with torch.inference_mode():
-            feat_map_shape = tuple(self.backbone(torch.zeros(size=(1, *input_shape))).shape[1:])
-        self.pool = SpatialSoftmax(feat_map_shape, num_kp=num_keypoints)
-
-    def forward(self, x):
-        return torch.flatten(self.pool(self.backbone(x)), start_dim=1)
-
-
 class MultiImageObsEncoder(ModuleAttrMixin):
     def __init__(
         self,
@@ -46,7 +24,7 @@ class MultiImageObsEncoder(ModuleAttrMixin):
         share_rgb_model: bool = False,
         # renormalize rgb input with imagenet normalization
         # assuming input in [0,1]
-        norm_mean_std: Optional[tuple[float, float]] = None,
+        imagenet_norm: bool = False,
     ):
         """
         Assumes rgb input: B,C,H,W
@@ -120,9 +98,10 @@ class MultiImageObsEncoder(ModuleAttrMixin):
                         this_normalizer = torchvision.transforms.CenterCrop(size=(h, w))
                 # configure normalizer
                 this_normalizer = nn.Identity()
-                if norm_mean_std is not None:
+                if imagenet_norm:
+                    # TODO(rcadene): move normalizer to dataset and env
                     this_normalizer = torchvision.transforms.Normalize(
-                        mean=norm_mean_std[0], std=norm_mean_std[1]
+                        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
                     )
 
                 this_transform = nn.Sequential(this_resizer, this_randomizer, this_normalizer)

lerobot/common/policies/diffusion/policy.py

@@ -7,7 +7,7 @@ import torch
 from lerobot.common.policies.abstract import AbstractPolicy
 from lerobot.common.policies.diffusion.diffusion_unet_image_policy import DiffusionUnetImagePolicy
 from lerobot.common.policies.diffusion.model.lr_scheduler import get_scheduler
-from lerobot.common.policies.diffusion.model.multi_image_obs_encoder import MultiImageObsEncoder, RgbEncoder
+from lerobot.common.policies.diffusion.model.multi_image_obs_encoder import MultiImageObsEncoder
 
 
 class DiffusionPolicy(AbstractPolicy):
@@ -38,7 +38,7 @@ class DiffusionPolicy(AbstractPolicy):
         self.cfg = cfg
 
         noise_scheduler = hydra.utils.instantiate(cfg_noise_scheduler)
-        rgb_model = RgbEncoder(input_shape=shape_meta.obs.image.shape, **cfg_rgb_model)
+        rgb_model = hydra.utils.instantiate(cfg_rgb_model)
         obs_encoder = MultiImageObsEncoder(
             rgb_model=rgb_model,
             **cfg_obs_encoder,

lerobot/common/policies/tdmpc/policy.py

@@ -128,11 +128,6 @@ class TDMPC(AbstractPolicy):
     def select_action(self, observation, step_count):
         t0 = step_count.item() == 0
 
-        # TODO(rcadene): remove unsqueeze hack...
-        if observation["image"].ndim == 3:
-            observation["image"] = observation["image"].unsqueeze(0)
-            observation["state"] = observation["state"].unsqueeze(0)
-
         obs = {
             # TODO(rcadene): remove contiguous hack...
             "rgb": observation["image"].contiguous(),
@@ -149,7 +144,7 @@ class TDMPC(AbstractPolicy):
         if self.cfg.mpc:
             a = self.plan(z, t0=t0, step=step)
         else:
-            a = self.model.pi(z, self.cfg.min_std * self.model.training).squeeze(0)
+            a = self.model.pi(z, self.cfg.min_std * self.model.training)
         return a
 
     @torch.no_grad()