0001-enable-openvino-inference-for-eval.patch

From 002e1ce362b74d40e242f8f44388808ede687128 Mon Sep 17 00:00:00 2001
From: "Liu, Baihe" <baihe.liu@intel.com>
Date: Fri, 28 Feb 2025 02:14:27 +0800
Subject: [PATCH 1/2] enable openvino inference for eval

---
 constants.py        |   3 +-
 imitate_episodes.py | 111 +++++++++++++++++++++++++++++++++++---------
 sim_env.py          |   2 +
 3 files changed, 93 insertions(+), 23 deletions(-)

diff --git a/constants.py b/constants.py
index f445350..61926a2 100644
--- a/constants.py
+++ b/constants.py
@@ -21,7 +21,8 @@ SIM_TASK_CONFIGS = {
         'dataset_dir': DATA_DIR + '/sim_insertion_scripted',
         'num_episodes': 50,
         'episode_len': 400,
-        'camera_names': ['top']
+        'camera_names': ['top', 'angle', 'left_wrist', 'right_wrist']
+        # 'camera_names': ['top']
     },
 
     'sim_insertion_human': {
diff --git a/imitate_episodes.py b/imitate_episodes.py
index 34f9a37..3fcc29b 100644
--- a/imitate_episodes.py
+++ b/imitate_episodes.py
@@ -7,6 +7,8 @@ import matplotlib.pyplot as plt
 from copy import deepcopy
 from tqdm import tqdm
 from einops import rearrange
+import time
+import openvino as ov
 
 from constants import DT
 from constants import PUPPET_GRIPPER_JOINT_OPEN
@@ -32,6 +34,8 @@ def main(args):
     batch_size_train = args['batch_size']
     batch_size_val = args['batch_size']
     num_epochs = args['num_epochs']
+    device = args['device']
+    print_time = args['print_time']
 
     # get task parameters
     is_sim = task_name[:4] == 'sim_'
@@ -85,14 +89,16 @@ def main(args):
         'seed': args['seed'],
         'temporal_agg': args['temporal_agg'],
         'camera_names': camera_names,
-        'real_robot': not is_sim
+        'real_robot': not is_sim,
+        'device': device,
+        'print_time': print_time
     }
 
     if is_eval:
         ckpt_names = [f'policy_best.ckpt']
         results = []
         for ckpt_name in ckpt_names:
-            success_rate, avg_return = eval_bc(config, ckpt_name, save_episode=True)
+            success_rate, avg_return = eval_bc(config, ckpt_name, save_episode=False)
             results.append([ckpt_name, success_rate, avg_return])
 
         for ckpt_name, success_rate, avg_return in results:
@@ -144,10 +150,17 @@ def get_image(ts, camera_names):
         curr_image = rearrange(ts.observation['images'][cam_name], 'h w c -> c h w')
         curr_images.append(curr_image)
     curr_image = np.stack(curr_images, axis=0)
-    curr_image = torch.from_numpy(curr_image / 255.0).float().cuda().unsqueeze(0)
+    curr_image = torch.from_numpy(curr_image / 255.0).float().cpu().unsqueeze(0)
     return curr_image
 
 
+def get_ov_model(model_path, device='CPU'):
+    core = ov.Core()
+    ov_model = core.read_model(model_path)
+    compiled_model = ov.compile_model(ov_model, device)
+    return compiled_model
+
+
 def eval_bc(config, ckpt_name, save_episode=True):
     set_seed(1000)
     ckpt_dir = config['ckpt_dir']
@@ -161,14 +174,17 @@ def eval_bc(config, ckpt_name, save_episode=True):
     task_name = config['task_name']
     temporal_agg = config['temporal_agg']
     onscreen_cam = 'angle'
+    device = config['device'].upper()
+    print_time = config['print_time']
 
     # load policy and stats
     ckpt_path = os.path.join(ckpt_dir, ckpt_name)
-    policy = make_policy(policy_class, policy_config)
-    loading_status = policy.load_state_dict(torch.load(ckpt_path))
-    print(loading_status)
-    policy.cuda()
-    policy.eval()
+    
+    if ckpt_path.endswith(".ckpt"):
+        ckpt_path = ckpt_path.replace('.ckpt', '.xml')
+    if not os.path.exists(ckpt_path):
+        raise FileNotFoundError("OpenVINO Model .xml is not found, please convert .ckpt to .xml through ov_convert.py")
+    policy = get_ov_model(ckpt_path, device=device)
     print(f'Loaded: {ckpt_path}')
     stats_path = os.path.join(ckpt_dir, f'dataset_stats.pkl')
     with open(stats_path, 'rb') as f:
