| | import rclpy |
| | from sensor_msgs.msg import JointState |
| | import numpy as np |
| | import time |
| | import h5py |
| | from real_robot_env import AnubisRobotEnv |
| | import cv2 |
| | from tqdm import tqdm |
| | import os |
| | from scripts.agilex_model import create_model |
| | import random |
| | from PIL import Image, ImageDraw, ImageFont |
| | import yaml |
| | import torch |
| |
|
| | class RDTInferenceRobotEnv(AnubisRobotEnv): |
| | def __init__(self, hz=20, max_timestep=500, task_name='', num_rollout=1, model_name='1st'): |
| | self.model_name = model_name |
| | self.checkpoint = f'/home/rllab/workspace/jellyho/RoboticsDiffusionTransformer/checkpoints/{task_name}' |
| | super().__init__(hz=hz, max_timestep=max_timestep, task_name=task_name, num_rollout=num_rollout) |
| | self.last_front_img = None |
| | self.right_wrist_img = None |
| | self.left_wrist_img = None |
| |
|
| | def bringup_model(self): |
| | with open('configs/base.yaml', "r") as fp: |
| | config = yaml.safe_load(fp) |
| | self.model = create_model( |
| | args=config, |
| | dtype=torch.bfloat16, |
| | pretrained=self.checkpoint, |
| | pretrained_text_encoder_name_or_path="google/t5-v1_1-xxl", |
| | pretrained_vision_encoder_name_or_path="google/siglip-so400m-patch14-384", |
| | control_frequency=20 |
| | ) |
| | print('model loaded') |
| |
|
| | def inference(self): |
| | proprio = self.obs['eef_pose'] |
| | if self.action_counter % 64 == 0: |
| | front_img = self.obs['agentview_image'] |
| | right_wrist_img = self.obs['rightview_image'] |
| | left_wrist_img = self.obs['leftview_image'] |
| | image_arrs = [ |
| | self.last_front_img, |
| | self.last_right_wrist_img, |
| | self.last_left_wrist_img, |
| | front_img, |
| | right_wrist_img, |
| | left_wrist_img |
| | ] |
| | images = [Image.fromarray(arr) if arr is not None else None for arr in image_arrs] |
| | proprio = torch.tensor(proprio).unsqueeze(0) |
| | with torch.inference_mode(): |
| | self.actions = self.model.step( |
| | proprio=proprio, |
| | images=images, |
| | instruction=self.instruction |
| | ).squeeze(0).cpu().numpy() |
| | idx = self.action_counter % 64 |
| | act = self.actions[idx] |
| | self.action_counter += 1 |
| | self.last_front_img = self.obs['agentview_image'] |
| | self.last_right_wrist_img = self.obs['rightview_image'] |
| | self.last_left_wrist_img = self.obs['leftview_image'] |
| | self.send_action(act) |
| |
|
| | if __name__ == '__main__': |
| | task_name = 'anubis_towel_kirby' |
| | rollout_num = 20 |
| | hz = 20 |
| | model_name = 'rdt' |
| | |
| | node = RDTInferenceRobotEnv( |
| | hz=hz, |
| | max_timestep=800, |
| | task_name=task_name, |
| | num_rollout=rollout_num, |
| | model_name=model_name |
| | ) |
| | while node.rollout_counter < rollout_num: |
| | try: |
| | img = cv2.cvtColor(node.obs['agentview_image'], cv2.COLOR_RGB2BGR) |
| | if node.start: |
| | node.window.show(img, text=node.instruction) |
| | else: |
| | node.window.show(img, overlay_img=node.overlay_img, text=node.instruction) |
| | node.last_front_img = node.obs['agentview_image'] |
| | node.last_right_wrist_img = node.obs['rightview_image'] |
| | node.last_left_wrist_img = node.obs['leftview_image'] |
| | except KeyboardInterrupt: |
| | node.ros_close() |
| | |
| | except Exception as e: |
| | print(f"An error occurred: {e}") |
| | node.ros_close() |
