Style & Formatting

.pre-commit-config.yaml

@@ -0,0 +1,33 @@
+exclude: ^(data/|tests/|diffusion_policy/)
+default_language_version:
+    python: python3.10
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.0
+    hooks:
+      - id: check-added-large-files
+      - id: debug-statements
+      - id: check-merge-conflict
+      - id: check-case-conflict
+      - id: check-yaml
+      - id: check-toml
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+  - repo: https://github.com/asottile/pyupgrade
+    rev: v3.15.1
+    hooks:
+    -   id: pyupgrade
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.2.2
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format
+  - repo: https://github.com/python-poetry/poetry
+    rev: 1.8.0
+    hooks:
+      - id: poetry-check
+      - id: poetry-lock
+        args:
+          - "--check"
+          - "--no-update"

README.md

@@ -10,7 +10,7 @@ conda activate lerobot
 
 [Install `poetry`](https://python-poetry.org/docs/#installation) (if you don't have it already)
 ```
-curl -sSL https://install.python-poetry.org | python3 -
+curl -sSL https://install.python-poetry.org | python -
 ```
 
 Install dependencies
@@ -26,6 +26,7 @@ export TMPDIR='~/tmp'
 
 Install `diffusion_policy` #HACK
 ```
+# from this directory
 git clone https://github.com/real-stanford/diffusion_policy
 cp -r diffusion_policy/diffusion_policy $(poetry env info -p)/lib/python3.10/site-packages/
 ```
@@ -107,11 +108,10 @@ eval_episodes=7
 
 **Style**
 ```
-isort lerobot && isort tests && black lerobot && black tests
-pylint lerobot && pylint tests  # not enforce for now
+pre-commit install
 ```
 
 **Tests**
 ```
 pytest -sx tests
-```
+```

lerobot/common/datasets/factory.py

@@ -70,6 +70,7 @@ def make_offline_buffer(cfg, sampler=None):
         offline_buffer = PushtExperienceReplay(
             "pusht",
             # download="force",
+            # TODO(aliberts): automate download
             download=False,
             streaming=False,
             root="data",

lerobot/common/datasets/pusht.py

@@ -1,7 +1,6 @@
 import os
-import pickle
 from pathlib import Path
-from typing import Any, Callable, Dict, Tuple
+from typing import Callable
 
 import einops
 import numpy as np
@@ -10,25 +9,25 @@
 import torch
 import torchrl
 import tqdm
-from diffusion_policy.common.replay_buffer import ReplayBuffer
-from diffusion_policy.env.pusht.pusht_env import pymunk_to_shapely
 from tensordict import TensorDict
 from torchrl.data.datasets.utils import _get_root_dir
 from torchrl.data.replay_buffers.replay_buffers import (
-    TensorDictPrioritizedReplayBuffer,
     TensorDictReplayBuffer,
 )
 from torchrl.data.replay_buffers.samplers import (
     Sampler,
-    SliceSampler,
-    SliceSamplerWithoutReplacement,
 )
 from torchrl.data.replay_buffers.storages import TensorStorage, _collate_id
 from torchrl.data.replay_buffers.writers import ImmutableDatasetWriter, Writer
 
+from diffusion_policy.common.replay_buffer import ReplayBuffer
+from diffusion_policy.env.pusht.pusht_env import pymunk_to_shapely
+
 # as define in env
 SUCCESS_THRESHOLD = 0.95  # 95% coverage,
 
+DEFAULT_TEE_MASK = pymunk.ShapeFilter.ALL_MASKS()
+
 
 def get_goal_pose_body(pose):
     mass = 1
@@ -53,7 +52,7 @@ def add_tee(
     angle,
     scale=30,
     color="LightSlateGray",
-    mask=pymunk.ShapeFilter.ALL_MASKS(),
+    mask=DEFAULT_TEE_MASK,
 ):
     mass = 1
     length = 4
@@ -87,7 +86,6 @@ def add_tee(
 
 
 class PushtExperienceReplay(TensorDictReplayBuffer):
-
     def __init__(
         self,
         dataset_id,
@@ -127,7 +125,7 @@ def __init__(
         if split_trajs:
             raise NotImplementedError
 
-        if self.download == True:
+        if self.download:
             raise NotImplementedError()
 
         if root is None:
@@ -193,18 +191,18 @@ def _download_and_preproc(self):
         # TODO(rcadene)
 
         # load
+        # TODO(aliberts): Dynamic paths
         zarr_path = (
             "/home/rcadene/code/diffusion_policy/data/pusht/pusht_cchi_v7_replay.zarr"
+            # "/home/simon/build/diffusion_policy/data/pusht/pusht_cchi_v7_replay.zarr"
         )
-        dataset_dict = ReplayBuffer.copy_from_path(
-            zarr_path
-        )  # , keys=['img', 'state', 'action'])
+        dataset_dict = ReplayBuffer.copy_from_path(zarr_path)  # , keys=['img', 'state', 'action'])
 
         episode_ids = dataset_dict.get_episode_idxs()
         num_episodes = dataset_dict.meta["episode_ends"].shape[0]
         total_frames = dataset_dict["action"].shape[0]
         assert len(
-            set([dataset_dict[key].shape[0] for key in dataset_dict.keys()])
+            {dataset_dict[key].shape[0] for key in dataset_dict}
         ), "Some data type dont have the same number of total frames."
 
         # TODO: verify that goal pose is expected to be fixed
@@ -245,9 +243,7 @@ def _download_and_preproc(self):
                 ]
                 space.add(*walls)
 
-                block_body = add_tee(
-                    space, block_pos[i].tolist(), block_angle[i].item()
-                )
+                block_body = add_tee(space, block_pos[i].tolist(), block_angle[i].item())
                 goal_geom = pymunk_to_shapely(goal_body, block_body.shapes)
                 block_geom = pymunk_to_shapely(block_body, block_body.shapes)
                 intersection_area = goal_geom.intersection(block_geom).area
@@ -278,11 +274,7 @@ def _download_and_preproc(self):
 
             if episode_id == 0:
                 # hack to initialize tensordict data structure to store episodes
-                td_data = (
-                    episode[0]
-                    .expand(total_frames)
-                    .memmap_like(self.root / self.dataset_id)
-                )
+                td_data = episode[0].expand(total_frames).memmap_like(self.root / self.dataset_id)
 
             td_data[idxtd : idxtd + len(episode)] = episode
 

lerobot/common/datasets/simxarm.py

@@ -1,15 +1,14 @@
 import os
 import pickle
 from pathlib import Path
-from typing import Any, Callable, Dict, Tuple
+from typing import Callable
 
 import torch
 import torchrl
 import tqdm
 from tensordict import TensorDict
 from torchrl.data.datasets.utils import _get_root_dir
 from torchrl.data.replay_buffers.replay_buffers import (
-    TensorDictPrioritizedReplayBuffer,
     TensorDictReplayBuffer,
 )
 from torchrl.data.replay_buffers.samplers import (
@@ -22,7 +21,6 @@
 
 
 class SimxarmExperienceReplay(TensorDictReplayBuffer):
-
     available_datasets = [
         "xarm_lift_medium",
     ]
@@ -77,15 +75,11 @@ def __init__(
 
         if num_slices is not None or slice_len is not None:
             if sampler is not None:
-                raise ValueError(
-                    "`num_slices` and `slice_len` are exclusive with the `sampler` argument."
-                )
+                raise ValueError("`num_slices` and `slice_len` are exclusive with the `sampler` argument.")
 
             if replacement:
                 if not self.shuffle:
-                    raise RuntimeError(
-                        "shuffle=False can only be used when replacement=False."
-                    )
+                    raise RuntimeError("shuffle=False can only be used when replacement=False.")
                 sampler = SliceSampler(
                     num_slices=num_slices,
                     slice_len=slice_len,
@@ -130,7 +124,7 @@ def _download_and_preproc(self):
 
         # load
         dataset_dir = Path("data") / self.dataset_id
-        dataset_path = dataset_dir / f"buffer.pkl"
+        dataset_path = dataset_dir / "buffer.pkl"
         print(f"Using offline dataset '{dataset_path}'")
         with open(dataset_path, "rb") as f:
             dataset_dict = pickle.load(f)
@@ -150,12 +144,8 @@ def _download_and_preproc(self):
 
             image = torch.tensor(dataset_dict["observations"]["rgb"][idx0:idx1])
             state = torch.tensor(dataset_dict["observations"]["state"][idx0:idx1])
-            next_image = torch.tensor(
-                dataset_dict["next_observations"]["rgb"][idx0:idx1]
-            )
-            next_state = torch.tensor(
-                dataset_dict["next_observations"]["state"][idx0:idx1]
-            )
+            next_image = torch.tensor(dataset_dict["next_observations"]["rgb"][idx0:idx1])
+            next_state = torch.tensor(dataset_dict["next_observations"]["state"][idx0:idx1])
             next_reward = torch.tensor(dataset_dict["rewards"][idx0:idx1])
             next_done = torch.tensor(dataset_dict["dones"][idx0:idx1])
 
@@ -176,11 +166,7 @@ def _download_and_preproc(self):
 
             if episode_id == 0:
                 # hack to initialize tensordict data structure to store episodes
-                td_data = (
-                    episode[0]
-                    .expand(total_frames)
-                    .memmap_like(self.root / self.dataset_id)
-                )
+                td_data = episode[0].expand(total_frames).memmap_like(self.root / self.dataset_id)
 
             td_data[idx0:idx1] = episode
 

lerobot/common/envs/pusht.py

@@ -1,7 +1,6 @@
 import importlib
 from typing import Optional
 
-import numpy as np
 import torch
 from tensordict import TensorDict
 from torchrl.data.tensor_specs import (
@@ -20,7 +19,6 @@
 
 
 class PushtEnv(EnvBase):
-
     def __init__(
         self,
         frame_skip: int = 1,
@@ -46,7 +44,8 @@ def __init__(
         if not _has_gym:
             raise ImportError("Cannot import gym.")
 
-        from diffusion_policy.env.pusht.pusht_env import PushTEnv
+        # TODO(rcadene) (PushTEnv is similar to PushTImageEnv, but without the image rendering, it's faster to iterate on)
+        # from diffusion_policy.env.pusht.pusht_env import PushTEnv
 
         if not from_pixels:
             raise NotImplementedError("Use PushTEnv, instead of PushTImageEnv")
@@ -71,14 +70,10 @@ def _format_raw_obs(self, raw_obs):
             obs = {"image": torch.from_numpy(raw_obs["image"])}
 
             if not self.pixels_only:
-                obs["state"] = torch.from_numpy(raw_obs["agent_pos"]).type(
-                    torch.float32
-                )
+                obs["state"] = torch.from_numpy(raw_obs["agent_pos"]).type(torch.float32)
         else:
             # TODO:
-            obs = {
-                "state": torch.from_numpy(raw_obs["observation"]).type(torch.float32)
-            }
+            obs = {"state": torch.from_numpy(raw_obs["observation"]).type(torch.float32)}
 
         obs = TensorDict(obs, batch_size=[])
         return obs
@@ -109,7 +104,7 @@ def _step(self, tensordict: TensorDict):
         # step expects shape=(4,) so we pad if necessary
         # TODO(rcadene): add info["is_success"] and info["success"] ?
         sum_reward = 0
-        for t in range(self.frame_skip):
+        for _ in range(self.frame_skip):
             raw_obs, reward, done, info = self._env.step(action)
             sum_reward += reward
 

lerobot/common/envs/simxarm.py

@@ -15,12 +15,13 @@
 
 from lerobot.common.utils import set_seed
 
+MAX_NUM_ACTIONS = 4
+
 _has_gym = importlib.util.find_spec("gym") is not None
 _has_simxarm = importlib.util.find_spec("simxarm") is not None and _has_gym
 
 
 class SimxarmEnv(EnvBase):
-
     def __init__(
         self,
         task,
@@ -52,18 +53,13 @@ def __init__(
         from simxarm import TASKS
 
         if self.task not in TASKS:
-            raise ValueError(
-                f"Unknown task {self.task}. Must be one of {list(TASKS.keys())}"
-            )
+            raise ValueError(f"Unknown task {self.task}. Must be one of {list(TASKS.keys())}")
 
         self._env = TASKS[self.task]["env"]()
 
-        MAX_NUM_ACTIONS = 4
         num_actions = len(TASKS[self.task]["action_space"])
         self._action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(num_actions,))
-        self._action_padding = np.zeros(
-            (MAX_NUM_ACTIONS - num_actions), dtype=np.float32
-        )
+        self._action_padding = np.zeros((MAX_NUM_ACTIONS - num_actions), dtype=np.float32)
         if "w" not in TASKS[self.task]["action_space"]:
             self._action_padding[-1] = 1.0
 
@@ -75,9 +71,7 @@ def render(self, mode="rgb_array", width=384, height=384):
 
     def _format_raw_obs(self, raw_obs):
         if self.from_pixels:
-            image = self.render(
-                mode="rgb_array", width=self.image_size, height=self.image_size
-            )
+            image = self.render(mode="rgb_array", width=self.image_size, height=self.image_size)
             image = image.transpose(2, 0, 1)  # (H, W, C) -> (C, H, W)
             image = torch.tensor(image.copy(), dtype=torch.uint8)
 
@@ -114,7 +108,7 @@ def _step(self, tensordict: TensorDict):
         action = np.concatenate([action, self._action_padding])
         # TODO(rcadene): add info["is_success"] and info["success"] ?
         sum_reward = 0
-        for t in range(self.frame_skip):
+        for _ in range(self.frame_skip):
             raw_obs, reward, done, info = self._env.step(action)
             sum_reward += reward
 

lerobot/common/envs/transforms.py

@@ -5,7 +5,6 @@
 
 
 class Prod(ObservationTransform):
-
     def __init__(self, in_keys: Sequence[NestedKey], prod: float):
         super().__init__()
         self.in_keys = in_keys