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 torchvision.transforms.functional import center_crop 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="Path to the pretrained model") 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 from transformers import AutoModelForVision2Seq, AutoProcessor DATA_STAT = {'mean': [ 0.00263866, 0.01804881, -0.02151551, -0.00384866, 0.00500441, -0.00057146, -0.26013601], 'std': [0.06639539, 0.1246438 , 0.09675793, 0.03351422, 0.04930534, 0.25787726, 0.96762997], 'max': [0.31303197, 0.77948809, 0.42906255, 0.20186238, 0.63990456, 0.99999917, 1. ], 'min': [-0.31464151, -0.64183694, -0.62718982, -0.5888508 , -0.97813392, -0.99999928, -1. ], 'q01': [-0.18656027, -0.31995443, -0.24702898, -0.18005923, -0.2164692 , -0.82366071, -1. ], 'q99': [0.18384692, 0.45547636, 0.27452313, 0.03571117, 0.1188747 , 0.85074112, 1. ]} MODEL_PATH = args.pretrained_path def make_policy(): device = torch.device('cuda') processor = AutoProcessor.from_pretrained(MODEL_PATH, trust_remote_code=True) vla = AutoModelForVision2Seq.from_pretrained( MODEL_PATH, attn_implementation="flash_attention_2", # [Optional] Requires `flash_attn` torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, trust_remote_code=True ).to(device) vla.norm_stats["maniskill"] = { "action": { "min": np.array(DATA_STAT["min"]), "max": np.array(DATA_STAT["max"]), "mean": np.array(DATA_STAT["mean"]), "std": np.array(DATA_STAT["std"]), "q01": np.array(DATA_STAT["q01"]), "q99": np.array(DATA_STAT["q99"]), } } vla = vla.eval() return vla, processor vla, processor = make_policy() success_counts = {} for env_id in task2lang.keys(): env = gym.make( env_id, obs_mode=args.obs_mode, control_mode="pd_ee_delta_pose", 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 ) MAX_EPISODE_STEPS = 400 total_episodes = args.num_traj success_count = 0 base_seed = 20241201 import tqdm for episode in tqdm.trange(total_episodes): obs_window = deque(maxlen=2) obs, _ = env.reset(seed = base_seed + episode) img = env.render().squeeze(0).detach().cpu().numpy() obs_window.append(img) global_steps = 0 video_frames = [] success_time = 0 done = False while global_steps < MAX_EPISODE_STEPS and not done: obs = obs_window[-1] image_arrs = [ obs_window[-1] ] images = [Image.fromarray(arr) for arr in image_arrs] original_size = images[0].size crop_scale = 0.9 sqrt_crop_scale = crop_scale sqrt_crop_scale = np.sqrt(crop_scale) images = [ center_crop( img, output_size=( int(sqrt_crop_scale * img.size[1]), int(sqrt_crop_scale * img.size[0]) ) ) for img in images ] images = [img.resize(original_size, Image.Resampling.BILINEAR) for img in images] # de-capitalize and remove trailing period instruction = task2lang[env_id].lower() prompt = f"In: What action should the robot take to {instruction}?\nOut:" inputs = processor(prompt, images).to("cuda:0", dtype=torch.bfloat16) actions = vla.predict_action(**inputs, unnorm_key="maniskill", do_sample=False)[None] for idx in range(actions.shape[0]): action = actions[idx] # print(action) # action = action * (np.array(DATA_STAT['std']) + 1e-8) + np.array(DATA_STAT['mean']) obs, reward, terminated, truncated, info = env.step(action) img = env.render().squeeze(0).detach().cpu().numpy() obs_window.append(img) video_frames.append(img) global_steps += 1 if terminated or truncated: assert "success" in info, sorted(info.keys()) if info['success']: success_count += 1 done = True break print(f"Trial {episode+1} finished, success: {info['success']}, steps: {global_steps}") success_counts[env_id] = success_count print(f"Task {env_id} finished, success: {success_count}/{total_episodes}") log_file = "results_ovla_all.log" with open(log_file, 'a') as f: f.write(f"{seed}:{success_counts}\n")