Support MTBench (#15)

This PR incorporates MTBench into the current codebase, as a good demonstration that shows how to use FastTD3 for multi-task setup. - Add support for MTBench along with its wrapper - Add support for per-task reward normalizer useful for multi-task RL, motivated by BRC paper (https://arxiv.org/abs/2505.23150v1)
2025-06-20 21:52:43 -07:00
parent 3facede77d
commit cef44108d8
8 changed files with 679 additions and 34 deletions
@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import torch
+from omegaconf import OmegaConf
+
+import isaacgym
+import isaacgymenvs
+
+
+class MTBenchEnv:
+    def __init__(
+        self,
+        task_name: str,
+        device_id: int,
+        num_envs: int,
+        seed: int,
+    ):
+        # NOTE: Currently, we only support Meta-World-v2 MT-10/MT-50 in MTBench
+        task_config = MTBENCH_MW2_CONFIG.copy()
+        if task_name == "meta-world-v2-mt10":
+            # MT-10 Setup
+            assert num_envs == 4096, "MT-10 only supports 4096 environments (for now)"
+            self.num_tasks = 10
+            task_config["env"]["tasks"] = [4, 16, 17, 18, 28, 31, 38, 40, 48, 49]
+            task_config["env"]["taskEnvCount"] = [410] * 6 + [409] * 4
+        elif task_name == "meta-world-v2-mt50":
+            # MT-50 Setup
+            self.num_tasks = 50
+            assert num_envs == 8192, "MT-50 only supports 8192 environments (for now)"
+            task_config["env"]["tasks"] = list(range(50))
+            task_config["env"]["taskEnvCount"] = [164] * 42 + [163] * 8  # 6888 + 1304
+        else:
+            raise ValueError(f"Unsupported task name: {task_name}")
+        task_config["env"]["numEnvs"] = num_envs
+        task_config["env"]["numObservations"] = 39 + self.num_tasks
+        task_config["env"]["seed"] = seed
+
+        # Convert dictionary to OmegaConf object
+        env_cfg = {"task": task_config}
+        env_cfg = OmegaConf.create(env_cfg)
+
+        self.env = isaacgymenvs.make(
+            task=env_cfg.task.name,
+            num_envs=num_envs,
+            sim_device=f"cuda:{device_id}",
+            rl_device=f"cuda:{device_id}",
+            seed=seed,
+            headless=True,
+            cfg=env_cfg,
+        )
+
+        self.num_envs = num_envs
+        self.asymmetric_obs = False
+        self.num_obs = self.env.observation_space.shape[0]
+        assert (
+            self.num_obs == 39 + self.num_tasks
+        ), "MTBench observation space is 39 + num_tasks (one-hot vector)"
+        self.num_privileged_obs = 0
+        self.num_actions = self.env.action_space.shape[0]
+        self.max_episode_steps = self.env.max_episode_length
+
+    def reset(self) -> torch.Tensor:
+        """Reset the environment."""
+        # TODO: Check if we need no_grad and detach here
+        with torch.no_grad():  # do we need this?
+            self.env.reset_idx(torch.arange(self.num_envs, device=self.env.device))
+            obs_dict = self.env.reset()
+            return obs_dict["obs"].detach()
+
+    def step(
+        self, actions: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict]:
+        """Step the environment."""
+        assert isinstance(actions, torch.Tensor)
+
+        # TODO: Check if we need no_grad and detach here
+        with torch.no_grad():
+            obs_dict, rew, dones, infos = self.env.step(actions.detach())
+            truncations = infos["time_outs"]
+            info_ret = {"time_outs": truncations.detach()}
+            if "episode" in infos:
+                info_ret["episode"] = infos["episode"]
+            # NOTE: There's really no way to get the raw observations from IsaacGym
+            # We just use the 'reset_obs' as next_obs, unfortunately.
+            info_ret["observations"] = {"raw": {"obs": obs_dict["obs"].detach()}}
+            return obs_dict["obs"].detach(), rew.detach(), dones.detach(), info_ret
+
+    def render(self):
+        raise NotImplementedError(
+            "We don't support rendering for IsaacLab environments"
+        )
+
+
+MTBENCH_MW2_CONFIG = {
+    "name": "meta-world-v2",
+    "physics_engine": "physx",
+    "env": {
+        "numEnvs": 1,
+        "envSpacing": 1.5,
+        "episodeLength": 150,
+        "enableDebugVis": False,
+        "clipObservations": 5.0,
+        "clipActions": 1.0,
+        "aggregateMode": 3,
+        "actionScale": 0.01,
+        "resetNoise": 0.15,
+        "tasks": [0],
+        "taskEnvCount": [4096],
+        "init_at_random_progress": True,
+        "exemptedInitAtRandomProgressTasks": [],
+        "taskEmbedding": True,
+        "taskEmbeddingType": "one_hot",
+        "seed": 42,
+        "cameraRenderingInterval": 5000,
+        "cameraWidth": 1024,
+        "cameraHeight": 1024,
+        "sparse_reward": False,
+        "termination_on_success": False,
+        "reward_scale": 1.0,
+        "fixed": False,
+        "numObservations": None,
+        "numActions": 4,
+    },
+    "enableCameraSensors": False,
+    "sim": {
+        "dt": 0.01667,
+        "substeps": 2,
+        "up_axis": "z",
+        "use_gpu_pipeline": True,
+        "gravity": [0.0, 0.0, -9.81],
+        "physx": {
+            "num_threads": 4,
+            "solver_type": 1,
+            "use_gpu": True,
+            "num_position_iterations": 8,
+            "num_velocity_iterations": 1,
+            "contact_offset": 0.005,
+            "rest_offset": 0.0,
+            "bounce_threshold_velocity": 0.2,
+            "max_depenetration_velocity": 1000.0,
+            "default_buffer_size_multiplier": 10.0,
+            "max_gpu_contact_pairs": 1048576,
+            "num_subscenes": 4,
+            "contact_collection": 0,
+        },
+    },
+    "task": {"randomize": False},
+}