content: import torch
  from torch import nn
  import torch.nn.functional as F
  import numpy as np
  import h5py
  import pathlib
  import os
  import argparse
  import matplotlib.pyplot as plt
  from PIL import Image
  import torchvision
  from torchvision import transforms
  # from visualize_episodes import visualize_joints
  from utils import set_seed, sample_box_pose
  # from imitate_episodes import get_image
  from sim_env import BOX_POSE
  from constants import DT
  from imitate_episodes import save_videos
  from einops import rearrange
  import time
  
  DT = 0.02
  import IPython
  e = IPython.embed
  
  # modified from https://github.com/jyopari/VINN/blob/main/nearest-neighbor-eval/handle_nn.ipynb
  
  def calculate_nearest_neighbors(curr_feature, support_inputs, support_targets, k, state_weight):
      has_skip = len(support_targets.shape) == 3
      if has_skip: # when there is action skip
          num_targets, skip, a_dim = support_targets.shape
          support_targets = support_targets.view((num_targets, -1))
  
      curr_vis_feature, curr_s_feature = curr_feature
      support_vis_feature, support_s_feature = support_inputs
  
      pairwise_dist_vis = torch.norm(curr_vis_feature - support_vis_feature, dim=1).unsqueeze(0)
      pairwise_dist_s = torch.norm(curr_s_feature - support_s_feature, dim=1).unsqueeze(0)
      pairwise_dist = pairwise_dist_vis + pairwise_dist_s * state_weight
  
      sorted_dist, index = torch.sort(pairwise_dist, dim=1) # sort the support axis
      permuted_support_targets = support_targets[index]
      topk_dist = pairwise_dist[:, :k]
      topk_support_targets = permuted_support_targets[:, :k]
      weights = F.softmax(-topk_dist, dim=1)
      weighted_support_targets = weights.unsqueeze(2) * topk_support_targets
      prediction = torch.sum(weighted_support_targets, dim=1)
  
      if has_skip:
          num_predictions = prediction.shape[0]
          prediction = prediction.reshape((num_predictions, skip, a_dim))
  
      return prediction
  
  
  def main(args):
      # TODO ######################
      k = None # for scripted box transfer
      skip = 100
      real_robot = True
      save_episode = True
      # TODO ######################
      onscreen_cam = 'main'
      state_dim = 14
      dataset_dir = args['dataset_dir']
      onscreen_render = args['onscreen_render']
      ckpt_dir = args['ckpt_dir']
      model_dir = args['model_dir']
      task_name = args['task_name']
  
      if 'insertion' in task_name:
          sim_episode_len = 400
          env_max_reward = 4
          ks = [None]
      elif 'transfer_cube' in task_name:
          sim_episode_len = 400
          env_max_reward = 4
          ks = [1, 1, 1]
          if 'human' in dataset_dir:
              state_weight = 5
          else:
              state_weight = 10
          print(f'{state_weight=}')
      elif task_name == 'ziploc_slide':
          env_max_reward = 1
          ks = [71]
          state_weight = 0
      elif task_name == 'aloha_mobile_wipe_wine':
          sim_episode_len = 1300
          env_max_reward = 4
          ks = [2, 2, 2]
          state_weight = 5
          print(f'{state_weight=}')
      else:
          raise NotImplementedError
  
      model_name = pathlib.PurePath(model_dir).name
      seed = int(model_name.split('-')[-1][:-3])
  
      repr_type = 'byol'
      if 'cotrain' in model_name:
          repr_type += '_cotrain'
      e() # make sure!
  
      k = ks[seed]
  
      if real_robot:
          BASE_DELAY = 15
          query_freq = skip - BASE_DELAY
  
      # load train data
      vis_features = []
      state_features = []
      Y = []
      for episode_id in range(0, 40):
          dataset_path = os.path.join(dataset_dir, f'episode_{episode_id}.hdf5')
          with h5py.File(dataset_path, 'r') as root:
              action = root['/action'][:]
              base_action = root['/base_action'][:]
              action = np.concatenate([action, base_action], axis=1)
              camera_names = list(root[f'/observations/images/'].keys())
  
          # Visual feature
          all_cam_feature = []
          for cam_name in camera_names:
              feature_dataset_path = os.path.join(dataset_dir, f'{repr_type}_features_seed{seed}_episode_{episode_id}.hdf5')
              with h5py.File(feature_dataset_path, 'r') as root:
                  cam_feature = root[f'/features/{cam_name}'][:]
                  all_cam_feature.append(cam_feature)
          vis_fea = np.concatenate(all_cam_feature, axis=1)
  
          ## State feature
          dataset_path = os.path.join(dataset_dir, f'episode_{episode_id}.hdf5')
          with h5py.File(dataset_path, 'r') as root:
              s_fea = root['/observations/qpos'][:]
  
          # stack actions together
          eps_len = len(action)
          indices = np.tile(np.arange(skip), eps_len).reshape(eps_len, skip) # each row is 0, 1, ... skip
          offset = np.expand_dims(np.arange(eps_len), axis=1)
          indices = indices + offset # row1: 0, 1, ... skip; row2: 1, 2, ... skip+1
          # indices will exceed eps_len, thus clamp to eps_len-1
          indices = np.clip(indices, 0, eps_len-1)
          # stack action
          action = action[indices] # new shape: eps_len, skip, a_dim
  
          vis_features.append(vis_fea)
          state_features.append(s_fea)
          Y.append(action)
  
      vis_features = np.concatenate(vis_features)
      state_features  = np.concatenate(state_features)
      Y = np.concatenate(Y)
      train_inputs = [torch.from_numpy(vis_features).cuda(), torch.from_numpy(state_features).cuda()]
      train_targets = torch.from_numpy(Y).cuda()
  
      set_seed(1000)
      feature_extractors = {}
      for cam_name in camera_names:
          resnet = torchvision.models.resnet18(pretrained=True)
          loading_status = resnet.load_state_dict(torch.load(model_dir.replace('DUMMY', cam_name)))
          print(cam_name, loading_status)
          resnet = nn.Sequential(*list(resnet.children())[:-1])
          resnet = resnet.cuda()
          resnet.eval()
          feature_extractors[cam_name] = resnet
  
  
  
      # load environment
      if real_robot:
          from aloha_scripts.real_env import make_real_env #### TODO TODO
          env = make_real_env(init_node=True, setup_robots=True, setup_base=True)
          max_timesteps = sim_episode_len
          camera_names = ['cam_high', 'cam_left_wrist', 'cam_right_wrist']
      else:
          from sim_env import make_sim_env
          env = make_sim_env(task_name)
          max_timesteps = sim_episode_len
  
  
      num_rollouts = 50
      episode_returns = []
      max_rewards = []
      for rollout_id in range(num_rollouts):
          ### set task
          BOX_POSE[0] = sample_box_pose() # used in sim reset
          ts = env.reset()
  
          ### evaluation loop
          qpos_history = torch.zeros((1, max_timesteps, state_dim)).cuda()
          image_list = [] # for visualization
          qpos_list = []
          target_qpos_list = []
          rewards = []
          with torch.inference_mode():
              for t in range(sim_episode_len):
                  start_time = time.time()
                  if t % 100 == 0: print(t)
                  if t % query_freq == 0:
                      ### process previous timestep to get qpos and image_list
                      obs = ts.observation
                      if 'images' in obs:
                          image_list.append(obs['images'])
                      else:
                          image_list.append({'main': obs['image']})
                      qpos_numpy = np.array(obs['qpos'])
                      # qpos = pre_process(qpos_numpy)
                      qpos = torch.from_numpy(qpos_numpy).float().cuda().unsqueeze(0)
                      qpos_history[:, t] = qpos
                      _, curr_image_raw = get_image(ts, camera_names)
  
                      image_size = 120
                      transform = transforms.Compose([
                          transforms.Resize(image_size),  # will scale the image
                          transforms.CenterCrop(image_size),
                          transforms.ToTensor(),
                          transforms.Lambda(expand_greyscale),
                          transforms.Normalize(
                              mean=torch.tensor([0.485, 0.456, 0.406]),
                              std=torch.tensor([0.229, 0.224, 0.225])),
                      ])
  
                      all_cam_features = []
                      for cam_id, curr_image in enumerate(curr_image_raw):
                          curr_image = Image.fromarray(curr_image) # TODO only one camera
                          curr_image = transform(curr_image)
                          curr_image = curr_image.unsqueeze(dim=0).cuda()
                          curr_image_feature = feature_extractors[camera_names[cam_id]](curr_image)
                          curr_image_feature = curr_image_feature.squeeze(3).squeeze(2)
                          all_cam_features.append(curr_image_feature)
                      curr_image_feature = torch.cat(all_cam_features, dim=1)
  
                      ### Visual feature
                      # curr_feature = curr_image_feature
  
                      ### State feature
                      # curr_feature = qpos
  
                      ### Both features
                      curr_feature = [curr_image_feature, qpos]
  
                      action = calculate_nearest_neighbors(curr_feature, train_inputs, train_targets, k, state_weight) # TODO use this
                      action = action.squeeze(0).cpu().numpy()
                      action = np.concatenate([action[:-BASE_DELAY, :-2], action[BASE_DELAY:, -2:]], axis=1)
                      print(f'Query: {(time.time() - start_time):.3f}s')
  
                  curr_action = action[t % query_freq]
                  target_qpos = curr_action[:-2]
                  base_action = curr_action[-2:]
  
                  # ### SAFETY
                  # max_a = 0.05
                  # curr_qpos = qpos.squeeze().cpu().numpy()
                  # target_qpos = target_qpos.clip(curr_qpos - max_a, curr_qpos + max_a)
                  # ### SAFETY
  
                  ### step the environment
                  ts = env.step(target_qpos, base_action=base_action)
                  duration = time.time() - start_time
                  # print(f'{duration:.3f}')
                  time.sleep(max(0, DT - duration))
  
                  ### save things for visualization
                  qpos_list.append(qpos_numpy)
                  target_qpos_list.append(target_qpos)
                  rewards.append(ts.reward)
  
                  # if real_robot and t != 0 and t % 60 == 0:
                  #    e()
              plt.close()
          if real_robot:
              env.puppet_bot_left.dxl.robot_set_operating_modes("single", "gripper", "position")
              env.puppet_bot_right.dxl.robot_set_operating_modes("single", "gripper", "position")
              env.puppet_bot_left.dxl.robot_set_operating_modes("single", "gripper", "pwm")
              env.puppet_bot_right.dxl.robot_set_operating_modes("single", "gripper", "pwm")
  
          rewards = np.array(rewards)
          episode_return = np.sum(rewards[rewards!=None])
          episode_returns.append(episode_return)
          max_reward = np.max(rewards)
          max_rewards.append(max_reward)
  
          print(f'{episode_return=}, {max_reward=}')
          if save_episode:
              save_videos(image_list, DT, video_path=os.path.join(ckpt_dir, f'video{rollout_id}.mp4'))
              # visualize_joints(qpos_list, target_qpos_list, plot_path=os.path.join(ckpt_dir, f'qpos{rollout_id}.png'))
              # visualize_joints(qpos_list, example_qpos, plot_path=os.path.join(ckpt_dir, f'qpos_reference{rollout_id}.png'), label_overwrite=("policy", "dataset"))
  
      success_rate = np.mean(np.array(max_rewards) == env_max_reward)
      avg_return = np.mean(episode_returns)
      summary_str = f'\nSuccess rate: {success_rate}\nAverage return: {avg_return}\n\n'
      for r in range(env_max_reward+1):
          more_or_equal_r = (np.array(max_rewards) >= r).sum()
          more_or_equal_r_rate = more_or_equal_r / num_rollouts
          summary_str += f'Reward >= {r}: {more_or_equal_r}/{num_rollouts} = {more_or_equal_r_rate*100}%\n'
  
      print(summary_str)
  
      # save success rate to txt
      result_file_name = f'result_{skip}_{k}' + '.txt'
      with open(os.path.join(ckpt_dir, result_file_name), 'w') as f:
          f.write(summary_str)
          f.write(repr(episode_returns))
          f.write('\n\n')
          f.write(repr(max_rewards))
  
      return success_rate, avg_return
  
  
  
  def get_image(ts, camera_names):
      if 'images' in ts.observation:
          curr_images = []
          for cam_name in camera_names:
              curr_image = rearrange(ts.observation['images'][cam_name], 'h w c -> c h w')
              curr_images.append(curr_image)
          curr_image_raw = np.stack(curr_images, axis=0)
      else:
          curr_image_raw = rearrange(ts.observation['image'], 'h w c -> c h w')
      curr_image = torch.from_numpy(curr_image_raw / 255.0).float().cuda().unsqueeze(0)
      curr_image_raw = rearrange(curr_image_raw, 'b c h w -> b h w c')
      return curr_image, curr_image_raw
  
  
  def expand_greyscale(t):
      return t.expand(3, -1, -1)
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser()
      parser.add_argument('--onscreen_render', action='store_true')
      parser.add_argument('--dataset_dir', action='store', type=str, help='The text to parse.', required=True)
      parser.add_argument('--model_dir', action='store', type=str, help='model_dir', required=True)
      parser.add_argument('--task_name', action='store', type=str, help='task_name', required=True)
      parser.add_argument('--ckpt_dir', action='store', type=str, help='The text to parse.', required=True)
      main(vars(parser.parse_args()))
  