diff --git a/README.md b/README.md
index 2d4df622..cf21337e 100644
--- a/README.md
+++ b/README.md
@@ -1 +1,11 @@
-# lerobot
+# LeRobot
+
+## Installation
+
+Install dependencies using `conda`:
+
+```
+conda env create -f environment.yaml
+conda activate lerobot
+```
+
diff --git a/environment.yaml b/environment.yaml
new file mode 100644
index 00000000..d9dea113
--- /dev/null
+++ b/environment.yaml
@@ -0,0 +1,23 @@
+name: lerobot
+dependencies:
+ - python=3.8.16
+ - pytorch::pytorch=1.13.1
+ - pytorch::torchvision=0.14.1
+ - nvidia::cudatoolkit=11.7
+ - anaconda::pip
+ - pip:
+ - cython==0.29.33
+ - mujoco==2.3.2
+ - mujoco-py==2.1.2.14
+ - termcolor
+ - omegaconf
+ - gym==0.21.0
+ - dm-env==1.6
+ - pandas
+ - wandb
+ - moviepy
+ - imageio
+ - gdown
+ # - -e benchmarks/d4rl
+ # TODO: verify this works
+ - git+https://github.com/nicklashansen/simxarm.git@main#egg=simxarm
diff --git a/lerobot/__init__.py b/lerobot/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/lerobot/__version__.py b/lerobot/__version__.py
new file mode 100644
index 00000000..6c8e6b97
--- /dev/null
+++ b/lerobot/__version__.py
@@ -0,0 +1 @@
+__version__ = "0.0.0"
diff --git a/lerobot/configs/default.yaml b/lerobot/configs/default.yaml
new file mode 100644
index 00000000..7c861d7b
--- /dev/null
+++ b/lerobot/configs/default.yaml
@@ -0,0 +1,2 @@
+seed: 1337
+log_dir: logs/2024_01_26_train
diff --git a/lerobot/scripts/download.py b/lerobot/scripts/download.py
new file mode 100644
index 00000000..ef66c5aa
--- /dev/null
+++ b/lerobot/scripts/download.py
@@ -0,0 +1,21 @@
+import os
+import zipfile
+
+import gdown
+
+
+def download():
+ url = "https://drive.google.com/uc?id=1nhxpykGtPDhmQKm-_B8zBSywVRdgeVya"
+ download_path = "data.zip"
+ gdown.download(url, download_path, quiet=False)
+ print("Extracting...")
+ with zipfile.ZipFile(download_path, "r") as zip_f:
+ for member in zip_f.namelist():
+ if member.startswith("data/xarm") and member.endswith(".pkl"):
+ print(member)
+ zip_f.extract(member=member)
+ os.remove(download_path)
+
+
+if __name__ == "__main__":
+ download()
diff --git a/lerobot/scripts/eval.py b/lerobot/scripts/eval.py
new file mode 100644
index 00000000..3b0f61ef
--- /dev/null
+++ b/lerobot/scripts/eval.py
@@ -0,0 +1,71 @@
+from pathlib import Path
+
+import hydra
+import imageio
+import numpy as np
+import torch
+from termcolor import colored
+
+from ..lib.envs import make_env
+from ..lib.utils import set_seed
+
+
+def eval_agent(
+ env, agent, num_episodes: int, save_video: bool = False, video_path: Path = None
+):
+ """Evaluate a trained agent and optionally save a video."""
+ if save_video:
+ assert video_path is not None
+ assert video_path.suffix == ".mp4"
+ episode_rewards = []
+ episode_successes = []
+ episode_lengths = []
+ for i in range(num_episodes):
+ obs, done, ep_reward, t = env.reset(), False, 0, 0
+ ep_success = False
+ if save_video:
+ frames = []
+ while not done:
+ action = agent.act(obs, t0=t == 0, eval_mode=True, step=step)
+ obs, reward, done, info = env.step(action.cpu().numpy())
+ ep_reward += reward
+ if "success" in info and info["success"]:
+ ep_success = True
+ if save_video:
+ frame = env.render(
+ mode="rgb_array",
+ # TODO(rcadene): make height, width, camera_id configurable
+ height=384,
+ width=384,
+ camera_id=0,
+ )
+ frames.append(frame)
+ t += 1
+ episode_rewards.append(float(ep_reward))
+ episode_successes.append(float(ep_success))
+ episode_lengths.append(t)
+ if save_video:
+ frames = np.stack(frames).transpose(0, 3, 1, 2)
+ video_path.parent.mkdir(parents=True, exist_ok=True)
+ # TODO(rcadene): make fps configurable
+ imageio.mimsave(video_path, frames, fps=15)
+ return {
+ "episode_reward": np.nanmean(episode_rewards),
+ "episode_success": np.nanmean(episode_successes),
+ "episode_length": np.nanmean(episode_lengths),
+ }
+
+
+@hydra.main(version_base=None, config_name="default", config_path="../configs")
+def eval(cfg: dict):
+ assert torch.cuda.is_available()
+ set_seed(cfg.seed)
+ print(colored("Log dir:", "yellow", attrs=["bold"]), cfg.log_dir)
+
+ env = make_env(cfg)
+
+ eval_metrics = eval_agent(env, agent, num_episodes=10, save_video=True)
+
+
+if __name__ == "__main__":
+ eval()
diff --git a/lerobot/scripts/train.py b/lerobot/scripts/train.py
new file mode 100644
index 00000000..2a3c7970
--- /dev/null
+++ b/lerobot/scripts/train.py
@@ -0,0 +1,16 @@
+import hydra
+import torch
+from termcolor import colored
+
+from ..lib.utils import set_seed
+
+
+@hydra.main(version_base=None, config_name="default", config_path="../configs")
+def train(cfg: dict):
+ assert torch.cuda.is_available()
+ set_seed(cfg.seed)
+ print(colored("Work dir:", "yellow", attrs=["bold"]), cfg.log_dir)
+
+
+if __name__ == "__main__":
+ train()
diff --git a/lerobot/scripts/visualize.py b/lerobot/scripts/visualize.py
new file mode 100644
index 00000000..64d24504
--- /dev/null
+++ b/lerobot/scripts/visualize.py
@@ -0,0 +1,80 @@
+import pickle
+from pathlib import Path
+
+import imageio
+import simxarm
+
+if __name__ == "__main__":
+
+ task = "lift"
+ dataset_dir = Path(f"data/xarm_{task}_medium")
+ dataset_path = dataset_dir / f"buffer.pkl"
+ print(f"Using offline dataset '{dataset_path}'")
+ with open(dataset_path, "rb") as f:
+ dataset_dict = pickle.load(f)
+
+ required_keys = [
+ "observations",
+ "next_observations",
+ "actions",
+ "rewards",
+ "dones",
+ "masks",
+ ]
+ for k in required_keys:
+ if k not in dataset_dict and k[:-1] in dataset_dict:
+ dataset_dict[k] = dataset_dict.pop(k[:-1])
+
+ out_dir = Path("tmp/2023_01_26_xarm_lift_medium")
+ out_dir.mkdir(parents=True, exist_ok=True)
+
+ frames = dataset_dict["observations"]["rgb"][:100]
+ frames = frames.transpose(0, 2, 3, 1)
+ imageio.mimsave(out_dir / "test.mp4", frames, fps=30)
+
+ frames = []
+ cfg = {}
+
+ env = simxarm.make(
+ task=task,
+ obs_mode="all",
+ image_size=84,
+ action_repeat=cfg.get("action_repeat", 1),
+ frame_stack=cfg.get("frame_stack", 1),
+ seed=1,
+ )
+
+ obs = env.reset()
+ frame = env.render(mode="rgb_array", width=384, height=384)
+ frames.append(frame)
+
+ # def is_first_obs(obs):
+ # nonlocal first_obs
+ # print(((dataset_dict["observations"]["state"][i]-obs["state"])**2).sum())
+ # print(((dataset_dict["observations"]["rgb"][i]-obs["rgb"])**2).sum())
+
+ for i in range(25):
+ action = dataset_dict["actions"][i]
+
+ print(f"#{i}")
+ # print(obs["state"])
+ # print(dataset_dict["observations"]["state"][i])
+ print(((dataset_dict["observations"]["state"][i] - obs["state"]) ** 2).sum())
+ print(((dataset_dict["observations"]["rgb"][i] - obs["rgb"]) ** 2).sum())
+
+ obs, reward, done, info = env.step(action)
+ frame = env.render(mode="rgb_array", width=384, height=384)
+ frames.append(frame)
+
+ print(reward)
+ print(dataset_dict["rewards"][i])
+
+ print(done)
+ print(dataset_dict["dones"][i])
+
+ if dataset_dict["dones"][i]:
+ obs = env.reset()
+ frame = env.render(mode="rgb_array", width=384, height=384)
+ frames.append(frame)
+
+ # imageio.mimsave(out_dir / 'test_rollout.mp4', frames, fps=60)
diff --git a/setup.py b/setup.py
new file mode 100644
index 00000000..f503e843
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,159 @@
+"""A setuptools based setup module.
+
+See:
+https://packaging.python.org/en/latest/distributing.html
+https://github.com/pypa/sampleproject
+"""
+
+# To use a consistent encoding
+from codecs import open
+from os import path
+
+# Always prefer setuptools over distutils
+from setuptools import find_packages, setup
+
+here = path.abspath(path.dirname(__file__))
+
+# Get the long description from the README file
+with open(path.join(here, "README.md"), encoding="utf-8") as f:
+ long_description = f.read()
+
+# Arguments marked as "Required" below must be included for upload to PyPI.
+# Fields marked as "Optional" may be commented out.
+
+# https://stackoverflow.com/questions/458550/standard-way-to-embed-version-into-python-package/16084844#16084844
+exec(open(path.join(here, "lerobot", "__version__.py")).read())
+setup(
+ # This is the name of your project. The first time you publish this
+ # package, this name will be registered for you. It will determine how
+ # users can install this project, e.g.:
+ #
+ # $ pip install sampleproject
+ #
+ # And where it will live on PyPI: https://pypi.org/project/sampleproject/
+ #
+ # There are some restrictions on what makes a valid project name
+ # specification here:
+ # https://packaging.python.org/specifications/core-metadata/#name
+ name="lerobot", # Required
+ # Versions should comply with PEP 440:
+ # https://www.python.org/dev/peps/pep-0440/
+ #
+ # For a discussion on single-sourcing the version across setup.py and the
+ # project code, see
+ # https://packaging.python.org/en/latest/single_source_version.html
+ version=__version__, # noqa: F821 # Required
+ # This is a one-line description or tagline of what your project does. This
+ # corresponds to the "Summary" metadata field:
+ # https://packaging.python.org/specifications/core-metadata/#summary
+ description="Le robot is learning", # Required
+ # This is an optional longer description of your project that represents
+ # the body of text which users will see when they visit PyPI.
+ #
+ # Often, this is the same as your README, so you can just read it in from
+ # that file directly (as we have already done above)
+ #
+ # This field corresponds to the "Description" metadata field:
+ # https://packaging.python.org/specifications/core-metadata/#description-optional
+ long_description=long_description, # Optional
+ # This should be a valid link to your project's main homepage.
+ #
+ # This field corresponds to the "Home-Page" metadata field:
+ # https://packaging.python.org/specifications/core-metadata/#home-page-optional
+ url="https://github.com/cadene/lerobot", # Optional
+ # This should be your name or the name of the organization which owns the
+ # project.
+ author="Remi Cadene", # Optional
+ # This should be a valid email address corresponding to the author listed
+ # above.
+ author_email="re.cadene@gmail.com", # Optional
+ # Classifiers help users find your project by categorizing it.
+ #
+ # For a list of valid classifiers, see
+ # https://pypi.python.org/pypi?%3Aaction=list_classifiers
+ classifiers=[ # Optional
+ # How mature is this project? Common values are
+ # 3 - Alpha
+ # 4 - Beta
+ # 5 - Production/Stable
+ "Development Status :: 3 - Alpha",
+ # Indicate who your project is intended for
+ "Intended Audience :: Developers",
+ "Topic :: Software Development :: Build Tools",
+ # Pick your license as you wish
+ "License :: OSI Approved :: MIT License",
+ # Specify the Python versions you support here. In particular, ensure
+ # that you indicate whether you support Python 2, Python 3 or both.
+ "Programming Language :: Python :: 3.7",
+ ],
+ # This field adds keywords for your project which will appear on the
+ # project page. What does your project relate to?
+ #
+ # Note that this is a string of words separated by whitespace, not a list.
+ keywords="pytorch framework bootstrap deep learning scaffolding", # Optional
+ # You can just specify package directories manually here if your project is
+ # simple. Or you can use find_packages().
+ #
+ # Alternatively, if you just want to distribute a single Python file, use
+ # the `py_modules` argument instead as follows, which will expect a file
+ # called `my_module.py` to exist:
+ #
+ # py_modules=["my_module"],
+ #
+ packages=find_packages(
+ exclude=[
+ "data",
+ "logs",
+ ]
+ ),
+ # This field lists other packages that your project depends on to run.
+ # Any package you put here will be installed by pip when your project is
+ # installed, so they must be valid existing projects.
+ #
+ # For an analysis of "install_requires" vs pip's requirements files see:
+ # https://packaging.python.org/en/latest/requirements.html
+ install_requires=[
+ "torch",
+ "numpy",
+ "argparse",
+ ],
+ # List additional groups of dependencies here (e.g. development
+ # dependencies). Users will be able to install these using the "extras"
+ # syntax, for example:
+ #
+ # $ pip install sampleproject[dev]
+ #
+ # Similar to `install_requires` above, these must be valid existing
+ # projects.
+ # extras_require={ # Optional
+ # 'dev': ['check-manifest'],
+ # 'test': ['coverage'],
+ # },
+ # If there are data files included in your packages that need to be
+ # installed, specify them here.
+ #
+ # If using Python 2.6 or earlier, then these have to be included in
+ # MANIFEST.in as well.
+ # package_data={ # Optional
+ # 'sample': ['package_data.dat'],
+ # },
+ include_package_data=True,
+ # Although 'package_data' is the preferred approach, in some case you may
+ # need to place data files outside of your packages. See:
+ # http://docs.python.org/3.4/distutils/setupscript.html#installing-additional-files
+ #
+ # In this case, 'data_file' will be installed into '<sys.prefix>/my_data'
+ # data_files=[('my_data', ['data/data_file'])], # Optional
+ # To provide executable scripts, use entry points in preference to the
+ # "scripts" keyword. Entry points provide cross-platform support and allow
+ # `pip` to create the appropriate form of executable for the target
+ # platform.
+ #
+ # For example, the following would provide a command called `sample` which
+ # executes the function `main` from this package when invoked:
+ # entry_points={ # Optional
+ # 'console_scripts': [
+ # 'sample=sample:main',
+ # ],
+ # },
+)
diff --git a/test/test_envs.py b/test/test_envs.py
new file mode 100644
index 00000000..49b1cae2
--- /dev/null
+++ b/test/test_envs.py
@@ -0,0 +1,53 @@
+import pytest
+from tensordict import TensorDict
+from torchrl.envs.utils import check_env_specs, step_mdp
+
+from lerobot.lib.envs import SimxarmEnv
+
+
+@pytest.mark.parametrize(
+ "task,from_pixels,pixels_only",
+ [
+ ("lift", False, False),
+ ("lift", True, False),
+ ("lift", True, True),
+ ("reach", False, False),
+ ("reach", True, False),
+ ("push", False, False),
+ ("push", True, False),
+ ("peg_in_box", False, False),
+ ("peg_in_box", True, False),
+ ],
+)
+def test_simxarm(task, from_pixels, pixels_only):
+ env = SimxarmEnv(
+ task,
+ from_pixels=from_pixels,
+ pixels_only=pixels_only,
+ image_size=84 if from_pixels else None,
+ )
+ check_env_specs(env)
+
+ print("observation_spec:", env.observation_spec)
+ print("action_spec:", env.action_spec)
+ print("reward_spec:", env.reward_spec)
+
+ td = env.reset()
+ print("reset tensordict", td)
+
+ td = env.rand_step(td)
+ print("random step tensordict", td)
+
+ def simple_rollout(steps=100):
+ # preallocate:
+ data = TensorDict({}, [steps])
+ # reset
+ _data = env.reset()
+ for i in range(steps):
+ _data["action"] = env.action_spec.rand()
+ _data = env.step(_data)
+ data[i] = _data
+ _data = step_mdp(_data, keep_other=True)
+ return data
+
+ print("data from rollout:", simple_rollout(100))