Integrate diffusion policy

.github/workflows/test.yml

@@ -69,10 +69,7 @@ jobs:
           key: venv-${{ steps.setup-python.outputs.python-version }}-${{ env.POETRY_VERSION }}-${{ hashFiles('**/poetry.lock') }}
       - name: Install dependencies
         if: steps.restore-dependencies-cache.outputs.cache-hit != 'true'
-        run: |
-          poetry install --no-interaction --no-root
-          git clone https://github.com/real-stanford/diffusion_policy
-          cp -r diffusion_policy/diffusion_policy $(poetry env info -p)/lib/python3.10/site-packages/
+        run: poetry install --no-interaction --no-root
       - name: Save cached venv
         if: |
             steps.restore-dependencies-cache.outputs.cache-hit != 'true' &&

.gitignore

@@ -1,6 +1,3 @@
-# Custom
-diffusion_policy
-
 # Logging
 logs
 tmp

.pre-commit-config.yaml

@@ -1,4 +1,4 @@
-exclude: ^(data/|tests/|diffusion_policy/)
+exclude: ^(data/|tests/)
 default_language_version:
     python: python3.10
 repos:

README.md

@@ -24,12 +24,6 @@ mkdir ~/tmp
 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/
-```
 
 ## Usage
 

lerobot/common/datasets/pusht.py

@@ -8,20 +8,19 @@
 import torch
 import torchrl
 import tqdm
-from diffusion_policy.common.replay_buffer import ReplayBuffer as DiffusionPolicyReplayBuffer
-from diffusion_policy.env.pusht.pusht_env import pymunk_to_shapely
 from tensordict import TensorDict
 from torchrl.data.replay_buffers.samplers import SliceSampler
 from torchrl.data.replay_buffers.storages import TensorStorage
 from torchrl.data.replay_buffers.writers import Writer
 
 from lerobot.common.datasets.abstract import AbstractExperienceReplay
 from lerobot.common.datasets.utils import download_and_extract_zip
+from lerobot.common.envs.pusht.pusht_env import pymunk_to_shapely
+from lerobot.common.policies.diffusion.replay_buffer import ReplayBuffer as DiffusionPolicyReplayBuffer
 
 # as define in env
 SUCCESS_THRESHOLD = 0.95  # 95% coverage,
 
-DEFAULT_TEE_MASK = pymunk.ShapeFilter.ALL_MASKS()
 PUSHT_URL = "https://diffusion-policy.cs.columbia.edu/data/training/pusht.zip"
 PUSHT_ZARR = Path("pusht/pusht_cchi_v7_replay.zarr")
 
@@ -49,8 +48,10 @@ def add_tee(
     angle,
     scale=30,
     color="LightSlateGray",
-    mask=DEFAULT_TEE_MASK,
+    mask=None,
 ):
+    if mask is None:
+        mask = pymunk.ShapeFilter.ALL_MASKS()
     mass = 1
     length = 4
     vertices1 = [

lerobot/common/envs/factory.py

@@ -18,7 +18,7 @@ def make_env(cfg, transform=None):
         kwargs["task"] = cfg.env.task
         clsfunc = SimxarmEnv
     elif cfg.env.name == "pusht":
-        from lerobot.common.envs.pusht import PushtEnv
+        from lerobot.common.envs.pusht.env import PushtEnv
 
         # assert kwargs["seed"] > 200, "Seed 0-200 are used for the demonstration dataset, so we don't want to seed the eval env with this range."
 

lerobot/common/envs/pusht.py → lerobot/common/envs/pusht/env.py

@@ -17,7 +17,6 @@
 from lerobot.common.utils import set_seed
 
 _has_gym = importlib.util.find_spec("gym") is not None
-_has_diffpolicy = importlib.util.find_spec("diffusion_policy") is not None and _has_gym
 
 
 class PushtEnv(EnvBase):
@@ -45,17 +44,15 @@ def __init__(
         if from_pixels:
             assert image_size
 
-        if not _has_diffpolicy:
-            raise ImportError("Cannot import diffusion_policy.")
         if not _has_gym:
             raise ImportError("Cannot import gym.")
 
         # 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
+        # from lerobot.common.envs.pusht.pusht_env import PushTEnv
 
         if not from_pixels:
             raise NotImplementedError("Use PushTEnv, instead of PushTImageEnv")
-        from diffusion_policy.env.pusht.pusht_image_env import PushTImageEnv
+        from lerobot.common.envs.pusht.pusht_image_env import PushTImageEnv
 
         self._env = PushTImageEnv(render_size=self.image_size)