Hugging Face's logo

Timsty
/
groot_deployment

ONNX
Model card Files
xet
 Community
groot_deployment / Isaac-GR00T /gr00t /eval /sim /SimplerEnv /simpler_env.py
Timsty's picture
Timsty
Add files using upload-large-folder tool
b6a9d87 verified
Raw History Blame Contribute Delete
11.9 kB
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import cv2
import gymnasium as gym
from gymnasium.envs.registration import register
import numpy as np
import simpler_env
from simpler_env.utils.env.observation_utils import get_image_from_maniskill2_obs_dict
from transforms3d import euler as te, quaternions as tq
class GoogleFractalEnv(gym.Env):
def __init__(self, env_name: str, image_size: tuple[int, int]):
env = simpler_env.make(env_name)
env._max_episode_steps = 10000
self.env = env
agent_space = env.observation_space["agent"]
print("[SimplerEnv] agent space keys:", list(agent_space.spaces.keys()))
# assert False
obs_low = env.observation_space["agent"]["eef_pos"].low
obs_high = env.observation_space["agent"]["eef_pos"].high
self.observation_space = gym.spaces.Dict(
{
"video.image": gym.spaces.Box(
low=0,
high=255,
shape=(image_size[0], image_size[1], 3),
dtype=np.uint8,
),
"state.x": gym.spaces.Box(low=obs_low[0], high=obs_high[0], shape=(1,)),
"state.y": gym.spaces.Box(low=obs_low[1], high=obs_high[1], shape=(1,)),
"state.z": gym.spaces.Box(low=obs_low[2], high=obs_high[2], shape=(1,)),
"state.rx": gym.spaces.Box(low=obs_low[3], high=obs_high[3], shape=(1,)),
"state.ry": gym.spaces.Box(low=obs_low[4], high=obs_high[4], shape=(1,)),
"state.rz": gym.spaces.Box(low=obs_low[5], high=obs_high[5], shape=(1,)),
"state.rw": gym.spaces.Box(low=obs_low[6], high=obs_high[6], shape=(1,)),
"state.gripper": gym.spaces.Box(low=obs_low[7], high=obs_high[7], shape=(1,)),
"annotation.human.action.task_description": gym.spaces.Text(max_length=512),
}
)
action_low = env.action_space.low
action_high = env.action_space.high
self.action_space = gym.spaces.Dict(
{
"action.x": gym.spaces.Box(low=action_low[0], high=action_high[0], shape=(1,)),
"action.y": gym.spaces.Box(low=action_low[1], high=action_high[1], shape=(1,)),
"action.z": gym.spaces.Box(low=action_low[2], high=action_high[2], shape=(1,)),
"action.roll": gym.spaces.Box(low=action_low[3], high=action_high[3], shape=(1,)),
"action.pitch": gym.spaces.Box(low=action_low[4], high=action_high[4], shape=(1,)),
"action.yaw": gym.spaces.Box(low=action_low[5], high=action_high[5], shape=(1,)),
"action.gripper": gym.spaces.Box(
low=action_low[6], high=action_high[6], shape=(1,)
),
}
)
self.image_size = image_size
self.previous_gripper_action = None
self.sticky_action_is_on = False
self.sticky_gripper_action = 0.0
self.gripper_action_repeat = 0
self.sticky_gripper_num_repeat = 15
def reset(self, seed=None, options=None):
self.previous_gripper_action = None
self.sticky_action_is_on = False
self.sticky_gripper_action = 0.0
self.gripper_action_repeat = 0
observation, info = self.env.reset(seed=int(seed) if seed is not None else None)
observation = self._process_observation(observation)
info["success"] = False
return observation, info
def step(self, action):
action_vector = np.concatenate(
[
action["action.x"],
action["action.y"],
action["action.z"],
action["action.roll"],
action["action.pitch"],
action["action.yaw"],
self._postprocess_gripper(action["action.gripper"]),
],
axis=0,
)
observation, reward, done, truncated, info = self.env.step(action_vector)
observation = self._process_observation(observation)
info["success"] = bool(info.get("success", False))
return observation, reward, done, truncated, info
def _process_observation(self, obs):
img = get_image_from_maniskill2_obs_dict(self.env, obs)
proprio = obs["agent"]["eef_pos"]
qunat_xyzw = np.roll(proprio[3:7], -1)
gripper_closedness = 1 - proprio[7]
return {
"video.image": cv2.resize(img, (self.image_size[1], self.image_size[0])),
"state.x": [proprio[0]],
"state.y": [proprio[1]],
"state.z": [proprio[2]],
"state.rx": [qunat_xyzw[0]],
"state.ry": [qunat_xyzw[1]],
"state.rz": [qunat_xyzw[2]],
"state.rw": [qunat_xyzw[3]],
"state.gripper": [gripper_closedness],
"annotation.human.action.task_description": self.env.unwrapped.get_language_instruction(),
}
def _postprocess_gripper(self, current_gripper_action: float) -> float:
current_gripper_action = (current_gripper_action * 2) - 1 # [0,1] -> [-1,1]
relative_gripper_action = -current_gripper_action
if np.abs(relative_gripper_action) > 0.5 and self.sticky_action_is_on is False:
self.sticky_action_is_on = True
self.sticky_gripper_action = relative_gripper_action
if self.sticky_action_is_on:
self.gripper_action_repeat += 1
relative_gripper_action = self.sticky_gripper_action
if self.gripper_action_repeat == self.sticky_gripper_num_repeat:
self.sticky_action_is_on = False
self.gripper_action_repeat = 0
self.sticky_gripper_action = 0.0
return relative_gripper_action
class WidowXBridgeEnv(gym.Env):
def __init__(self, env_name: str, image_size: tuple[int, int]):
env = simpler_env.make(env_name)
env._max_episode_steps = 10000
self.env = env
obs_low = env.observation_space["agent"]["eef_pos"].low
obs_high = env.observation_space["agent"]["eef_pos"].high
self.observation_space = gym.spaces.Dict(
{
"video.image_0": gym.spaces.Box(
low=0,
high=255,
shape=(image_size[0], image_size[1], 3),
dtype=np.uint8,
),
"state.x": gym.spaces.Box(low=obs_low[0], high=obs_high[0], shape=(1,)),
"state.y": gym.spaces.Box(low=obs_low[1], high=obs_high[1], shape=(1,)),
"state.z": gym.spaces.Box(low=obs_low[2], high=obs_high[2], shape=(1,)),
"state.roll": gym.spaces.Box(low=obs_low[3], high=obs_high[3], shape=(1,)),
"state.pitch": gym.spaces.Box(low=obs_low[4], high=obs_high[4], shape=(1,)),
"state.yaw": gym.spaces.Box(low=obs_low[5], high=obs_high[5], shape=(1,)),
"state.pad": gym.spaces.Box(low=obs_low[6], high=obs_high[6], shape=(1,)),
"state.gripper": gym.spaces.Box(low=obs_low[7], high=obs_high[7], shape=(1,)),
"annotation.human.action.task_description": gym.spaces.Text(max_length=512),
}
)
action_low = env.action_space.low
action_high = env.action_space.high
self.action_space = gym.spaces.Dict(
{
"action.x": gym.spaces.Box(low=action_low[0], high=action_high[0], shape=(1,)),
"action.y": gym.spaces.Box(low=action_low[1], high=action_high[1], shape=(1,)),
"action.z": gym.spaces.Box(low=action_low[2], high=action_high[2], shape=(1,)),
"action.roll": gym.spaces.Box(low=action_low[3], high=action_high[3], shape=(1,)),
"action.pitch": gym.spaces.Box(low=action_low[4], high=action_high[4], shape=(1,)),
"action.yaw": gym.spaces.Box(low=action_low[5], high=action_high[5], shape=(1,)),
"action.gripper": gym.spaces.Box(
low=action_low[6], high=action_high[6], shape=(1,)
),
}
)
self.image_size = image_size
# Bridge orientation adjustment
self.default_rot = np.array([[0, 0, 1.0], [0, 1.0, 0], [-1.0, 0, 0]])
def reset(self, seed=None, options=None):
observation, info = self.env.reset(seed=int(seed) if seed is not None else None)
observation = self._process_observation(observation)
info["success"] = False
return observation, info
def step(self, action):
action_vector = np.concatenate(
[
action["action.x"],
action["action.y"],
action["action.z"],
action["action.roll"],
action["action.pitch"],
action["action.yaw"],
self._postprocess_gripper(action["action.gripper"]),
],
axis=0,
)
observation, reward, done, truncated, info = self.env.step(action_vector)
observation = self._process_observation(observation)
info["success"] = bool(info.get("success", False))
return observation, reward, done, truncated, info
def _process_observation(self, obs):
img = get_image_from_maniskill2_obs_dict(self.env, obs)
proprio = obs["agent"]["eef_pos"]
rm_bridge = tq.quat2mat(proprio[3:7])
rpy_bridge_converted = te.mat2euler(rm_bridge @ self.default_rot.T)
return {
"video.image_0": cv2.resize(img, (self.image_size[1], self.image_size[0])),
"state.x": [proprio[0]],
"state.y": [proprio[1]],
"state.z": [proprio[2]],
"state.roll": [rpy_bridge_converted[0]],
"state.pitch": [rpy_bridge_converted[1]],
"state.yaw": [rpy_bridge_converted[2]],
"state.pad": [0],
"state.gripper": [proprio[7]],
"annotation.human.action.task_description": self.env.unwrapped.get_language_instruction(),
}
def _postprocess_gripper(self, action):
# trained with [0, 1], 0 close, 1 open -> convert to SimplerEnv [-1, 1]
return 2.0 * (action > 0.5) - 1.0
def register_simpler_envs():
# Google/Fractal
for env_name in [
"google_robot_pick_coke_can",
"google_robot_pick_object",
"google_robot_move_near",
"google_robot_open_drawer",
"google_robot_close_drawer",
"google_robot_place_in_closed_drawer",
]:
register(
id=f"simpler_env_google/{env_name}",
entry_point="gr00t.eval.sim.SimplerEnv.simpler_env:GoogleFractalEnv",
kwargs={"env_name": env_name, "image_size": (256, 320)},
)
# WidowX/Bridge
for env_name in [
"widowx_spoon_on_towel",
"widowx_carrot_on_plate",
"widowx_stack_cube",
"widowx_put_eggplant_in_basket",
"widowx_put_eggplant_in_sink",
"widowx_open_drawer",
"widowx_close_drawer",
]:
register(
id=f"simpler_env_widowx/{env_name}",
entry_point="gr00t.eval.sim.SimplerEnv.simpler_env:WidowXBridgeEnv",
kwargs={"env_name": env_name, "image_size": (256, 256)},
)