Fixes build errors due to name conflicts

reppo_alg/env_utils/torch_wrappers/humanoid_bench_env.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import gymnasium as gym
+import numpy as np
+import torch
+from gymnasium.wrappers import TimeLimit
+from loguru import logger as log
+from stable_baselines3.common.vec_env import SubprocVecEnv
+
+# Disable all logging below CRITICAL level
+log.remove()
+log.add(lambda msg: False, level="CRITICAL")
+
+
+def make_env(env_name, rank, render_mode=None, seed=0):
+    """
+    Utility function for multiprocessed env.
+
+    :param rank: (int) index of the subprocess
+    :param seed: (int) the inital seed for RNG
+    """
+
+    if env_name in [
+        "h1hand-push-v0",
+        "h1-push-v0",
+        "h1hand-cube-v0",
+        "h1cube-v0",
+        "h1hand-basketball-v0",
+        "h1-basketball-v0",
+        "h1hand-kitchen-v0",
+        "h1-kitchen-v0",
+    ]:
+        max_episode_steps = 500
+    else:
+        max_episode_steps = 1000
+
+    def _init():
+        env = gym.make(env_name, render_mode=render_mode)
+        env = TimeLimit(env, max_episode_steps=max_episode_steps)
+        env.unwrapped.seed(seed + rank)
+
+        return env
+
+    return _init
+
+
+class HumanoidBenchEnv:
+    """Wraps HumanoidBench environment to support parallel environments."""
+
+    def __init__(self, env_name, num_envs=1, render_mode=None, device=None):
+        # NOTE: HumanoidBench action space is already normalized to [-1, 1]
+        device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.sim_device = device
+        self.num_envs = num_envs
+
+        # Create the base environment
+        self.envs = SubprocVecEnv(
+            [make_env(env_name, i, render_mode=render_mode) for i in range(num_envs)]
+        )
+
+        if env_name in [
+            "h1hand-push-v0",
+            "h1-push-v0",
+            "h1hand-cube-v0",
+            "h1cube-v0",
+            "h1hand-basketball-v0",
+            "h1-basketball-v0",
+            "h1hand-kitchen-v0",
+            "h1-kitchen-v0",
+        ]:
+            self.max_episode_steps = 500
+        else:
+            self.max_episode_steps = 1000
+
+        # For compatibility with MuJoCo Playground
+        self.asymmetric_obs = False  # For comptatibility with MuJoCo Playground
+        self.num_obs = self.envs.observation_space.shape[-1]
+        self.num_actions = self.envs.action_space.shape[-1]
+
+    def reset(self):
+        """Reset the environment."""
+        observations = self.envs.reset()
+        observations = torch.from_numpy(observations).to(
+            device=self.sim_device, dtype=torch.float
+        )
+        return observations
+
+    def render(self):
+        assert self.num_envs == 1, (
+            "Currently only supports single environment rendering"
+        )
+        return self.envs.render()
+
+    def step(self, actions):
+        assert isinstance(actions, torch.Tensor)
+        actions = actions.cpu().numpy()
+
+        observations, rewards, dones, raw_infos = self.envs.step(actions)
+
+        # This will be used for getting 'true' next observations
+        infos = dict()
+        infos["observations"] = {"raw": {"obs": observations.copy()}}
+        truncateds = np.zeros_like(dones)
+        for i in range(self.num_envs):
+            if raw_infos[i].get("TimeLimit.truncated", False):
+                truncateds[i] = True
+                infos["observations"]["raw"]["obs"][i] = raw_infos[i][
+                    "terminal_observation"
+                ]
+
+        observations = torch.from_numpy(observations).to(
+            device=self.sim_device, dtype=torch.float
+        )
+        rewards = torch.from_numpy(rewards).to(
+            device=self.sim_device, dtype=torch.float
+        )
+        dones = torch.from_numpy(dones).to(device=self.sim_device)
+        truncateds = torch.from_numpy(truncateds).to(device=self.sim_device)
+        infos["observations"]["raw"]["obs"] = torch.from_numpy(
+            infos["observations"]["raw"]["obs"]
+        ).to(device=self.sim_device, dtype=torch.float)
+        infos["time_outs"] = truncateds
+
+        return observations, rewards, dones, infos

