Initial Public Release
This commit is contained in:
Younggyo Seo
@@ -0,0 +1,128 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import gymnasium as gym
|
||||
|
||||
import humanoid_bench
|
||||
from gymnasium.wrappers import TimeLimit
|
||||
from stable_baselines3.common.vec_env import SubprocVecEnv
|
||||
import numpy as np
|
||||
import torch
|
||||
from loguru import logger as log
|
||||
|
||||
# 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():
|
||||
import humanoid_bench
|
||||
|
||||
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
|
||||
@@ -0,0 +1,82 @@
|
||||
from typing import Optional
|
||||
|
||||
import gymnasium as gym
|
||||
import torch
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
app_launcher = AppLauncher(headless=True)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
import isaaclab_tasks
|
||||
from isaaclab_tasks.utils.parse_cfg import parse_env_cfg
|
||||
|
||||
|
||||
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"
|
||||
)
|
||||
@@ -0,0 +1,136 @@
|
||||
from mujoco_playground import registry
|
||||
from mujoco_playground import wrapper_torch
|
||||
|
||||
import jax
|
||||
import mujoco
|
||||
|
||||
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, None
|
||||
|
||||
def render_trajectory(self, trajectory):
|
||||
scene_option = mujoco.MjvOption()
|
||||
scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = False
|
||||
scene_option.flags[mujoco.mjtVisFlag.mjVIS_PERTFORCE] = False
|
||||
scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = False
|
||||
|
||||
frames = self.env.render(
|
||||
trajectory,
|
||||
camera="track" if "Joystick" in self.env_name else None,
|
||||
height=480,
|
||||
width=640,
|
||||
scene_option=scene_option,
|
||||
)
|
||||
return frames
|
||||
|
||||
|
||||
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)
|
||||
if use_tuned_reward:
|
||||
# NOTE: Tuned reward for G1. Used for producing Figure 7 in the paper.
|
||||
assert env_name in ["G1JoystickRoughTerrain", "G1JoystickFlatTerrain"]
|
||||
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
|
||||
|
||||
is_humanoid_task = env_name in [
|
||||
"G1JoystickRoughTerrain",
|
||||
"G1JoystickFlatTerrain",
|
||||
"T1JoystickRoughTerrain",
|
||||
"T1JoystickFlatTerrain",
|
||||
]
|
||||
|
||||
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 = wrapper_torch.RSLRLBraxWrapper(
|
||||
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
|
||||
)
|
||||
|
||||
render_env_cfg = registry.get_default_config(env_name)
|
||||
if is_humanoid_task and not use_push_randomization:
|
||||
render_env_cfg.push_config.enable = False
|
||||
render_env_cfg.push_config.magnitude_range = [0.0, 0.0]
|
||||
render_env = registry.load(env_name, config=render_env_cfg)
|
||||
render_env = PlaygroundEvalEnvWrapper(
|
||||
render_env, render_env_cfg.episode_length, env_name, 1, seed
|
||||
)
|
||||
|
||||
return train_env, eval_env, render_env
|
||||
Reference in New Issue
Block a user