|
|
from typing import Callable, List, Type |
|
|
import gymnasium as gym |
|
|
import numpy as np |
|
|
from mani_skill.envs.sapien_env import BaseEnv |
|
|
from mani_skill.utils import common, gym_utils |
|
|
import argparse |
|
|
import yaml |
|
|
import torch |
|
|
from collections import deque |
|
|
from PIL import Image |
|
|
import cv2 |
|
|
import imageio |
|
|
from functools import partial |
|
|
|
|
|
from diffusion_policy.workspace.robotworkspace import RobotWorkspace |
|
|
|
|
|
def parse_args(args=None): |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument("-e", "--env-id", type=str, default="PickCube-v1", help=f"Environment to run motion planning solver on. ") |
|
|
parser.add_argument("-o", "--obs-mode", type=str, default="rgb", help="Observation mode to use. Usually this is kept as 'none' as observations are not necesary to be stored, they can be replayed later via the mani_skill.trajectory.replay_trajectory script.") |
|
|
parser.add_argument("-n", "--num-traj", type=int, default=25, help="Number of trajectories to generate.") |
|
|
parser.add_argument("--only-count-success", action="store_true", help="If true, generates trajectories until num_traj of them are successful and only saves the successful trajectories/videos") |
|
|
parser.add_argument("--reward-mode", type=str) |
|
|
parser.add_argument("-b", "--sim-backend", type=str, default="auto", help="Which simulation backend to use. Can be 'auto', 'cpu', 'gpu'") |
|
|
parser.add_argument("--render-mode", type=str, default="rgb_array", help="can be 'sensors' or 'rgb_array' which only affect what is saved to videos") |
|
|
parser.add_argument("--vis", action="store_true", help="whether or not to open a GUI to visualize the solution live") |
|
|
parser.add_argument("--save-video", action="store_true", help="whether or not to save videos locally") |
|
|
parser.add_argument("--traj-name", type=str, help="The name of the trajectory .h5 file that will be created.") |
|
|
parser.add_argument("--shader", default="default", type=str, help="Change shader used for rendering. Default is 'default' which is very fast. Can also be 'rt' for ray tracing and generating photo-realistic renders. Can also be 'rt-fast' for a faster but lower quality ray-traced renderer") |
|
|
parser.add_argument("--record-dir", type=str, default="demos", help="where to save the recorded trajectories") |
|
|
parser.add_argument("--num-procs", type=int, default=1, help="Number of processes to use to help parallelize the trajectory replay process. This uses CPU multiprocessing and only works with the CPU simulation backend at the moment.") |
|
|
parser.add_argument("--random_seed", type=int, default=0, help="Random seed for the environment.") |
|
|
parser.add_argument("--pretrained_path", type=str, default=None, help="Random seed for the environment.") |
|
|
|
|
|
return parser.parse_args() |
|
|
|
|
|
task2lang = { |
|
|
"PegInsertionSide-v1": "Pick up a orange-white peg and insert the orange end into the box with a hole in it.", |
|
|
"PickCube-v1": "Grasp a red cube and move it to a target goal position.", |
|
|
"StackCube-v1": "Pick up a red cube and stack it on top of a green cube and let go of the cube without it falling.", |
|
|
"PlugCharger-v1": "Pick up one of the misplaced shapes on the board/kit and insert it into the correct empty slot.", |
|
|
"PushCube-v1": "Push and move a cube to a goal region in front of it." |
|
|
} |
|
|
import random |
|
|
import os |
|
|
|
|
|
args = parse_args() |
|
|
seed = args.random_seed |
|
|
random.seed(seed) |
|
|
os.environ['PYTHONHASHSEED'] = str(seed) |
|
|
np.random.seed(seed) |
|
|
torch.manual_seed(seed) |
|
|
torch.cuda.manual_seed(seed) |
|
|
torch.backends.cudnn.deterministic = True |
|
|
torch.backends.cudnn.benchmark = False |
|
|
|
|
|
env_id = args.env_id |
|
|
env = gym.make( |
|
|
env_id, |
|
|
obs_mode=args.obs_mode, |
|
|
control_mode="pd_joint_pos", |
|
|
render_mode=args.render_mode, |
|
|
reward_mode="dense" if args.reward_mode is None else args.reward_mode, |
|
|
sensor_configs=dict(shader_pack=args.shader), |
|
|
human_render_camera_configs=dict(shader_pack=args.shader), |
|
|
viewer_camera_configs=dict(shader_pack=args.shader), |
|
|
sim_backend=args.sim_backend |
|
|
) |
|
|
|
|
|
from diffusion_policy.workspace.robotworkspace import RobotWorkspace |
|
|
import hydra |
|
|
import dill |
|
|
|
|
|
checkpoint_path = args.pretrained_path |
|
|
print(f"Loading policy from {checkpoint_path}. Task is {task2lang[env_id]}") |
|
|
|
|
|
def get_policy(output_dir, device): |
|
|
|
|
|
|
|
|
payload = torch.load(open(checkpoint_path, 'rb'), pickle_module=dill) |
|
|
cfg = payload['cfg'] |
|
|
cls = hydra.utils.get_class(cfg._target_) |
|
|
workspace = cls(cfg, output_dir=output_dir) |
|
|
workspace: RobotWorkspace |
|
|
workspace.load_payload(payload, exclude_keys=None, include_keys=None) |
|
|
|
|
|
|
|
|
policy = workspace.model |
|
|
if cfg.training.use_ema: |
|
|
policy = workspace.ema_model |
|
|
|
|
|
device = torch.device(device) |
|
|
policy.to(device) |
|
|
policy.eval() |
|
|
|
|
|
return policy |
|
|
|
|
|
policy = get_policy('./', device = 'cuda') |
|
|
MAX_EPISODE_STEPS = 400 |
|
|
total_episodes = args.num_traj |
|
|
success_count = 0 |
|
|
base_seed = 20241201 |
|
|
instr = task2lang[env_id] |
|
|
import tqdm |
|
|
|
|
|
DATA_STAT = {'state_min': [-0.7463043928146362, -0.0801204964518547, -0.4976441562175751, -2.657780647277832, -0.5742632150650024, 1.8309762477874756, -2.2423808574676514, 0.0, 0.0], 'state_max': [0.7645499110221863, 1.4967026710510254, 0.4650936424732208, -0.3866899907588959, 0.5505855679512024, 3.2900545597076416, 2.5737812519073486, 0.03999999910593033, 0.03999999910593033], 'action_min': [-0.7472005486488342, -0.08631071448326111, -0.4995281398296356, -2.658363103866577, -0.5751323103904724, 1.8290787935256958, -2.245187997817993, -1.0], 'action_max': [0.7654682397842407, 1.4984270334243774, 0.46786263585090637, -0.38181185722351074, 0.5517147779464722, 3.291581630706787, 2.575840711593628, 1.0], 'action_std': [0.2199309915304184, 0.18780815601348877, 0.13044124841690063, 0.30669933557510376, 0.1340624988079071, 0.24968451261520386, 0.9589747190475464, 0.9827960729598999], 'action_mean': [-0.00885344110429287, 0.5523102879524231, -0.007564723491668701, -2.0108158588409424, 0.004714342765510082, 2.615924596786499, 0.08461848646402359, -0.19301606714725494]} |
|
|
|
|
|
state_min = torch.tensor(DATA_STAT['state_min']).cuda() |
|
|
state_max = torch.tensor(DATA_STAT['state_max']).cuda() |
|
|
action_min = torch.tensor(DATA_STAT['action_min']).cuda() |
|
|
action_max = torch.tensor(DATA_STAT['action_max']).cuda() |
|
|
|
|
|
for episode in tqdm.trange(total_episodes): |
|
|
obs_window = deque(maxlen=2) |
|
|
obs, _ = env.reset(seed = episode + base_seed) |
|
|
|
|
|
img = env.render().cuda().float() |
|
|
proprio = obs['agent']['qpos'][:].cuda() |
|
|
proprio = (proprio - state_min) / (state_max - state_min) * 2 - 1 |
|
|
obs_window.append({ |
|
|
'agent_pos': proprio, |
|
|
"head_cam": img.permute(0, 3, 1, 2), |
|
|
}) |
|
|
obs_window.append({ |
|
|
'agent_pos': proprio, |
|
|
"head_cam": img.permute(0, 3, 1, 2), |
|
|
}) |
|
|
|
|
|
global_steps = 0 |
|
|
video_frames = [] |
|
|
|
|
|
success_time = 0 |
|
|
done = False |
|
|
|
|
|
while global_steps < MAX_EPISODE_STEPS and not done: |
|
|
obs = obs_window[-1] |
|
|
actions = policy.predict_action(obs) |
|
|
actions = actions['action_pred'].squeeze(0) |
|
|
actions = (actions + 1) / 2 * (action_max - action_min) + action_min |
|
|
actions = actions.detach().cpu().numpy() |
|
|
actions = actions[:8] |
|
|
for idx in range(actions.shape[0]): |
|
|
action = actions[idx] |
|
|
obs, reward, terminated, truncated, info = env.step(action) |
|
|
img = env.render().cuda().float() |
|
|
proprio = obs['agent']['qpos'][:].cuda() |
|
|
proprio = (proprio - state_min) / (state_max - state_min) * 2 - 1 |
|
|
obs_window.append({ |
|
|
'agent_pos': proprio, |
|
|
"head_cam": img.permute(0, 3, 1, 2), |
|
|
}) |
|
|
video_frames.append(env.render().squeeze(0).detach().cpu().numpy()) |
|
|
global_steps += 1 |
|
|
if terminated or truncated: |
|
|
assert "success" in info, sorted(info.keys()) |
|
|
if info['success']: |
|
|
done = True |
|
|
success_count += 1 |
|
|
break |
|
|
print(f"Trial {episode+1} finished, success: {info['success']}, steps: {global_steps}") |
|
|
|
|
|
success_rate = success_count / total_episodes * 100 |
|
|
print(f"Tested {total_episodes} episodes, success rate: {success_rate:.2f}%") |
|
|
log_file = f"results_dp_{checkpoint_path.split('/')[-1].split('.')[0]}.txt" |
|
|
with open(log_file, 'a') as f: |
|
|
f.write(f"{args.env_id}:{seed}:{success_count}\n") |
