From 98f886974376603ad845b8473a274e3bdc1756a4 Mon Sep 17 00:00:00 2001
From: Simon Alibert <alibert.sim@gmail.com>
Date: Wed, 28 Feb 2024 10:57:08 +0100
Subject: [PATCH] WIP

---
 README.md                          | 35 ++++++++++++++++++++++++++++++
 lerobot/common/datasets/factory.py |  3 ++-
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 925d143a..6850e003 100644
--- a/README.md
+++ b/README.md
@@ -9,6 +9,41 @@ conda env create -f environment.yaml
 conda activate lerobot
 ```
 
+Install `torchrl`, `tensordict` and `diffusion_policy` dev builds
+```
+cd path/to/root
+git clone https://github.com/pytorch/tensordict
+git clone https://github.com/pytorch/rl
+git clone https://github.com/real-stanford/diffusion_policy
+cd tensordict
+python setup.py develop
+cd ../rl
+python setup.py develop
+cd ../diffusion_policy
+python setup.py develop
+```
+
+Install additional modules
+```
+pip install \
+    hydra \
+    termcolor \
+    einops \
+    pygame \
+    pymunk \
+    zarr \
+    gym \
+    shapely \
+    opencv-python \
+    scikit-image \
+    mpmath==1.3.0 \
+```
+
+Fix Hydra
+```
+pip install hydra-core --upgrade
+```
+
 **dev**
 
 ```
diff --git a/lerobot/common/datasets/factory.py b/lerobot/common/datasets/factory.py
index 54155478..6b8037fe 100644
--- a/lerobot/common/datasets/factory.py
+++ b/lerobot/common/datasets/factory.py
@@ -2,7 +2,8 @@ import torch
 
 from lerobot.common.datasets.pusht import PushtExperienceReplay
 from lerobot.common.datasets.simxarm import SimxarmExperienceReplay
-from rl.torchrl.data.replay_buffers.samplers import PrioritizedSliceSampler
+# from rl.torchrl.data.replay_buffers.samplers import PrioritizedSliceSampler
+from torchrl.data.replay_buffers import PrioritizedSliceSampler
 
 # TODO(rcadene): implement