reppo_alg/env_utils/torch_wrappers/isaaclab_env.py

@@ -0,0 +1,81 @@
+from typing import Optional
+
+import gymnasium as gym
+import torch
+from isaaclab.app import AppLauncher
+from isaaclab_tasks.utils.parse_cfg import parse_env_cfg
+
+app_launcher = AppLauncher(headless=True)
+simulation_app = app_launcher.app
+
+
+
+class IsaacLabEnv:
+    """Wrapper for IsaacLab environments to be compatible with MuJoCo Playground"""
+
+    def __init__(
+        self,
+        task_name: str,
+        device: str,
+        num_envs: int,
+        seed: int,
+        action_bounds: Optional[float] = None,
+    ):
+        env_cfg = parse_env_cfg(
+            task_name,
+            device=device,
+            num_envs=num_envs,
+        )
+        env_cfg.seed = seed
+        self.seed = seed
+        self.envs = gym.make(task_name, cfg=env_cfg, render_mode=None)
+
+        self.num_envs = self.envs.unwrapped.num_envs
+        self.max_episode_steps = self.envs.unwrapped.max_episode_length
+        self.action_bounds = action_bounds
+        self.num_obs = self.envs.unwrapped.single_observation_space["policy"].shape[0]
+        self.asymmetric_obs = "critic" in self.envs.unwrapped.single_observation_space
+        if self.asymmetric_obs:
+            self.num_privileged_obs = self.envs.unwrapped.single_observation_space[
+                "critic"
+            ].shape[0]
+        else:
+            self.num_privileged_obs = 0
+        self.num_actions = self.envs.unwrapped.single_action_space.shape[0]
+
+    def reset(self, random_start_init: bool = True) -> torch.Tensor:
+        obs_dict, _ = self.envs.reset()
+        # NOTE: decorrelate episode horizons like RSL‑RL
+        if random_start_init:
+            self.envs.unwrapped.episode_length_buf = torch.randint_like(
+                self.envs.unwrapped.episode_length_buf, high=int(self.max_episode_steps)
+            )
+        return obs_dict["policy"]
+
+    def reset_with_critic_obs(self) -> tuple[torch.Tensor, torch.Tensor]:
+        obs_dict, _ = self.envs.reset()
+        return obs_dict["policy"], obs_dict["critic"]
+
+    def step(
+        self, actions: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict]:
+        if self.action_bounds is not None:
+            actions = torch.clamp(actions, -1.0, 1.0) * self.action_bounds
+        obs_dict, rew, terminations, truncations, infos = self.envs.step(actions)
+        dones = (terminations | truncations).to(dtype=torch.long)
+        obs = obs_dict["policy"]
+        critic_obs = obs_dict["critic"] if self.asymmetric_obs else None
+        info_ret = {"time_outs": truncations, "observations": {"critic": critic_obs}}
+        # NOTE: There's really no way to get the raw observations from IsaacLab
+        # We just use the 'reset_obs' as next_obs, unfortunately.
+        # See https://github.com/isaac-sim/IsaacLab/issues/1362
+        info_ret["observations"]["raw"] = {
+            "obs": obs,
+            "critic_obs": critic_obs,
+        }
+        return obs, rew, dones, info_ret
+
+    def render(self):
+        raise NotImplementedError(
+            "We don't support rendering for IsaacLab environments"
+        )

reppo_alg/env_utils/torch_wrappers/maniskill_wrapper.py