@@ -179,9 +195,23 @@ def eval_bc(config, ckpt_name, save_episode=True):
 
     # load environment
     if real_robot:
-        from aloha_scripts.robot_utils import move_grippers # requires aloha
-        from aloha_scripts.real_env import make_real_env # requires aloha
-        env = make_real_env(init_node=True)
+        from aloha.robot_utils import move_grippers # requires aloha
+        from aloha.real_env import make_real_env # requires aloha
+        from interbotix_common_modules.common_robot.robot import (
+            create_interbotix_global_node,
+            get_interbotix_global_node,
+            robot_startup,
+        )
+        from interbotix_common_modules.common_robot.exceptions import InterbotixException
+        try:
+            node = get_interbotix_global_node()
+        except:
+            node = create_interbotix_global_node('aloha')
+        env = make_real_env(node=node, setup_base=False)
+        try:
+            robot_startup(node)
+        except InterbotixException:
+            pass
         env_max_reward = 0
     else:
         from sim_env import make_sim_env
@@ -195,7 +225,7 @@ def eval_bc(config, ckpt_name, save_episode=True):
 
     max_timesteps = int(max_timesteps * 1) # may increase for real-world tasks
 
-    num_rollouts = 50
+    num_rollouts = 10
     episode_returns = []
     highest_rewards = []
     for rollout_id in range(num_rollouts):
@@ -216,22 +246,30 @@ def eval_bc(config, ckpt_name, save_episode=True):
 
         ### evaluation loop
         if temporal_agg:
-            all_time_actions = torch.zeros([max_timesteps, max_timesteps+num_queries, state_dim]).cuda()
-
-        qpos_history = torch.zeros((1, max_timesteps, state_dim)).cuda()
+            all_time_actions = torch.zeros([max_timesteps, max_timesteps+num_queries, state_dim]).cpu()
+        qpos_history = np.zeros((max_timesteps, state_dim))
         image_list = [] # for visualization
         qpos_list = []
         target_qpos_list = []
         rewards = []
+
         with torch.inference_mode():
+            latencies_all = [] 
+            time0 = time.time()
             for t in range(max_timesteps):
+                latencies = []
+
+                onscreen_time = time.time()
+
                 ### update onscreen render and wait for DT
                 if onscreen_render:
                     image = env._physics.render(height=480, width=640, camera_id=onscreen_cam)
                     plt_img.set_data(image)
                     plt.pause(DT)
+                latencies.append(time.time()-onscreen_time)  # onscreen render time
 
                 ### process previous timestep to get qpos and image_list
+                pre_time = time.time()
                 obs = ts.observation
                 if 'images' in obs:
                     image_list.append(obs['images'])
@@ -239,15 +277,25 @@ def eval_bc(config, ckpt_name, save_episode=True):
                     image_list.append({'main': obs['image']})
                 qpos_numpy = np.array(obs['qpos'])
                 qpos = pre_process(qpos_numpy)
-                qpos = torch.from_numpy(qpos).float().cuda().unsqueeze(0)
-                qpos_history[:, t] = qpos
+                qpos = torch.from_numpy(qpos).float().cpu().unsqueeze(0)
+
+                # qpos_history[:, t] = qpos
                 curr_image = get_image(ts, camera_names)
+                latencies.append(time.time()-pre_time)  # preprocess time
 
                 ### query policy
+                query_time = time.time()
                 if config['policy_class'] == "ACT":
+                    model_time = time.time()
                     if t % query_frequency == 0:
-                        all_actions = policy(qpos, curr_image)
+                        input = {
+                            'qpos': qpos,
+                            'tensor.1': curr_image,
+                        }
+                        all_actions = policy(input)[0]  # np array
+                    latencies.append(time.time()-model_time)  # model infer time
                     if temporal_agg:
+                        all_actions = torch.from_numpy(all_actions)
                         all_time_actions[[t], t:t+num_queries] = all_actions
                         actions_for_curr_step = all_time_actions[:, t]
                         actions_populated = torch.all(actions_for_curr_step != 0, axis=1)
