yfzhoucs/TinyLanguageRobots

1from tiny_ur5 import TinyUR5Env2import yaml3from initializer import Initializer4import torch5import numpy as np6import skimage7import clip8from skimage import img_as_ubyte9import json10from recorder import NumpyEncoder11import cv212 13 14class TaskSuccess:15    def __init__(self, env: TinyUR5Env, task) -> None:16        self.env = env17        self.task = task18 19        self.target = task['target']20        self.target_init_pos = env.get_pos_xy(self.target)21        self.target_init_orientation = env.get_pos_orientation(self.target)22 23        self.task = task['action']24    25    def success(self):26        if self.task == 'push_forward':27            if self.env.get_pos_xy(self.target)[1] - self.target_init_pos[1] > 20:28                return True29        elif self.task == 'push_backward':30            if self.env.get_pos_xy(self.target)[1] - self.target_init_pos[1] < -20:31                return True32        elif self.task == 'push_left':33            if self.env.get_pos_xy(self.target)[0] - self.target_init_pos[0] < -20:34                return True35        elif self.task == 'push_right':36            if self.env.get_pos_xy(self.target)[0] - self.target_init_pos[0] > 20:37                return True38        elif self.task == 'rotate_clock':39            # if self.env.get_pos_orientation(self.target) - self.target_init_orientation > 0.3:40            if self.env.get_pos_orientation(self.target) - self.target_init_orientation < -0.3:41                return True42        elif self.task == 'rotate_counterclock':43            # if self.env.get_pos_orientation(self.target) - self.target_init_orientation < -0.3:44            if self.env.get_pos_orientation(self.target) - self.target_init_orientation > 0.3:45                return True46        47        return False48 49 50 51 52class ModelTester:53    def __init__(self, yaml_file, model, model_forward_fn, device, method, show_cv2, show_human, time_upper_bound=500) -> None:54        self.yaml_file = yaml_file55        self.model = model56        self.model_forward_fn = model_forward_fn57        self.device = device58        self.time_upper_bound = time_upper_bound59        self.method = method60        self.show_cv2 = show_cv261        self.show_human = show_human62    63 64    def test_1_rollout(self, test_id):65        with open(self.yaml_file, "r") as stream:66            try:67                config = yaml.safe_load(stream)68                # print(config, type(config))69            except yaml.YAMLError as exc:70                print(exc)71        72        initializer = Initializer(config)73 74        config, task = initializer.get_config_and_task()75        sentence = initializer.get_sentence()76        print(sentence)77 78        env = TinyUR5Env(render_mode='human', config=config)79        if self.show_human:80            env.render()81        img = env.render('rgb_array')82        if self.show_cv2:83            cv2.imshow('tiny_ur5', cv2.cvtColor(img, cv2.COLOR_RGB2BGR))84 85        task_success_judge = TaskSuccess(env, task)86 87        time_step = 088        while time_step < self.time_upper_bound:89            actions = self.model_forward_fn(env, self.model, sentence, self.method, self.device)90            # for i in range(actions.shape[-1]):91            for i in range(60):92                action = actions[:, i]93                observation, reward, done, info = env.step(action, eef_z=80)94                # env.render()95                img = env.render('rgb_array')96                if self.show_human:97                    env.render()98                if self.show_cv2:99                    cv2.imshow('tiny_ur5', cv2.cvtColor(img, cv2.COLOR_RGB2BGR))100                    cv2.waitKey(1)101                time_step += 1102 103                success = task_success_judge.success()104                if success:105                    env.close()106                    return True, task107 108            print(time_step)109        env.close()110        return False, task111    112    def test(self, test_num:int, name):113        tasks_states = []114        for i in range(test_num):115            success, task = self.test_1_rollout(i)116            task['success'] = success117            tasks_states.append(task)118        119        with open(f'results_{name}_{test_num}.json', 'w') as f:120            json.dump(tasks_states, f, cls=NumpyEncoder, indent=4)121 122 123def model_forward_fn(env, model, sentence, method, device):124    img = env.render('rgb_array')125    img = img[::-1, :, :3]126    img = skimage.transform.resize(img, (224, 224))127    img = img_as_ubyte(img) / 255128    # skimage.io.imsave('tmp.png', img_as_ubyte(img))129    # img = skimage.io.imread('tmp.png')[::-1,:,:3] / 255130    img = torch.tensor(img, dtype=torch.float32).unsqueeze(0).to(device)131    sentence = clip.tokenize([sentence]).to(device)132 133    def _joints_to_sin_cos_(joints):134        sin_cos_joints = [0] * 8135        for i in range(len(joints)):136            sin_cos_joints[i * 2] = np.sin(joints[i])137            sin_cos_joints[i * 2 + 1] = np.cos(joints[i])138        return sin_cos_joints139    140    def _sin_cos_to_joint_(sin, cos):141        angle = np.arctan(sin / cos)142        if cos < 0:143            if sin > 0:144                angle = angle + np.pi145            else:146                angle = angle - np.pi147        return angle148 149    def _sin_cos_to_joints_(sin_cos):150        joints = [0] * 4151        for i in range(len(joints)):152            joints[i] = _sin_cos_to_joint_(sin_cos[i * 2], sin_cos[i * 2 + 1])153        return joints154    155    def _seq_sin_cos_to_joint_(sin_cos_seq):156        joints = []157        for i in range(sin_cos_seq.shape[-1]):158            action = _sin_cos_to_joints_(sin_cos_seq[:, i])159            joints.append(action)160        joints = np.transpose(np.array(joints))161        return joints162 163    if method == 'bcz':164        phis = torch.tensor(np.linspace(0.0, 1.0, 60, dtype=np.float32)) \165            .unsqueeze(0).unsqueeze(0).repeat(1, 8, 1).to(device)166        action = model(img, sentence, phis)167        # return joints_trajectory_pred[0].detach().cpu().numpy()168    elif method == 'ours':169 170        171        joint_angles = torch.tensor(_joints_to_sin_cos_(env.robot_joints)).unsqueeze(0).to(device)172        phis = torch.tensor(np.linspace(0.0, 1.0, 60, dtype=np.float32)) \173            .unsqueeze(0).unsqueeze(0).repeat(1, 8, 1).to(device)174        stage = 3175        target_position_pred, ee_pos_pred, \176            displacement_pred, attn_map, attn_map2, \177            attn_map3, attn_map4, action = \178            model(img, joint_angles, sentence, phis, stage)179        180    action = action.detach().cpu().numpy()[0]181    action = _seq_sin_cos_to_joint_(action)182    return action183 184 185def load_model(ckpt, method, device):186    if method == 'bcz':187        from models.film_model import Backbone188        # model = Backbone(img_size=224, num_traces_out=4, embedding_size=256, num_weight_points=10, input_nc=3, device=device)189        model = Backbone(img_size=224, num_traces_out=8, embedding_size=256, num_weight_points=12, input_nc=3, device=device)190        model.load_state_dict(torch.load(ckpt, map_location=device)['model'], strict=True)191        # model = model.cpu()192        model = model.to(device)193        return model194    elif method == 'ours':195        from models.backbone_rgbd_sub_attn_tinyur5 import Backbone196        model = Backbone(img_size=224, embedding_size=256, num_traces_out=2, num_joints=8, num_weight_points=12, input_nc=3, device=device)197        model.load_state_dict(torch.load(ckpt, map_location=device)['model'], strict=True)198        model = model.to(device)199        return model200 201 202def calculate_success_rate(filename):203    results = json.load(open(filename))204 205    success = 0206    for i in range(len(results)):207        if results[i]['success'] == True:208            success += 1209    print(success / len(results))210 211if __name__ == '__main__':212    device = torch.device('cpu')213 214    # # # BCZ215    # method = 'bcz'216    # # ckpt = '/share/yzhou298/ckpts/tinyur5/train-baseline-bcz-film-resnet-huberloss/200000.pth'217    # ckpt = '/share/yzhou298/ckpts/tinyur5/train-baseline-bcz-film-resnet-huberloss-2-larger-dataset-corrected-rotation/200000.pth'218    # # ckpt = '/share/yzhou298/ckpts/tinyur5/train-baseline-bcz-film-resnet-huberloss-long-inst/60000.pth'219    220    # # # Ours221    # # method = 'ours'222    # # # ckpt = '/share/yzhou298/ckpts/tinyur5/train-tinyur5-rgb-sub-attn-range/90000.pth'223    # # # ckpt = '/share/yzhou298/ckpts/tinyur5/train-tinyur5-rgb-sub-attn-range-larger-dataset/120000.pth'224    # # # ckpt = '/share/yzhou298/ckpts/tinyur5/train-tinyur5-rgb-sub-attn-range-larger-dataset/340000.pth'225    # # ckpt = '/share/yzhou298/ckpts/tinyur5/train-tinyur5-rgb-sub-attn-range-larger-dataset-corrected-rotation/310000.pth'226    227    228    # model = load_model(ckpt, method, device)229    # modeltester = ModelTester('config.yaml', model, model_forward_fn, device, method=method, show_cv2=False, show_human=False)230    # # modeltester.test_1_rollout(0)231    # modeltester.test(100, method+'_correct_rotation_200000')232 233    # calculate_success_rate('results_ours_100.json')234    # calculate_success_rate('results_bcz_100.json')235    # calculate_success_rate('results_bcz_correct_rotation_100.json')236    calculate_success_rate('results_bcz_correct_rotation_200000_100.json')