From 1144819c29b30f84917dc61e288b4a6bf79c6763 Mon Sep 17 00:00:00 2001
From: Cadene <re.cadene@gmail.com>
Date: Mon, 29 Jan 2024 12:49:30 +0000
Subject: [PATCH] First real commit, simxarm env added with torchrl!

---
 README.md                    |  12 ++-
 environment.yaml             |  23 +++++
 lerobot/__init__.py          |   0
 lerobot/__version__.py       |   1 +
 lerobot/configs/default.yaml |   2 +
 lerobot/scripts/download.py  |  21 +++++
 lerobot/scripts/eval.py      |  71 ++++++++++++++++
 lerobot/scripts/train.py     |  16 ++++
 lerobot/scripts/visualize.py |  80 ++++++++++++++++++
 setup.py                     | 159 +++++++++++++++++++++++++++++++++++
 test/test_envs.py            |  53 ++++++++++++
 11 files changed, 437 insertions(+), 1 deletion(-)
 create mode 100644 environment.yaml
 create mode 100644 lerobot/__init__.py
 create mode 100644 lerobot/__version__.py
 create mode 100644 lerobot/configs/default.yaml
 create mode 100644 lerobot/scripts/download.py
 create mode 100644 lerobot/scripts/eval.py
 create mode 100644 lerobot/scripts/train.py
 create mode 100644 lerobot/scripts/visualize.py
 create mode 100644 setup.py
 create mode 100644 test/test_envs.py

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))