@@ -255,7 +303,7 @@ def eval_bc(config, ckpt_name, save_episode=True):
                         k = 0.01
                         exp_weights = np.exp(-k * np.arange(len(actions_for_curr_step)))
                         exp_weights = exp_weights / exp_weights.sum()
-                        exp_weights = torch.from_numpy(exp_weights).cuda().unsqueeze(dim=1)
+                        exp_weights = torch.from_numpy(exp_weights).to(qpos.device).unsqueeze(dim=1)
                         raw_action = (actions_for_curr_step * exp_weights).sum(dim=0, keepdim=True)
                     else:
                         raw_action = all_actions[:, t % query_frequency]
@@ -263,23 +311,40 @@ def eval_bc(config, ckpt_name, save_episode=True):
                     raw_action = policy(qpos, curr_image)
                 else:
                     raise NotImplementedError
+                latencies.append(time.time()-query_time)  # query policy time
 
                 ### post-process actions
-                raw_action = raw_action.squeeze(0).cpu().numpy()
+                post_time = time.time()
+                if not temporal_agg:
+                    raw_action = raw_action.squeeze(0)  # numpy array
+                else:
+                    raw_action = raw_action.squeeze(0).cpu().numpy()
                 action = post_process(raw_action)
                 target_qpos = action
+                latencies.append(time.time()-post_time)
 
                 ### step the environment
+                env_time = time.time()
                 ts = env.step(target_qpos)
+                latencies.append(time.time()-env_time)
 
                 ### for visualization
                 qpos_list.append(qpos_numpy)
                 target_qpos_list.append(target_qpos)
                 rewards.append(ts.reward)
-
+                if print_time:
+                    print(f'screen render:{latencies[0]:.9f}s, process image:{(latencies[1]):.9f}s, model inference:{latencies[2]:.9f}, query policy:{(latencies[3]):.9f}s, post process:{(latencies[4]):.9f}, env:{(latencies[5]):.9f}')
+                latencies_all.append(latencies)
+
+            print(f'Avg fps {max_timesteps / (time.time() - time0)}')
+            if print_time:
+                latencies_all = np.array(latencies_all)
+                average_latency = np.mean(latencies_all[1:], axis=0)
+                print(f'==================Rollout {rollout_id} Avg Latency:==================\n \
+                    screen render:{average_latency[0]:.9f}s, process image:{average_latency[1]:.9f}s, model inference:{average_latency[2]:.9f}, query policy:{average_latency[3]:.9f}s, post process:{average_latency[4]:.9f}, env:{average_latency[5]:.9f}')
             plt.close()
         if real_robot:
-            move_grippers([env.puppet_bot_left, env.puppet_bot_right], [PUPPET_GRIPPER_JOINT_OPEN] * 2, move_time=0.5)  # open
+            move_grippers([env.puppet_bot_left, env.puppet_bot_right], [PUPPET_GRIPPER_JOINT_OPEN] * 2, moving_time=0.5)  # open
             pass
 
         rewards = np.array(rewards)
@@ -424,6 +489,8 @@ if __name__ == '__main__':
     parser.add_argument('--seed', action='store', type=int, help='seed', required=True)
     parser.add_argument('--num_epochs', action='store', type=int, help='num_epochs', required=True)
     parser.add_argument('--lr', action='store', type=float, help='lr', required=True)
+    parser.add_argument('--device', action='store', type=str, help='device:CPU/GPU', required=True)
+    parser.add_argument('--print_time', action='store_true', help='print time log in eval', required=False)
 
     # for ACT
     parser.add_argument('--kl_weight', action='store', type=int, help='KL Weight', required=False)
diff --git a/sim_env.py b/sim_env.py
index b79b935..857261c 100644
--- a/sim_env.py
+++ b/sim_env.py
@@ -108,6 +108,8 @@ class BimanualViperXTask(base.Task):
         obs['env_state'] = self.get_env_state(physics)
         obs['images'] = dict()
         obs['images']['top'] = physics.render(height=480, width=640, camera_id='top')
+        obs['images']['left_wrist'] = physics.render(height=480, width=640, camera_id='left_wrist')
+        obs['images']['right_wrist'] = physics.render(height=480, width=640, camera_id='right_wrist')
         obs['images']['angle'] = physics.render(height=480, width=640, camera_id='angle')
         obs['images']['vis'] = physics.render(height=480, width=640, camera_id='front_close')
 
-- 
2.34.1