File size: 8,453 Bytes
ac29381 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 | """
Example script for running 10 rollouts of a DROID policy on the example environment.
Usage:
First, make sure you download the simulation assets and unpack them into the root directory of this package.
Then, in a separate terminal, launch the policy server on localhost:8000
-- make sure to set XLA_PYTHON_CLIENT_MEM_FRACTION to avoid JAX hogging all the GPU memory.
For example, to launch a pi0-FAST-DROID policy (with joint position control),
run the command below in a separate terminal from the openpi "karl/droid_policies" branch:
XLA_PYTHON_CLIENT_MEM_FRACTION=0.5 uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_fast_droid_jointpos --policy.dir=s3://openpi-assets-simeval/pi0_fast_droid_jointpos
Finally, run the evaluation script:
python run_eval.py --episodes 10 --headless
"""
import uuid
import tyro
import argparse
import gymnasium as gym
import torch
import cv2
import mediapy
import numpy as np
from datetime import datetime
from pathlib import Path
from PIL import Image
from tqdm import tqdm
from openpi_client import image_tools
from sim_evals.inference.abstract_client import InferenceClient
from policy_client import WebsocketClientPolicy
class DreamZeroJointPosClient(InferenceClient):
def __init__(self,
remote_host:str = "localhost",
remote_port:int = 6000,
open_loop_horizon:int = 8,
) -> None:
self.client = WebsocketClientPolicy(remote_host, remote_port)
self.open_loop_horizon = open_loop_horizon
self.actions_from_chunk_completed = 0
self.pred_action_chunk = None
self.session_id = str(uuid.uuid4())
def visualize(self, request: dict):
"""
Return the camera views how the model sees it
"""
curr_obs = self._extract_observation(request)
right_img = image_tools.resize_with_pad(curr_obs["right_image"], 224, 224)
wrist_img = image_tools.resize_with_pad(curr_obs["wrist_image"], 224, 224)
left_img = image_tools.resize_with_pad(curr_obs["left_image"], 224, 224)
combined = np.concatenate([right_img, wrist_img, left_img], axis=1)
return combined
def reset(self):
self.actions_from_chunk_completed = 0
self.pred_action_chunk = None
self.session_id = str(uuid.uuid4())
def infer(self, obs: dict, instruction: str) -> dict:
"""
Infer the next action from the policy in a server-client setup
"""
curr_obs = self._extract_observation(obs)
if (
self.actions_from_chunk_completed == 0
or self.actions_from_chunk_completed >= self.open_loop_horizon
):
self.actions_from_chunk_completed = 0
request_data = {
"observation/exterior_image_0_left": image_tools.resize_with_pad(curr_obs["right_image"], 180, 320),
"observation/exterior_image_1_left": image_tools.resize_with_pad(curr_obs["left_image"], 180, 320),
"observation/wrist_image_left": image_tools.resize_with_pad(curr_obs["wrist_image"], 180, 320),
"observation/joint_position": curr_obs["joint_position"].astype(np.float64),
"observation/cartesian_position": np.zeros((6,), dtype=np.float64), # dummy cartesian position
"observation/gripper_position": curr_obs["gripper_position"].astype(np.float64),
"prompt": instruction,
"session_id": self.session_id,
}
for k, v in request_data.items():
print(f"{k}: {v.shape if not isinstance(v, str) else v}")
result = self.client.infer(request_data)
actions = result["actions"] if isinstance(result, dict) else result
assert len(actions.shape) == 2, f"Expected 2D array, got shape {actions.shape}"
assert actions.shape[-1] == 8, f"Expected 8 action dimensions (7 joints + 1 gripper), got {actions.shape[-1]}"
self.pred_action_chunk = actions
action = self.pred_action_chunk[self.actions_from_chunk_completed]
self.actions_from_chunk_completed += 1
# binarize gripper action
if action[-1].item() > 0.5:
action = np.concatenate([action[:-1], np.ones((1,))])
else:
action = np.concatenate([action[:-1], np.zeros((1,))])
img1 = image_tools.resize_with_pad(curr_obs["right_image"], 224, 224)
img2 = image_tools.resize_with_pad(curr_obs["wrist_image"], 224, 224)
img3 = image_tools.resize_with_pad(curr_obs["left_image"], 224, 224)
both = np.concatenate([img1, img2, img3], axis=1)
return {"action": action, "viz": both}
def _extract_observation(self, obs_dict, *, save_to_disk=False):
# Assign images
right_image = obs_dict["policy"]["external_cam"][0].clone().detach().cpu().numpy()
left_image = obs_dict["policy"]["external_cam_2"][0].clone().detach().cpu().numpy()
wrist_image = obs_dict["policy"]["wrist_cam"][0].clone().detach().cpu().numpy()
# Capture proprioceptive state
robot_state = obs_dict["policy"]
joint_position = robot_state["arm_joint_pos"].clone().detach().cpu().numpy()
gripper_position = robot_state["gripper_pos"].clone().detach().cpu().numpy()
if save_to_disk:
combined_image = np.concatenate([right_image, wrist_image], axis=1)
combined_image = Image.fromarray(combined_image)
combined_image.save("robot_camera_views.png")
return {
"right_image": right_image,
"left_image": left_image,
"wrist_image": wrist_image,
"joint_position": joint_position,
"gripper_position": gripper_position,
}
def main(
episodes: int = 10,
scene: int = 1,
headless: bool = True,
host: str = "localhost",
port: int = 6000,
):
# launch omniverse app with arguments (inside function to prevent overriding tyro)
from isaaclab.app import AppLauncher
parser = argparse.ArgumentParser(description="Tutorial on creating an empty stage.")
AppLauncher.add_app_launcher_args(parser)
args_cli, _ = parser.parse_known_args()
args_cli.enable_cameras = True
args_cli.headless = headless
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app
# All IsaacLab dependent modules should be imported after the app is launched
import sim_evals.environments # noqa: F401
from isaaclab_tasks.utils import parse_env_cfg
# Initialize the env
env_cfg = parse_env_cfg(
"DROID",
device=args_cli.device,
num_envs=1,
use_fabric=True,
)
instruction = None
match scene:
case 1:
instruction = "put the cube in the bowl"
case 2:
instruction = "pick up the can and put it in the mug"
case 3:
instruction = "put the banana in the bin"
case _:
raise ValueError(f"Scene {scene} not supported")
env_cfg.set_scene(scene)
env = gym.make("DROID", cfg=env_cfg)
obs, _ = env.reset()
obs, _ = env.reset() # need second render cycle to get correctly loaded materials
client = DreamZeroJointPosClient(remote_host=host, remote_port=port)
video_dir = Path("runs") / datetime.now().strftime("%Y-%m-%d") / datetime.now().strftime("%H-%M-%S")
video_dir.mkdir(parents=True, exist_ok=True)
video = []
ep = 0
max_steps = env.env.max_episode_length
with torch.no_grad():
for ep in range(episodes):
for _ in tqdm(range(max_steps), desc=f"Episode {ep+1}/{episodes}"):
ret = client.infer(obs, instruction)
if not headless:
cv2.imshow("Right Camera", cv2.cvtColor(ret["viz"], cv2.COLOR_RGB2BGR))
cv2.waitKey(1)
video.append(ret["viz"])
action = torch.tensor(ret["action"])[None]
obs, _, term, trunc, _ = env.step(action)
if term or trunc:
break
client.reset()
mediapy.write_video(
video_dir / f"episode_{ep}.mp4",
video,
fps=15,
)
video = []
env.close()
simulation_app.close()
if __name__ == "__main__":
args = tyro.cli(main)
|