|
| 1 | +import os |
| 2 | + |
| 3 | +os.environ["TOKENIZERS_PARALLELISM"] = "false" |
| 4 | + |
| 5 | +import collections |
| 6 | +import copy |
| 7 | +import dataclasses |
| 8 | +import os.path as osp |
| 9 | +import time |
| 10 | +from datetime import datetime |
| 11 | +from pathlib import Path |
| 12 | + |
| 13 | +import cv2 |
| 14 | +import deoxys.utils.transform_utils as dft |
| 15 | +import imageio |
| 16 | +import numpy as np |
| 17 | +import torch |
| 18 | +import tqdm |
| 19 | +import tyro |
| 20 | +from deoxys import config_root |
| 21 | +from deoxys.experimental.motion_utils import reset_joints_to |
| 22 | +from deoxys.franka_interface import FrankaInterface |
| 23 | +from deoxys.utils import YamlConfig |
| 24 | +from PIL import Image |
| 25 | +from realsense_camera import MultiCamera |
| 26 | +from transformers import AutoModel, AutoProcessor |
| 27 | + |
| 28 | +# Add your camera serial numbers here |
| 29 | +EGO_CAMERA = "213522070137" |
| 30 | +THIRD_CAMERA = "243222074139" |
| 31 | + |
| 32 | +reset_joint_positions = [ |
| 33 | + 0.0760389047913384, |
| 34 | + -1.0362613022620384, |
| 35 | + -0.054254247684777324, |
| 36 | + -2.383951857286591, |
| 37 | + -0.004505598470154735, |
| 38 | + 1.3820559157131187, |
| 39 | + 0.784935455988679, |
| 40 | +] |
| 41 | + |
| 42 | + |
| 43 | +def save_rollout_video(rollout_images, save_dir): |
| 44 | + """Saves an MP4 replay of an episode.""" |
| 45 | + date_time = time.strftime("%Y_%m_%d-%H_%M_%S") |
| 46 | + mp4_path = Path(save_dir) / f"{date_time}.mp4" |
| 47 | + video_writer = imageio.get_writer(mp4_path, fps=5) |
| 48 | + for img in rollout_images: |
| 49 | + video_writer.append_data(img) |
| 50 | + video_writer.close() |
| 51 | + print(f"Saved rollout MP4 at path {mp4_path}") |
| 52 | + |
| 53 | + |
| 54 | +@dataclasses.dataclass |
| 55 | +class Args: |
| 56 | + ################################################################################################################# |
| 57 | + # Model parameters |
| 58 | + ################################################################################################################# |
| 59 | + resize_size: int = 224 |
| 60 | + replan_steps: int = 5 |
| 61 | + model_path: str = "" |
| 62 | + repo_id: str = "" |
| 63 | + task: str = "" |
| 64 | + |
| 65 | + ################################################################################################################# |
| 66 | + # Utils |
| 67 | + ################################################################################################################# |
| 68 | + video_out_path: Path = Path("experiments/7_franka/videos") # Path to save videos |
| 69 | + max_timesteps: int = 300 # Number of timesteps to run |
| 70 | + |
| 71 | + |
| 72 | +def convert_gripper_action(action): |
| 73 | + action[-1] = 1 - action[-1] |
| 74 | + if action[-1] < 0.5: |
| 75 | + action[-1] = -1 |
| 76 | + |
| 77 | + return action |
| 78 | + |
| 79 | + |
| 80 | +def get_robot_interface(): |
| 81 | + robot_interface = FrankaInterface(osp.join(config_root, "charmander.yml")) |
| 82 | + controller_cfg = YamlConfig(osp.join(config_root, "osc-pose-controller.yml")).as_easydict() |
| 83 | + controller_type = "OSC_POSE" |
| 84 | + |
| 85 | + return robot_interface, controller_cfg, controller_type |
| 86 | + |
| 87 | + |
| 88 | +def main(args: Args): |
| 89 | + multi_camera = MultiCamera() |
| 90 | + robot_interface, controller_cfg, controller_type = get_robot_interface() |
| 91 | + |
| 92 | + model = ( |
| 93 | + AutoModel.from_pretrained(args.model_path, dtype=torch.bfloat16, trust_remote_code=True).eval().cuda() |
| 94 | + ) |
| 95 | + |
| 96 | + processor = AutoProcessor.from_pretrained(args.model_path, trust_remote_code=True) |
| 97 | + |
| 98 | + while True: |
| 99 | + action_plan = collections.deque() |
| 100 | + input("Press Enter to start episode ...") |
| 101 | + reset_joints_to(robot_interface, reset_joint_positions) |
| 102 | + |
| 103 | + replay_images = [] |
| 104 | + bar = tqdm.tqdm( |
| 105 | + range(args.max_timesteps), |
| 106 | + position=0, |
| 107 | + leave=True, |
| 108 | + ncols=80, |
| 109 | + desc="Rollout steps", |
| 110 | + ) |
| 111 | + |
| 112 | + for _ in bar: |
| 113 | + try: |
| 114 | + images = multi_camera.get_frame() |
| 115 | + last_state = robot_interface._state_buffer[-1] |
| 116 | + last_gripper_state = robot_interface._gripper_state_buffer[-1] |
| 117 | + frame, _ = images[THIRD_CAMERA] |
| 118 | + ego_frame, _ = images[EGO_CAMERA] |
| 119 | + |
| 120 | + if not action_plan: |
| 121 | + frame = copy.deepcopy(frame) |
| 122 | + ego_frame = copy.deepcopy(ego_frame) |
| 123 | + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| 124 | + ego_frame = cv2.cvtColor(ego_frame, cv2.COLOR_BGR2RGB) |
| 125 | + replay_images.append(frame) |
| 126 | + frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| 127 | + ego_frame_rgb = cv2.cvtColor(ego_frame, cv2.COLOR_BGR2RGB) |
| 128 | + replay_images.append(frame_rgb.copy()) |
| 129 | + |
| 130 | + eef_pose = np.asarray(last_state.O_T_EE, dtype=np.float32).reshape(4, 4).T |
| 131 | + eef_state = np.concatenate( |
| 132 | + ( |
| 133 | + eef_pose[:3, -1], |
| 134 | + dft.mat2euler( |
| 135 | + eef_pose[:3, :-1], |
| 136 | + ), |
| 137 | + ), |
| 138 | + axis=-1, |
| 139 | + ) |
| 140 | + gripper_state = np.array([last_gripper_state.width]) |
| 141 | + state_data = np.concatenate([eef_state.flatten(), np.array([0]), gripper_state]) |
| 142 | + |
| 143 | + frame_resized = cv2.resize(frame_rgb, (args.resize_size, args.resize_size)) |
| 144 | + ego_frame_resized = cv2.resize(ego_frame_rgb, (args.resize_size, args.resize_size)) |
| 145 | + |
| 146 | + ego_frame = Image.fromarray(ego_frame_resized) |
| 147 | + frame = Image.fromarray(frame_resized) |
| 148 | + |
| 149 | + # NOTE: Change the keys to match your model |
| 150 | + batch = { |
| 151 | + "observation.images.image": [frame], |
| 152 | + "observation.images.wrist_image": [ego_frame], |
| 153 | + "observation.state": [state_data], |
| 154 | + "task": [args.task], |
| 155 | + "repo_id": [args.repo_id], |
| 156 | + } |
| 157 | + ov_out = processor.select_action( |
| 158 | + model, |
| 159 | + batch, |
| 160 | + ) |
| 161 | + action_chunk = ov_out.action[0].numpy() |
| 162 | + assert len(action_chunk) >= args.replan_steps, ( |
| 163 | + f"We want to replan every {args.replan_steps} steps, but policy only predicts {len(action_chunk)} steps." |
| 164 | + ) |
| 165 | + action_plan.extend(action_chunk[: args.replan_steps]) |
| 166 | + |
| 167 | + pred_action_chunk = action_plan.popleft() |
| 168 | + action = pred_action_chunk |
| 169 | + rotation_matrix = dft.euler2mat(action[3:6]) |
| 170 | + quat = dft.mat2quat(rotation_matrix) |
| 171 | + axis_angle = dft.quat2axisangle(quat) |
| 172 | + action[3:6] = axis_angle |
| 173 | + action = convert_gripper_action(action) |
| 174 | + |
| 175 | + robot_interface.control( |
| 176 | + controller_type=controller_type, action=action, controller_cfg=controller_cfg |
| 177 | + ) |
| 178 | + |
| 179 | + except KeyboardInterrupt: |
| 180 | + break |
| 181 | + |
| 182 | + # saving video |
| 183 | + video_save_path = args.video_out_path / args.task.replace(" ", "_") |
| 184 | + video_save_path.mkdir(parents=True, exist_ok=True) |
| 185 | + curr_time = datetime.now().strftime("%Y_%m_%d_%H_%M_%S") |
| 186 | + save_path = video_save_path / f"{curr_time}.mp4" |
| 187 | + video = np.stack(replay_images) |
| 188 | + imageio.mimsave(save_path, video, fps=20) |
| 189 | + |
| 190 | + if input("Do one more eval (default y)? [y/n]").lower() == "n": |
| 191 | + break |
| 192 | + |
| 193 | + |
| 194 | +if __name__ == "__main__": |
| 195 | + args = tyro.cli(Args) |
| 196 | + main(args) |
0 commit comments