@@ -0,0 +1,89 @@
+from gymnasium import Wrapper
+import torch
+
+
+class ManiSkillWrapper(Wrapper):
+    """
+    A wrapper for ManiSkill environments to ensure compatibility with the expected API.
+    This wrapper is used to handle the ManiSkill environments in a way that is consistent
+    with the other environments in the codebase.
+    """
+
+    def __init__(self, env, max_episode_steps: int, partial_reset, device: str):
+        super().__init__(env)
+        self.action_space = env.action_space
+        self.observation_space = env.observation_space
+        self.metadata = env.metadata
+        self.asymmetric_obs = False
+        self.max_episode_steps = max_episode_steps
+
+        self.partial_reset = partial_reset
+
+        self.returns = torch.zeros(env.num_envs, dtype=torch.float32, device=device)
+        self.episode_len = torch.zeros(env.num_envs, dtype=torch.float32, device=device)
+        self.success = torch.zeros(env.num_envs, dtype=torch.float32, device=device)
+
+    @property
+    def unwrapped(self):
+        """
+        Returns the underlying environment.
+        """
+        return self.env
+
+    @property
+    def num_actions(self):
+        """
+        Returns the number of actions in the action space.
+        """
+        return self.action_space.shape[1]
+
+    @property
+    def num_obs(self):
+        """
+        Returns the number of observations in the observation space.
+        """
+        return self.observation_space.shape[1]
+
+    def reset(self, seed=None, options=dict()):
+        """
+        Resets the environment and returns the initial observation.
+        """
+        return self.env.reset(seed=seed, options=options)
+
+    def step(self, action):
+        """
+        Takes a step in the environment with the given action.
+        Returns the next observation, reward, done, and info.
+        """
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        if "final_info" in info:
+            self.returns = (
+                info["final_info"]["episode"]["return"] * info["_final_info"].float()
+                + (1.0 - info["_final_info"].float()) * self.returns
+            )
+            self.episode_len = (
+                info["final_info"]["episode"]["episode_len"]
+                * info["_final_info"].float()
+                + (1.0 - info["_final_info"].float()) * self.episode_len
+            )
+            self.success = (
+                info["final_info"]["episode"]["success_once"]
+                * info["_final_info"].float()
+                + (1.0 - info["_final_info"].float()) * self.success
+            )
+        info["log_info"] = {
+            "return": self.returns,
+            "episode_len": self.episode_len,
+            "success": self.success,
+        }
+        if self.partial_reset:
+            # maniskill continues bootstrap on terminated, which playground does on truncated.
+            # This unifies the interfaces in a very hacky way
+            done = torch.zeros_like(
+                terminated, dtype=torch.bool, device=terminated.device
+            )
+            truncated = torch.logical_or(terminated, truncated)
+        else:
+            done = torch.logical_or(terminated, truncated)
+            truncated = torch.zeros_like(done, dtype=torch.bool, device=done.device)
+        return obs, reward, done, truncated, info

reppo_alg/env_utils/torch_wrappers/mtbench_env.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import isaacgymenvs
+import torch
+from omegaconf import OmegaConf
+
+
+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))
+            self.env.cumulatives["rewards"][:] = 0
+            self.env.cumulatives["success"][:] = 0
+            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},
+}

reppo_alg/env_utils/torch_wrappers/mujoco_playground_env.py

