First real commit, simxarm env added with torchrl!

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

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

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

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)