@@ -0,0 +1,138 @@
+import jax
+from mujoco_playground import registry, wrapper_torch
+
+jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
+jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
+jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
+
+
+class PlaygroundEvalEnvWrapper:
+    def __init__(self, eval_env, max_episode_steps, env_name, num_eval_envs, seed):
+        """
+        Wrapper used for evaluation / rendering environments.
+        Note that this is different from training environments that are
+        wrapped with RSLRLBraxWrapper.
+        """
+        self.env = eval_env
+        self.env_name = env_name
+        self.num_envs = num_eval_envs
+        self.jit_reset = jax.jit(jax.vmap(self.env.reset))
+        self.jit_step = jax.jit(jax.vmap(self.env.step))
+
+        if isinstance(self.env.unwrapped.observation_size, dict):
+            self.asymmetric_obs = True
+        else:
+            self.asymmetric_obs = False
+
+        self.key = jax.random.PRNGKey(seed)
+        self.key_reset = jax.random.split(self.key, num_eval_envs)
+        self.max_episode_steps = max_episode_steps
+
+    def reset(self):
+        self.state = self.jit_reset(self.key_reset)
+        if self.asymmetric_obs:
+            obs = wrapper_torch._jax_to_torch(self.state.obs["state"])
+        else:
+            obs = wrapper_torch._jax_to_torch(self.state.obs)
+        return obs
+
+    def step(self, actions):
+        self.state = self.jit_step(self.state, wrapper_torch._torch_to_jax(actions))
+        if self.asymmetric_obs:
+            next_obs = wrapper_torch._jax_to_torch(self.state.obs["state"])
+        else:
+            next_obs = wrapper_torch._jax_to_torch(self.state.obs)
+        rewards = wrapper_torch._jax_to_torch(self.state.reward)
+        dones = wrapper_torch._jax_to_torch(self.state.done)
+        return next_obs, rewards, dones, dones, None
+
+
+class RandomizeInitialWrapper(wrapper_torch.RSLRLBraxWrapper):
+    """
+    Wrapper to randomize the initial state of the environment.
+    This is useful for domain randomization experiments.
+    """
+
+    def reset(self):
+        print("Resetting environment with randomization")
+        obs = super().reset()
+        self.env_state.info["steps"] = jax.random.randint(
+            self.key, self.env_state.info["steps"].shape, 0, 1000
+        ).astype(jax.numpy.float32)
+        print(obs)
+        return obs
+
+    def reset_with_critic_obs(self):
+        print("Resetting environment with randomization and critic obs")
+        obs, critic_obs = super().reset_with_critic_obs()
+        self.env_state.info["steps"] = jax.random.randint(
+            self.key, self.env_state.info["steps"].shape, 0, 1000
+        ).astype(jax.numpy.float32)
+        return obs, critic_obs
+
+    def step(self, action):
+        obs, reward, done, info = super().step(action)
+        return obs, reward, done, done, info
+
+
+def make_env(
+    env_name,
+    seed,
+    num_envs,
+    num_eval_envs,
+    device_rank,
+    use_tuned_reward=False,
+    use_domain_randomization=False,
+    use_push_randomization=False,
+):
+    # Make training environment
+    train_env_cfg = registry.get_default_config(env_name)
+    is_humanoid_task = env_name in [
+        "G1JoystickRoughTerrain",
+        "G1JoystickFlatTerrain",
+        "T1JoystickRoughTerrain",
+        "T1JoystickFlatTerrain",
+    ]
+
+    if use_tuned_reward and is_humanoid_task:
+        # NOTE: Tuned reward for G1. Used for producing Figure 7 in the paper.
+        # Somehow it works reasonably for T1 as well.
+        # However, see `sim2real.md` for sim-to-real RL with Booster T1
+        train_env_cfg.reward_config.scales.energy = -5e-5
+        train_env_cfg.reward_config.scales.action_rate = -1e-1
+        train_env_cfg.reward_config.scales.torques = -1e-3
+        train_env_cfg.reward_config.scales.pose = -1.0
+        train_env_cfg.reward_config.scales.tracking_ang_vel = 1.25
+        train_env_cfg.reward_config.scales.tracking_lin_vel = 1.25
+        train_env_cfg.reward_config.scales.feet_phase = 1.0
+        train_env_cfg.reward_config.scales.ang_vel_xy = -0.3
+        train_env_cfg.reward_config.scales.orientation = -5.0
+
+    if is_humanoid_task and not use_push_randomization:
+        train_env_cfg.push_config.enable = False
+        train_env_cfg.push_config.magnitude_range = [0.0, 0.0]
+    randomizer = (
+        registry.get_domain_randomizer(env_name) if use_domain_randomization else None
+    )
+    raw_env = registry.load(env_name, config=train_env_cfg)
+    train_env = RandomizeInitialWrapper(
+        raw_env,
+        num_envs,
+        seed,
+        train_env_cfg.episode_length,
+        train_env_cfg.action_repeat,
+        randomization_fn=randomizer,
+        device_rank=device_rank,
+    )
+
+    # Make evaluation environment
+    eval_env_cfg = registry.get_default_config(env_name)
+    if is_humanoid_task and not use_push_randomization:
+        eval_env_cfg.push_config.enable = False
+        eval_env_cfg.push_config.magnitude_range = [0.0, 0.0]
+    eval_env = registry.load(env_name, config=eval_env_cfg)
+    eval_env = PlaygroundEvalEnvWrapper(
+        eval_env, eval_env_cfg.episode_length, env_name, num_eval_envs, seed
+    )
+
+    return train_env, eval_env