diff --git a/aloha-devel/act/detr/README.md b/aloha-devel/act/detr/README.md new file mode 100644 index 0000000000000000000000000000000000000000..500b1b8d01108f8ff99b2c505a58cdd43a546fee --- /dev/null +++ b/aloha-devel/act/detr/README.md @@ -0,0 +1,9 @@ +This part of the codebase is modified from DETR https://github.com/facebookresearch/detr under APACHE 2.0. + + @article{Carion2020EndtoEndOD, + title={End-to-End Object Detection with Transformers}, + author={Nicolas Carion and Francisco Massa and Gabriel Synnaeve and Nicolas Usunier and Alexander Kirillov and Sergey Zagoruyko}, + journal={ArXiv}, + year={2020}, + volume={abs/2005.12872} + } \ No newline at end of file diff --git a/aloha-devel/act/detr/setup.py b/aloha-devel/act/detr/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..55d18c0db74e27a687b6d6e0fb236e1cbb801f20 --- /dev/null +++ b/aloha-devel/act/detr/setup.py @@ -0,0 +1,10 @@ +from distutils.core import setup +from setuptools import find_packages + +setup( + name='detr', + version='0.0.0', + packages=find_packages(), + license='MIT License', + long_description=open('README.md').read(), +) \ No newline at end of file diff --git a/aloha-devel/act/detr/util/__pycache__/__init__.cpython-38.pyc b/aloha-devel/act/detr/util/__pycache__/__init__.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f2e8436cf4024a687a170210539e7f524774235 Binary files /dev/null and b/aloha-devel/act/detr/util/__pycache__/__init__.cpython-38.pyc differ diff --git a/aloha-devel/act/inference.py b/aloha-devel/act/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..d26b43fcdf1598d44041800c8a07e7c25de89ca7 --- /dev/null +++ b/aloha-devel/act/inference.py @@ -0,0 +1,766 @@ +#!/home/lin/software/miniconda3/envs/aloha/bin/python +# -- coding: UTF-8 +""" +#!/usr/bin/python3 +""" + +import torch +import numpy as np +import os +import pickle +import argparse +from einops import rearrange + +from utils import compute_dict_mean, set_seed, detach_dict # helper functions +from policy import ACTPolicy, CNNMLPPolicy, DiffusionPolicy +import collections +from collections import deque + +import rospy +from std_msgs.msg import Header +from geometry_msgs.msg import Twist +from sensor_msgs.msg import JointState, Image +from nav_msgs.msg import Odometry +from cv_bridge import CvBridge +import time +import threading +import math +import threading + + +import sys +sys.path.append("./") + +task_config = {'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist']} + +inference_thread = None +inference_lock = threading.Lock() +inference_actions = None +inference_timestep = None + + +def actions_interpolation(args, pre_action, actions, stats): + steps = np.concatenate((np.array(args.arm_steps_length), np.array(args.arm_steps_length)), axis=0) + pre_process = lambda s_qpos: (s_qpos - stats['qpos_mean']) / stats['qpos_std'] + post_process = lambda a: a * stats['qpos_std'] + stats['qpos_mean'] + result = [pre_action] + post_action = post_process(actions[0]) + # print("pre_action:", pre_action[7:]) + # print("actions_interpolation1:", post_action[:, 7:]) + max_diff_index = 0 + max_diff = -1 + for i in range(post_action.shape[0]): + diff = 0 + for j in range(pre_action.shape[0]): + if j == 6 or j == 13: + continue + diff += math.fabs(pre_action[j] - post_action[i][j]) + if diff > max_diff: + max_diff = diff + max_diff_index = i + + for i in range(max_diff_index, post_action.shape[0]): + step = max([math.floor(math.fabs(result[-1][j] - post_action[i][j])/steps[j]) for j in range(pre_action.shape[0])]) + inter = np.linspace(result[-1], post_action[i], step+2) + result.extend(inter[1:]) + while len(result) < args.chunk_size+1: + result.append(result[-1]) + result = np.array(result)[1:args.chunk_size+1] + # print("actions_interpolation2:", result.shape, result[:, 7:]) + result = pre_process(result) + result = result[np.newaxis, :] + return result + + +def get_model_config(args): + # 设置随机种子，你可以确保在相同的初始条件下，每次运行代码时生成的随机数序列是相同的。 + set_seed(1) + + # 如果是ACT策略 + # fixed parameters + if args.policy_class == 'ACT': + policy_config = {'lr': args.lr, + 'lr_backbone': args.lr_backbone, + 'backbone': args.backbone, + 'masks': args.masks, + 'weight_decay': args.weight_decay, + 'dilation': args.dilation, + 'position_embedding': args.position_embedding, + 'loss_function': args.loss_function, + 'chunk_size': args.chunk_size, # 查询 + 'camera_names': task_config['camera_names'], + 'use_depth_image': args.use_depth_image, + 'use_robot_base': args.use_robot_base, + 'kl_weight': args.kl_weight, # kl散度权重 + 'hidden_dim': args.hidden_dim, # 隐藏层维度 + 'dim_feedforward': args.dim_feedforward, + 'enc_layers': args.enc_layers, + 'dec_layers': args.dec_layers, + 'nheads': args.nheads, + 'dropout': args.dropout, + 'pre_norm': args.pre_norm + } + elif args.policy_class == 'CNNMLP': + policy_config = {'lr': args.lr, + 'lr_backbone': args.lr_backbone, + 'backbone': args.backbone, + 'masks': args.masks, + 'weight_decay': args.weight_decay, + 'dilation': args.dilation, + 'position_embedding': args.position_embedding, + 'loss_function': args.loss_function, + 'chunk_size': 1, # 查询 + 'camera_names': task_config['camera_names'], + 'use_depth_image': args.use_depth_image, + 'use_robot_base': args.use_robot_base + } + + elif args.policy_class == 'Diffusion': + policy_config = {'lr': args.lr, + 'lr_backbone': args.lr_backbone, + 'backbone': args.backbone, + 'masks': args.masks, + 'weight_decay': args.weight_decay, + 'dilation': args.dilation, + 'position_embedding': args.position_embedding, + 'loss_function': args.loss_function, + 'chunk_size': args.chunk_size, # 查询 + 'camera_names': task_config['camera_names'], + 'use_depth_image': args.use_depth_image, + 'use_robot_base': args.use_robot_base, + 'observation_horizon': args.observation_horizon, + 'action_horizon': args.action_horizon, + 'num_inference_timesteps': args.num_inference_timesteps, + 'ema_power': args.ema_power + } + else: + raise NotImplementedError + + config = { + 'ckpt_dir': args.ckpt_dir, + 'ckpt_name': args.ckpt_name, + 'ckpt_stats_name': args.ckpt_stats_name, + 'episode_len': args.max_publish_step, + 'state_dim': args.state_dim, + 'policy_class': args.policy_class, + 'policy_config': policy_config, + 'temporal_agg': args.temporal_agg, + 'camera_names': task_config['camera_names'], + } + return config + + +def make_policy(policy_class, policy_config): + if policy_class == 'ACT': + policy = ACTPolicy(policy_config) + elif policy_class == 'CNNMLP': + policy = CNNMLPPolicy(policy_config) + elif policy_class == 'Diffusion': + policy = DiffusionPolicy(policy_config) + else: + raise NotImplementedError + return policy + + +def get_image(observation, camera_names): + curr_images = [] + for cam_name in camera_names: + curr_image = rearrange(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) + return curr_image + + +def get_depth_image(observation, camera_names): + curr_images = [] + for cam_name in camera_names: + curr_images.append(observation['images_depth'][cam_name]) + curr_image = np.stack(curr_images, axis=0) + curr_image = torch.from_numpy(curr_image / 255.0).float().cuda().unsqueeze(0) + return curr_image + + +def inference_process(args, config, ros_operator, policy, stats, t, pre_action): + global inference_lock + global inference_actions + global inference_timestep + print_flag = True + pre_pos_process = lambda s_qpos: (s_qpos - stats['qpos_mean']) / stats['qpos_std'] + pre_action_process = lambda next_action: (next_action - stats["action_mean"]) / stats["action_std"] + rate = rospy.Rate(args.publish_rate) + while True and not rospy.is_shutdown(): + result = ros_operator.get_frame() + if not result: + if print_flag: + print("syn fail") + print_flag = False + rate.sleep() + continue + print_flag = True + (img_front, img_left, img_right, img_front_depth, img_left_depth, img_right_depth, + puppet_arm_left, puppet_arm_right, robot_base) = result + obs = collections.OrderedDict() + image_dict = dict() + + image_dict[config['camera_names'][0]] = img_front + image_dict[config['camera_names'][1]] = img_left + image_dict[config['camera_names'][2]] = img_right + + + obs['images'] = image_dict + + if args.use_depth_image: + image_depth_dict = dict() + image_depth_dict[config['camera_names'][0]] = img_front_depth + image_depth_dict[config['camera_names'][1]] = img_left_depth + image_depth_dict[config['camera_names'][2]] = img_right_depth + obs['images_depth'] = image_depth_dict + + obs['qpos'] = np.concatenate( + (np.array(puppet_arm_left.position), np.array(puppet_arm_right.position)), axis=0) + obs['qvel'] = np.concatenate( + (np.array(puppet_arm_left.velocity), np.array(puppet_arm_right.velocity)), axis=0) + obs['effort'] = np.concatenate( + (np.array(puppet_arm_left.effort), np.array(puppet_arm_right.effort)), axis=0) + if args.use_robot_base: + obs['base_vel'] = [robot_base.twist.twist.linear.x, robot_base.twist.twist.angular.z] + obs['qpos'] = np.concatenate((obs['qpos'], obs['base_vel']), axis=0) + else: + obs['base_vel'] = [0.0, 0.0] + # qpos_numpy = np.array(obs['qpos']) + + # 归一化处理qpos 并转到cuda + qpos = pre_pos_process(obs['qpos']) + qpos = torch.from_numpy(qpos).float().cuda().unsqueeze(0) + # 当前图像curr_image获取图像 + curr_image = get_image(obs, config['camera_names']) + curr_depth_image = None + if args.use_depth_image: + curr_depth_image = get_depth_image(obs, config['camera_names']) + start_time = time.time() + all_actions = policy(curr_image, curr_depth_image, qpos) + end_time = time.time() + print("model cost time: ", end_time -start_time) + inference_lock.acquire() + inference_actions = all_actions.cpu().detach().numpy() + if pre_action is None: + pre_action = obs['qpos'] + # print("obs['qpos']:", obs['qpos'][7:]) + if args.use_actions_interpolation: + inference_actions = actions_interpolation(args, pre_action, inference_actions, stats) + inference_timestep = t + inference_lock.release() + break + + +def model_inference(args, config, ros_operator, save_episode=True): + global inference_lock + global inference_actions + global inference_timestep + global inference_thread + set_seed(1000) + + # 1 创建模型数据 继承nn.Module + policy = make_policy(config['policy_class'], config['policy_config']) + # print("model structure\n", policy.model) + + # 2 加载模型权重 + ckpt_path = os.path.join(config['ckpt_dir'], config['ckpt_name']) + state_dict = torch.load(ckpt_path) + new_state_dict = {} + for key, value in state_dict.items(): + if key in ["model.is_pad_head.weight", "model.is_pad_head.bias"]: + continue + if key in ["model.input_proj_next_action.weight", "model.input_proj_next_action.bias"]: + continue + new_state_dict[key] = value + loading_status = policy.deserialize(new_state_dict) + if not loading_status: + print("ckpt path not exist") + return False + + # 3 模型设置为cuda模式和验证模式 + policy.cuda() + policy.eval() + + # 4 加载统计值 + stats_path = os.path.join(config['ckpt_dir'], config['ckpt_stats_name']) + # 统计的数据 # 加载action_mean, action_std, qpos_mean, qpos_std 14维 + with open(stats_path, 'rb') as f: + stats = pickle.load(f) + + # 数据预处理和后处理函数定义 + pre_process = lambda s_qpos: (s_qpos - stats['qpos_mean']) / stats['qpos_std'] + post_process = lambda a: a * stats['qpos_std'] + stats['qpos_mean'] + + max_publish_step = config['episode_len'] + chunk_size = config['policy_config']['chunk_size'] + + # 发布基础的姿态 + left0 = [-0.00133514404296875, 0.00209808349609375, 0.01583099365234375, -0.032616615295410156, -0.00286102294921875, 0.00095367431640625, 3.557830810546875] + right0 = [-0.00133514404296875, 0.00438690185546875, 0.034523963928222656, -0.053597450256347656, -0.00476837158203125, -0.00209808349609375, 3.557830810546875] + left1 = [-0.00133514404296875, 0.00209808349609375, 0.01583099365234375, -0.032616615295410156, -0.00286102294921875, 0.00095367431640625, -0.3393220901489258] + right1 = [-0.00133514404296875, 0.00247955322265625, 0.01583099365234375, -0.032616615295410156, -0.00286102294921875, 0.00095367431640625, -0.3397035598754883] + + ros_operator.puppet_arm_publish_continuous(left0, right0) + input("Enter any key to continue :") + ros_operator.puppet_arm_publish_continuous(left1, right1) + action = None + # 推理 + with torch.inference_mode(): + while True and not rospy.is_shutdown(): + # 每个回合的步数 + t = 0 + max_t = 0 + rate = rospy.Rate(args.publish_rate) + if config['temporal_agg']: + all_time_actions = np.zeros([max_publish_step, max_publish_step + chunk_size, config['state_dim']]) + while t < max_publish_step and not rospy.is_shutdown(): + # start_time = time.time() + # query policy + if config['policy_class'] == "ACT": + if t >= max_t: + pre_action = action + inference_thread = threading.Thread(target=inference_process, + args=(args, config, ros_operator, + policy, stats, t, pre_action)) + inference_thread.start() + inference_thread.join() + inference_lock.acquire() + if inference_actions is not None: + inference_thread = None + all_actions = inference_actions + inference_actions = None + max_t = t + args.pos_lookahead_step + if config['temporal_agg']: + all_time_actions[[t], t:t + chunk_size] = all_actions + inference_lock.release() + if config['temporal_agg']: + actions_for_curr_step = all_time_actions[:, t] + actions_populated = np.all(actions_for_curr_step != 0, axis=1) + actions_for_curr_step = actions_for_curr_step[actions_populated] + k = 0.01 + exp_weights = np.exp(-k * np.arange(len(actions_for_curr_step))) + exp_weights = exp_weights / exp_weights.sum() + exp_weights = exp_weights[:, np.newaxis] + raw_action = (actions_for_curr_step * exp_weights).sum(axis=0, keepdims=True) + else: + if args.pos_lookahead_step != 0: + raw_action = all_actions[:, t % args.pos_lookahead_step] + else: + raw_action = all_actions[:, t % chunk_size] + else: + raise NotImplementedError + action = post_process(raw_action[0]) + left_action = action[:7] # 取7维度 + right_action = action[7:14] + ros_operator.puppet_arm_publish(left_action, right_action) # puppet_arm_publish_continuous_thread + if args.use_robot_base: + vel_action = action[14:16] + ros_operator.robot_base_publish(vel_action) + t += 1 + # end_time = time.time() + # print("publish: ", t) + # print("time:", end_time - start_time) + # print("left_action:", left_action) + # print("right_action:", right_action) + rate.sleep() + + +class RosOperator: + def __init__(self, args): + self.robot_base_deque = None + self.puppet_arm_right_deque = None + self.puppet_arm_left_deque = None + self.img_front_deque = None + self.img_right_deque = None + self.img_left_deque = None + self.img_front_depth_deque = None + self.img_right_depth_deque = None + self.img_left_depth_deque = None + self.bridge = None + self.puppet_arm_left_publisher = None + self.puppet_arm_right_publisher = None + self.robot_base_publisher = None + self.puppet_arm_publish_thread = None + self.puppet_arm_publish_lock = None + self.args = args + self.ctrl_state = False + self.ctrl_state_lock = threading.Lock() + self.init() + self.init_ros() + + def init(self): + self.bridge = CvBridge() + self.img_left_deque = deque() + self.img_right_deque = deque() + self.img_front_deque = deque() + self.img_left_depth_deque = deque() + self.img_right_depth_deque = deque() + self.img_front_depth_deque = deque() + self.puppet_arm_left_deque = deque() + self.puppet_arm_right_deque = deque() + self.robot_base_deque = deque() + self.puppet_arm_publish_lock = threading.Lock() + self.puppet_arm_publish_lock.acquire() + + def puppet_arm_publish(self, left, right): + joint_state_msg = JointState() + joint_state_msg.header = Header() + joint_state_msg.header.stamp = rospy.Time.now() # 设置时间戳 + joint_state_msg.name = ['joint0', 'joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6'] # 设置关节名称 + joint_state_msg.position = left + self.puppet_arm_left_publisher.publish(joint_state_msg) + joint_state_msg.position = right + self.puppet_arm_right_publisher.publish(joint_state_msg) + + def robot_base_publish(self, vel): + vel_msg = Twist() + vel_msg.linear.x = vel[0] + vel_msg.linear.y = 0 + vel_msg.linear.z = 0 + vel_msg.angular.x = 0 + vel_msg.angular.y = 0 + vel_msg.angular.z = vel[1] + self.robot_base_publisher.publish(vel_msg) + + def puppet_arm_publish_continuous(self, left, right): + rate = rospy.Rate(self.args.publish_rate) + left_arm = None + right_arm = None + while True and not rospy.is_shutdown(): + if len(self.puppet_arm_left_deque) != 0: + left_arm = list(self.puppet_arm_left_deque[-1].position) + if len(self.puppet_arm_right_deque) != 0: + right_arm = list(self.puppet_arm_right_deque[-1].position) + if left_arm is None or right_arm is None: + rate.sleep() + continue + else: + break + left_symbol = [1 if left[i] - left_arm[i] > 0 else -1 for i in range(len(left))] + right_symbol = [1 if right[i] - right_arm[i] > 0 else -1 for i in range(len(right))] + flag = True + step = 0 + while flag and not rospy.is_shutdown(): + if self.puppet_arm_publish_lock.acquire(False): + return + left_diff = [abs(left[i] - left_arm[i]) for i in range(len(left))] + right_diff = [abs(right[i] - right_arm[i]) for i in range(len(right))] + flag = False + for i in range(len(left)): + if left_diff[i] < self.args.arm_steps_length[i]: + left_arm[i] = left[i] + else: + left_arm[i] += left_symbol[i] * self.args.arm_steps_length[i] + flag = True + for i in range(len(right)): + if right_diff[i] < self.args.arm_steps_length[i]: + right_arm[i] = right[i] + else: + right_arm[i] += right_symbol[i] * self.args.arm_steps_length[i] + flag = True + joint_state_msg = JointState() + joint_state_msg.header = Header() + joint_state_msg.header.stamp = rospy.Time.now() # 设置时间戳 + joint_state_msg.name = ['joint0', 'joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6'] # 设置关节名称 + joint_state_msg.position = left_arm + self.puppet_arm_left_publisher.publish(joint_state_msg) + joint_state_msg.position = right_arm + self.puppet_arm_right_publisher.publish(joint_state_msg) + step += 1 + print("puppet_arm_publish_continuous:", step) + rate.sleep() + + def puppet_arm_publish_linear(self, left, right): + num_step = 100 + rate = rospy.Rate(200) + + left_arm = None + right_arm = None + + while True and not rospy.is_shutdown(): + if len(self.puppet_arm_left_deque) != 0: + left_arm = list(self.puppet_arm_left_deque[-1].position) + if len(self.puppet_arm_right_deque) != 0: + right_arm = list(self.puppet_arm_right_deque[-1].position) + if left_arm is None or right_arm is None: + rate.sleep() + continue + else: + break + + traj_left_list = np.linspace(left_arm, left, num_step) + traj_right_list = np.linspace(right_arm, right, num_step) + + for i in range(len(traj_left_list)): + traj_left = traj_left_list[i] + traj_right = traj_right_list[i] + traj_left[-1] = left[-1] + traj_right[-1] = right[-1] + joint_state_msg = JointState() + joint_state_msg.header = Header() + joint_state_msg.header.stamp = rospy.Time.now() # 设置时间戳 + joint_state_msg.name = ['joint0', 'joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6'] # 设置关节名称 + joint_state_msg.position = traj_left + self.puppet_arm_left_publisher.publish(joint_state_msg) + joint_state_msg.position = traj_right + self.puppet_arm_right_publisher.publish(joint_state_msg) + rate.sleep() + + def puppet_arm_publish_continuous_thread(self, left, right): + if self.puppet_arm_publish_thread is not None: + self.puppet_arm_publish_lock.release() + self.puppet_arm_publish_thread.join() + self.puppet_arm_publish_lock.acquire(False) + self.puppet_arm_publish_thread = None + self.puppet_arm_publish_thread = threading.Thread(target=self.puppet_arm_publish_continuous, args=(left, right)) + self.puppet_arm_publish_thread.start() + + def get_frame(self): + if len(self.img_left_deque) == 0 or len(self.img_right_deque) == 0 or len(self.img_front_deque) == 0 or \ + (self.args.use_depth_image and (len(self.img_left_depth_deque) == 0 or len(self.img_right_depth_deque) == 0 or len(self.img_front_depth_deque) == 0)): + return False + if self.args.use_depth_image: + frame_time = min([self.img_left_deque[-1].header.stamp.to_sec(), self.img_right_deque[-1].header.stamp.to_sec(), self.img_front_deque[-1].header.stamp.to_sec(), + self.img_left_depth_deque[-1].header.stamp.to_sec(), self.img_right_depth_deque[-1].header.stamp.to_sec(), self.img_front_depth_deque[-1].header.stamp.to_sec()]) + else: + frame_time = min([self.img_left_deque[-1].header.stamp.to_sec(), self.img_right_deque[-1].header.stamp.to_sec(), self.img_front_deque[-1].header.stamp.to_sec()]) + + if len(self.img_left_deque) == 0 or self.img_left_deque[-1].header.stamp.to_sec() < frame_time: + return False + if len(self.img_right_deque) == 0 or self.img_right_deque[-1].header.stamp.to_sec() < frame_time: + return False + if len(self.img_front_deque) == 0 or self.img_front_deque[-1].header.stamp.to_sec() < frame_time: + return False + if len(self.puppet_arm_left_deque) == 0 or self.puppet_arm_left_deque[-1].header.stamp.to_sec() < frame_time: + return False + if len(self.puppet_arm_right_deque) == 0 or self.puppet_arm_right_deque[-1].header.stamp.to_sec() < frame_time: + return False + if self.args.use_depth_image and (len(self.img_left_depth_deque) == 0 or self.img_left_depth_deque[-1].header.stamp.to_sec() < frame_time): + return False + if self.args.use_depth_image and (len(self.img_right_depth_deque) == 0 or self.img_right_depth_deque[-1].header.stamp.to_sec() < frame_time): + return False + if self.args.use_depth_image and (len(self.img_front_depth_deque) == 0 or self.img_front_depth_deque[-1].header.stamp.to_sec() < frame_time): + return False + if self.args.use_robot_base and (len(self.robot_base_deque) == 0 or self.robot_base_deque[-1].header.stamp.to_sec() < frame_time): + return False + + while self.img_left_deque[0].header.stamp.to_sec() < frame_time: + self.img_left_deque.popleft() + img_left = self.bridge.imgmsg_to_cv2(self.img_left_deque.popleft(), 'passthrough') + + while self.img_right_deque[0].header.stamp.to_sec() < frame_time: + self.img_right_deque.popleft() + img_right = self.bridge.imgmsg_to_cv2(self.img_right_deque.popleft(), 'passthrough') + + while self.img_front_deque[0].header.stamp.to_sec() < frame_time: + self.img_front_deque.popleft() + img_front = self.bridge.imgmsg_to_cv2(self.img_front_deque.popleft(), 'passthrough') + + while self.puppet_arm_left_deque[0].header.stamp.to_sec() < frame_time: + self.puppet_arm_left_deque.popleft() + puppet_arm_left = self.puppet_arm_left_deque.popleft() + + while self.puppet_arm_right_deque[0].header.stamp.to_sec() < frame_time: + self.puppet_arm_right_deque.popleft() + puppet_arm_right = self.puppet_arm_right_deque.popleft() + + img_left_depth = None + if self.args.use_depth_image: + while self.img_left_depth_deque[0].header.stamp.to_sec() < frame_time: + self.img_left_depth_deque.popleft() + img_left_depth = self.bridge.imgmsg_to_cv2(self.img_left_depth_deque.popleft(), 'passthrough') + + img_right_depth = None + if self.args.use_depth_image: + while self.img_right_depth_deque[0].header.stamp.to_sec() < frame_time: + self.img_right_depth_deque.popleft() + img_right_depth = self.bridge.imgmsg_to_cv2(self.img_right_depth_deque.popleft(), 'passthrough') + + img_front_depth = None + if self.args.use_depth_image: + while self.img_front_depth_deque[0].header.stamp.to_sec() < frame_time: + self.img_front_depth_deque.popleft() + img_front_depth = self.bridge.imgmsg_to_cv2(self.img_front_depth_deque.popleft(), 'passthrough') + + robot_base = None + if self.args.use_robot_base: + while self.robot_base_deque[0].header.stamp.to_sec() < frame_time: + self.robot_base_deque.popleft() + robot_base = self.robot_base_deque.popleft() + + return (img_front, img_left, img_right, img_front_depth, img_left_depth, img_right_depth, + puppet_arm_left, puppet_arm_right, robot_base) + + def img_left_callback(self, msg): + if len(self.img_left_deque) >= 2000: + self.img_left_deque.popleft() + self.img_left_deque.append(msg) + + def img_right_callback(self, msg): + if len(self.img_right_deque) >= 2000: + self.img_right_deque.popleft() + self.img_right_deque.append(msg) + + def img_front_callback(self, msg): + if len(self.img_front_deque) >= 2000: + self.img_front_deque.popleft() + self.img_front_deque.append(msg) + + def img_left_depth_callback(self, msg): + if len(self.img_left_depth_deque) >= 2000: + self.img_left_depth_deque.popleft() + self.img_left_depth_deque.append(msg) + + def img_right_depth_callback(self, msg): + if len(self.img_right_depth_deque) >= 2000: + self.img_right_depth_deque.popleft() + self.img_right_depth_deque.append(msg) + + def img_front_depth_callback(self, msg): + if len(self.img_front_depth_deque) >= 2000: + self.img_front_depth_deque.popleft() + self.img_front_depth_deque.append(msg) + + def puppet_arm_left_callback(self, msg): + if len(self.puppet_arm_left_deque) >= 2000: + self.puppet_arm_left_deque.popleft() + self.puppet_arm_left_deque.append(msg) + + def puppet_arm_right_callback(self, msg): + if len(self.puppet_arm_right_deque) >= 2000: + self.puppet_arm_right_deque.popleft() + self.puppet_arm_right_deque.append(msg) + + def robot_base_callback(self, msg): + if len(self.robot_base_deque) >= 2000: + self.robot_base_deque.popleft() + self.robot_base_deque.append(msg) + + def ctrl_callback(self, msg): + self.ctrl_state_lock.acquire() + self.ctrl_state = msg.data + self.ctrl_state_lock.release() + + def get_ctrl_state(self): + self.ctrl_state_lock.acquire() + state = self.ctrl_state + self.ctrl_state_lock.release() + return state + + def init_ros(self): + rospy.init_node('joint_state_publisher', anonymous=True) + rospy.Subscriber(self.args.img_left_topic, Image, self.img_left_callback, queue_size=1000, tcp_nodelay=True) + rospy.Subscriber(self.args.img_right_topic, Image, self.img_right_callback, queue_size=1000, tcp_nodelay=True) + rospy.Subscriber(self.args.img_front_topic, Image, self.img_front_callback, queue_size=1000, tcp_nodelay=True) + if self.args.use_depth_image: + rospy.Subscriber(self.args.img_left_depth_topic, Image, self.img_left_depth_callback, queue_size=1000, tcp_nodelay=True) + rospy.Subscriber(self.args.img_right_depth_topic, Image, self.img_right_depth_callback, queue_size=1000, tcp_nodelay=True) + rospy.Subscriber(self.args.img_front_depth_topic, Image, self.img_front_depth_callback, queue_size=1000, tcp_nodelay=True) + rospy.Subscriber(self.args.puppet_arm_left_topic, JointState, self.puppet_arm_left_callback, queue_size=1000, tcp_nodelay=True) + rospy.Subscriber(self.args.puppet_arm_right_topic, JointState, self.puppet_arm_right_callback, queue_size=1000, tcp_nodelay=True) + rospy.Subscriber(self.args.robot_base_topic, Odometry, self.robot_base_callback, queue_size=1000, tcp_nodelay=True) + self.puppet_arm_left_publisher = rospy.Publisher(self.args.puppet_arm_left_cmd_topic, JointState, queue_size=10) + self.puppet_arm_right_publisher = rospy.Publisher(self.args.puppet_arm_right_cmd_topic, JointState, queue_size=10) + self.robot_base_publisher = rospy.Publisher(self.args.robot_base_cmd_topic, Twist, queue_size=10) + + +def get_arguments(): + parser = argparse.ArgumentParser() + parser.add_argument('--ckpt_dir', action='store', type=str, help='ckpt_dir', required=True) + parser.add_argument('--task_name', action='store', type=str, help='task_name', default='aloha_mobile_dummy', required=False) + parser.add_argument('--max_publish_step', action='store', type=int, help='max_publish_step', default=10000, required=False) + parser.add_argument('--ckpt_name', action='store', type=str, help='ckpt_name', default='policy_best.ckpt', required=False) + parser.add_argument('--ckpt_stats_name', action='store', type=str, help='ckpt_stats_name', default='dataset_stats.pkl', required=False) + parser.add_argument('--policy_class', action='store', type=str, help='policy_class, capitalize', default='ACT', required=False) + parser.add_argument('--batch_size', action='store', type=int, help='batch_size', default=8, required=False) + parser.add_argument('--seed', action='store', type=int, help='seed', default=0, required=False) + parser.add_argument('--num_epochs', action='store', type=int, help='num_epochs', default=2000, required=False) + parser.add_argument('--lr', action='store', type=float, help='lr', default=1e-5, required=False) + parser.add_argument('--weight_decay', type=float, help='weight_decay', default=1e-4, required=False) + parser.add_argument('--dilation', action='store_true', + help="If true, we replace stride with dilation in the last convolutional block (DC5)", required=False) + parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'), + help="Type of positional embedding to use on top of the image features", required=False) + parser.add_argument('--masks', action='store_true', + help="Train segmentation head if the flag is provided") + parser.add_argument('--kl_weight', action='store', type=int, help='KL Weight', default=10, required=False) + parser.add_argument('--hidden_dim', action='store', type=int, help='hidden_dim', default=512, required=False) + parser.add_argument('--dim_feedforward', action='store', type=int, help='dim_feedforward', default=3200, required=False) + parser.add_argument('--temporal_agg', action='store', type=bool, help='temporal_agg', default=True, required=False) + + parser.add_argument('--state_dim', action='store', type=int, help='state_dim', default=14, required=False) + parser.add_argument('--lr_backbone', action='store', type=float, help='lr_backbone', default=1e-5, required=False) + parser.add_argument('--backbone', action='store', type=str, help='backbone', default='resnet18', required=False) + parser.add_argument('--loss_function', action='store', type=str, help='loss_function l1 l2 l1+l2', default='l1', required=False) + parser.add_argument('--enc_layers', action='store', type=int, help='enc_layers', default=4, required=False) + parser.add_argument('--dec_layers', action='store', type=int, help='dec_layers', default=7, required=False) + parser.add_argument('--nheads', action='store', type=int, help='nheads', default=8, required=False) + parser.add_argument('--dropout', default=0.1, type=float, help="Dropout applied in the transformer", required=False) + parser.add_argument('--pre_norm', action='store_true', required=False) + + parser.add_argument('--img_front_topic', action='store', type=str, help='img_front_topic', + default='/camera_f/color/image_raw', required=False) + parser.add_argument('--img_left_topic', action='store', type=str, help='img_left_topic', + default='/camera_l/color/image_raw', required=False) + parser.add_argument('--img_right_topic', action='store', type=str, help='img_right_topic', + default='/camera_r/color/image_raw', required=False) + + parser.add_argument('--img_front_depth_topic', action='store', type=str, help='img_front_depth_topic', + default='/camera_f/depth/image_raw', required=False) + parser.add_argument('--img_left_depth_topic', action='store', type=str, help='img_left_depth_topic', + default='/camera_l/depth/image_raw', required=False) + parser.add_argument('--img_right_depth_topic', action='store', type=str, help='img_right_depth_topic', + default='/camera_r/depth/image_raw', required=False) + + parser.add_argument('--puppet_arm_left_cmd_topic', action='store', type=str, help='puppet_arm_left_cmd_topic', + default='/master/joint_left', required=False) + parser.add_argument('--puppet_arm_right_cmd_topic', action='store', type=str, help='puppet_arm_right_cmd_topic', + default='/master/joint_right', required=False) + parser.add_argument('--puppet_arm_left_topic', action='store', type=str, help='puppet_arm_left_topic', + default='/puppet/joint_left', required=False) + parser.add_argument('--puppet_arm_right_topic', action='store', type=str, help='puppet_arm_right_topic', + default='/puppet/joint_right', required=False) + + parser.add_argument('--robot_base_topic', action='store', type=str, help='robot_base_topic', + default='/odom_raw', required=False) + parser.add_argument('--robot_base_cmd_topic', action='store', type=str, help='robot_base_topic', + default='/cmd_vel', required=False) + parser.add_argument('--use_robot_base', action='store', type=bool, help='use_robot_base', + default=False, required=False) + parser.add_argument('--publish_rate', action='store', type=int, help='publish_rate', + default=40, required=False) + parser.add_argument('--pos_lookahead_step', action='store', type=int, help='pos_lookahead_step', + default=0, required=False) + parser.add_argument('--chunk_size', action='store', type=int, help='chunk_size', + default=32, required=False) + parser.add_argument('--arm_steps_length', action='store', type=float, help='arm_steps_length', + default=[0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.2], required=False) + + parser.add_argument('--use_actions_interpolation', action='store', type=bool, help='use_actions_interpolation', + default=False, required=False) + parser.add_argument('--use_depth_image', action='store', type=bool, help='use_depth_image', + default=False, required=False) + + # for Diffusion + parser.add_argument('--observation_horizon', action='store', type=int, help='observation_horizon', default=1, required=False) + parser.add_argument('--action_horizon', action='store', type=int, help='action_horizon', default=8, required=False) + parser.add_argument('--num_inference_timesteps', action='store', type=int, help='num_inference_timesteps', default=10, required=False) + parser.add_argument('--ema_power', action='store', type=int, help='ema_power', default=0.75, required=False) + args = parser.parse_args() + return args + + +def main(): + args = get_arguments() + ros_operator = RosOperator(args) + config = get_model_config(args) + model_inference(args, config, ros_operator, save_episode=True) + + +if __name__ == '__main__': + main() +# python act/inference.py --ckpt_dir ~/train0314/ \ No newline at end of file diff --git a/aloha-devel/act/policy.py b/aloha-devel/act/policy.py new file mode 100644 index 0000000000000000000000000000000000000000..a2be7f37af38685648624a1d3a5c086fc9572ff1 --- /dev/null +++ b/aloha-devel/act/policy.py @@ -0,0 +1,166 @@ +import torch.nn as nn +from torch.nn import functional as F +import torchvision.transforms as transforms +from detr.main import build_ACT_model_and_optimizer, build_CNNMLP_model_and_optimizer, build_diffusion_model_and_optimizer + +import IPython +e = IPython.embed + + +class DiffusionPolicy(nn.Module): + def __init__(self, args_override): + super().__init__() + model, optimizer = build_diffusion_model_and_optimizer(args_override) + self.model = model + self.optimizer = optimizer + + def configure_optimizers(self): + return self.optimizer + + def __call__(self, image, depth_image, robot_state, actions=None, action_is_pad=None): + B = robot_state.shape[0] + if actions is not None: + noise, noise_pred = self.model(image, depth_image, robot_state, actions, action_is_pad) + # L2 loss + all_l2 = F.mse_loss(noise_pred, noise, reduction='none') + loss = (all_l2 * ~action_is_pad.unsqueeze(-1)).mean() + + loss_dict = {} + loss_dict['l2_loss'] = loss + loss_dict['loss'] = loss + return loss_dict, (noise, noise_pred) + else: # inference time + return self.model(image, depth_image, robot_state, actions, action_is_pad) + + def serialize(self): + return self.model.serialize() + + def deserialize(self, model_dict): + return self.model.deserialize(model_dict) + + +class ACTPolicy(nn.Module): + def __init__(self, args_override): + super().__init__() + model, optimizer = build_ACT_model_and_optimizer(args_override) + + self.model = model # CVAE decoder + self.optimizer = optimizer + self.kl_weight = args_override['kl_weight'] + self.loss_function = args_override['loss_function'] + + print(f'KL Weight {self.kl_weight}') + + def __call__(self, image, depth_image, robot_state, actions=None, action_is_pad=None): + + normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225]) + depth_normalize = transforms.Normalize(mean=[0.5], std=[0.5]) + + image = normalize(image) # 图像归一化 + if depth_image is not None: + depth_image = depth_normalize(depth_image) + + # 总共max个步 只取前model.num_queries个 + if actions is not None: # training time + actions = actions[:, :self.model.num_queries] + action_is_pad = action_is_pad[:, :self.model.num_queries] + + a_hat, (mu, logvar) = self.model(image, depth_image, robot_state, actions, action_is_pad) + + loss_dict = dict() + if self.loss_function == 'l1': + all_l1 = F.l1_loss(actions, a_hat, reduction='none') + elif self.loss_function == 'l2': + all_l1 = F.mse_loss(actions, a_hat, reduction='none') + else: + all_l1 = F.smooth_l1_loss(actions, a_hat, reduction='none') + + l1 = (all_l1 * ~action_is_pad.unsqueeze(-1)).mean() + + loss_dict['l1'] = l1 + if self.kl_weight != 0: + total_kld, dim_wise_kld, mean_kld = kl_divergence(mu, logvar) + loss_dict['kl'] = total_kld[0] + loss_dict['loss'] = loss_dict['l1'] + loss_dict['kl'] * self.kl_weight + else: + loss_dict['loss'] = loss_dict['l1'] + + return loss_dict, a_hat + else: # inference time + a_hat, (_, _) = self.model(image, depth_image, robot_state) # no action, sample from prior + return a_hat + + def configure_optimizers(self): + return self.optimizer + + def serialize(self): + return self.state_dict() + + def deserialize(self, model_dict): + return self.load_state_dict(model_dict) + + +class CNNMLPPolicy(nn.Module): + def __init__(self, args_override): + super().__init__() + model, optimizer = build_CNNMLP_model_and_optimizer(args_override) + self.model = model # decoder + self.optimizer = optimizer + self.loss_function = args_override['loss_function'] + + # 而 __call__ 在对象被调用时执行 + def __call__(self, image, depth_image, robot_state, actions=None, + action_is_pad=None): + env_state = None # TODO + + normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225]) + depth_normalize = transforms.Normalize(mean=[0.5], std=[0.5]) + image = normalize(image) # 图像归一化 + if depth_image is not None: + depth_image = depth_normalize(depth_image) + if actions is not None: # training time + actions = actions[:, 0] # 动作 + a_hat = self.model(image, depth_image, robot_state, actions, action_is_pad) + # 均方误差 + if self.loss_function == 'l1': + mse = F.l1_loss(actions, a_hat) + elif self.loss_function == 'l2': + mse = F.mse_loss(actions, a_hat) + else: + mse = F.smooth_l1_loss(actions, a_hat) + + loss_dict = dict() + loss_dict['mse'] = mse + loss_dict['loss'] = loss_dict['mse'] + return loss_dict, a_hat + + else: # inference time + a_hat = self.model(image, depth_image, robot_state) # no action, sample from prior + return a_hat + + def configure_optimizers(self): + return self.optimizer + + def serialize(self): + return self.state_dict() + + def deserialize(self, model_dict): + return self.load_state_dict(model_dict) + + +def kl_divergence(mu, logvar): + batch_size = mu.size(0) + assert batch_size != 0 + if mu.data.ndimension() == 4: + mu = mu.view(mu.size(0), mu.size(1)) + if logvar.data.ndimension() == 4: + logvar = logvar.view(logvar.size(0), logvar.size(1)) + + klds = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()) + total_kld = klds.sum(1).mean(0, True) + dimension_wise_kld = klds.mean(0) + mean_kld = klds.mean(1).mean(0, True) + + return total_kld, dimension_wise_kld, mean_kld diff --git a/aloha-devel/act/test_inference.py b/aloha-devel/act/test_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..5eadd7209052d09ac35f15a31f920e0e29ed08d7 --- /dev/null +++ b/aloha-devel/act/test_inference.py @@ -0,0 +1,180 @@ +import torch +import argparse +import os +from policy import ACTPolicy, CNNMLPPolicy, DiffusionPolicy +from train import make_policy + + +def test(args): + # a. Parse arguments is done outside + # b. Define TASK_CONFIGS and policy_config + # (Adapted from train.py) + TASK_CONFIGS = { + args.task_name: { + 'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist'], + } + } + task_config = TASK_CONFIGS[args.task_name] + camera_names = task_config['camera_names'] + + if args.policy_class == 'ACT': + policy_config = {'lr': args.lr, + 'lr_backbone': args.lr_backbone, + 'backbone': args.backbone, + 'masks': args.masks, + 'weight_decay': args.weight_decay, + 'dilation': args.dilation, + 'position_embedding': args.position_embedding, + 'loss_function': args.loss_function, + 'chunk_size': args.chunk_size, + 'camera_names': camera_names, + 'use_depth_image': args.use_depth_image, + 'use_robot_base': args.use_robot_base, + 'kl_weight': args.kl_weight, + 'hidden_dim': args.hidden_dim, + 'dim_feedforward': args.dim_feedforward, + 'enc_layers': args.enc_layers, + 'dec_layers': args.dec_layers, + 'nheads': args.nheads, + 'dropout': args.dropout, + 'pre_norm': args.pre_norm, + 'pretrain_backbone_path': args.pretrain_backbone_path, + } + elif args.policy_class == 'CNNMLP': + policy_config = {'lr': args.lr, + 'lr_backbone': args.lr_backbone, + 'backbone': args.backbone, + 'masks': args.masks, + 'weight_decay': args.weight_decay, + 'dilation': args.dilation, + 'position_embedding': args.position_embedding, + 'loss_function': args.loss_function, + 'chunk_size': 1, + 'camera_names': camera_names, + 'use_depth_image': args.use_depth_image, + 'use_robot_base': args.use_robot_base, + 'hidden_dim': args.hidden_dim, + 'pretrain_backbone_path': args.pretrain_backbone_path, + } + elif args.policy_class == 'Diffusion': + policy_config = {'lr': args.lr, + 'lr_backbone': args.lr_backbone, + 'backbone': args.backbone, + 'masks': args.masks, + 'weight_decay': args.weight_decay, + 'dilation': args.dilation, + 'position_embedding': args.position_embedding, + 'loss_function': args.loss_function, + 'chunk_size': args.chunk_size, + 'camera_names': camera_names, + 'use_depth_image': args.use_depth_image, + 'use_robot_base': args.use_robot_base, + 'observation_horizon': args.observation_horizon, + 'action_horizon': args.action_horizon, + 'num_inference_timesteps': args.num_inference_timesteps, + 'ema_power': args.ema_power, + 'hidden_dim': args.hidden_dim, + 'pretrain_backbone_path': args.pretrain_backbone_path, + } + else: + raise NotImplementedError + + # c. Create the policy + print(f"Loading checkpoint from: {args.ckpt_path}") + policy = make_policy(args.policy_class, policy_config, args.ckpt_path) + + # d. Move policy to GPU + policy.cuda() + + # e. Set to eval mode + policy.eval() + print("Policy loaded and in eval mode.") + + # f. Create dummy input tensors + batch_size = 1 + num_cam = len(camera_names) + image_data = torch.randn(batch_size, num_cam, 3, 480, 640).cuda() + qpos_data = torch.randn(batch_size, args.state_dim).cuda() + + depth_image_data = None + if args.use_depth_image: + depth_image_data = torch.randn(batch_size, num_cam, 1, 480, 640).cuda() + + print("Dummy data created and moved to GPU.") + + # g. Measure GPU memory before inference + torch.cuda.reset_peak_memory_stats() + start_mem = torch.cuda.memory_allocated() + print(f"Initial memory allocated: {start_mem / 1024**2:.2f} MB") + + # h. Perform inference + with torch.no_grad(): + print("Running forward pass...") + action = policy(image_data, depth_image_data, qpos_data) + print("Forward pass completed.") + + # i. Measure GPU memory after inference + end_mem = torch.cuda.memory_allocated() + peak_mem = torch.cuda.max_memory_allocated() + + print(f"Final memory allocated: {end_mem / 1024**2:.2f} MB") + print(f"Peak memory during inference: {peak_mem / 1024**2:.2f} MB") + print(f"Memory consumed by forward pass: {(peak_mem - start_mem) / 1024**2:.2f} MB") + if isinstance(action, tuple): + print(f"Output action shape: {action[0].shape}") + else: + print(f"Output action shape: {action.shape}") + + +def main(): + parser = argparse.ArgumentParser("Test Inference Memory", parents=[get_inference_args_parser()]) + parser.add_argument('--ckpt_path', action='store', type=str, help='path to checkpoint', required=True) + args = parser.parse_args() + test(args) + +def get_inference_args_parser(): + parser = argparse.ArgumentParser(add_help=False) + # Remove arguments that are not needed for inference testing + # and set sensible defaults. + parser.add_argument('--dataset_dir', action='store', type=str, help='dataset_dir', default='./dataset') + parser.add_argument('--task_name', action='store', type=str, help='task_name', default='aloha_mobile_dummy') + parser.add_argument('--policy_class', action='store', type=str, help='policy_class, capitalize, CNNMLP, ACT, Diffusion', default='ACT') + + # Model parameters + parser.add_argument('--kl_weight', action='store', type=int, help='KL Weight', default=10) + parser.add_argument('--chunk_size', action='store', type=int, help='chunk_size', default=32) + parser.add_argument('--hidden_dim', action='store', type=int, help='hidden_dim', default=512) + parser.add_argument('--dim_feedforward', action='store', type=int, help='dim_feedforward', default=3200) + parser.add_argument('--state_dim', action='store', type=int, help='state_dim', default=14) + parser.add_argument('--lr_backbone', action='store', type=float, help='lr_backbone', default=1e-5) + parser.add_argument('--backbone', action='store', type=str, help='backbone', default='resnet18') + parser.add_argument('--loss_function', action='store', type=str, help='loss_function l1 l2 l1+l2', default='l1') + parser.add_argument('--enc_layers', action='store', type=int, help='enc_layers', default=4) + parser.add_argument('--dec_layers', action='store', type=int, help='dec_layers', default=7) + parser.add_argument('--nheads', action='store', type=int, help='nheads', default=8) + parser.add_argument('--dropout', default=0.1, type=float, help="Dropout applied in the transformer") + parser.add_argument('--pre_norm', action='store_true') + parser.add_argument('--lr', action='store', type=float, help='lr', default=1e-5) + parser.add_argument('--weight_decay', type=float, help='weight_decay', default=1e-4) + parser.add_argument('--dilation', action='store_true', + help="If true, we replace stride with dilation in the last convolutional block (DC5)") + parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'), + help="Type of positional embedding to use on top of the image features") + parser.add_argument('--masks', action='store_true', + help="Train segmentation head if the flag is provided") + parser.add_argument('--pretrain_backbone_path', action='store', type=str, help='pretrain_backbone_path', default='') + + # for Diffusion + parser.add_argument('--observation_horizon', action='store', type=int, help='observation_horizon', default=1) + parser.add_argument('--action_horizon', action='store', type=int, help='action_horizon', default=8) + parser.add_argument('--num_inference_timesteps', action='store', type=int, help='num_inference_timesteps', default=10) + parser.add_argument('--ema_power', action='store', type=float, help='ema_power', default=0.75) # Changed type to float + + parser.add_argument('--use_robot_base', action='store', type=bool, help='use_robot_base', default=False) + parser.add_argument('--use_depth_image', action='store', type=bool, help='use_depth_image', default=False) + + return parser + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/aloha-devel/act/train.py b/aloha-devel/act/train.py new file mode 100644 index 0000000000000000000000000000000000000000..1d9e3b0ba286e0f910dd928b171ad2f4e57efceb --- /dev/null +++ b/aloha-devel/act/train.py @@ -0,0 +1,376 @@ +import torch +import numpy as np +import os +import pickle +import argparse +import matplotlib.pyplot as plt +from copy import deepcopy +from tqdm import tqdm + +from utils import load_data +from utils import compute_dict_mean, set_seed, detach_dict +from policy import ACTPolicy, CNNMLPPolicy, DiffusionPolicy + +import sys +sys.path.append("./") + + +def train(args): + set_seed(1) + + DATA_DIR = os.path.expanduser(args.dataset_dir) + + TASK_CONFIGS = { + args.task_name: { + 'dataset_dir': os.path.join(DATA_DIR, args.task_name), + 'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist'], + 'num_episodes': args.num_episodes + } + } + + task_config = TASK_CONFIGS[args.task_name] + + dataset_dir = task_config['dataset_dir'] + num_episodes = task_config['num_episodes'] + camera_names = task_config['camera_names'] + + # fixed parameters + if args.policy_class == 'ACT': + policy_config = {'lr': args.lr, + 'lr_backbone': args.lr_backbone, + 'backbone': args.backbone, + 'masks': args.masks, + 'weight_decay': args.weight_decay, + 'dilation': args.dilation, + 'position_embedding': args.position_embedding, + 'loss_function': args.loss_function, + 'chunk_size': args.chunk_size, # chunk_size + 'camera_names': camera_names, + 'use_depth_image': args.use_depth_image, + 'use_robot_base': args.use_robot_base, + 'kl_weight': args.kl_weight, # kl + 'hidden_dim': args.hidden_dim, # Hidden dim + 'dim_feedforward': args.dim_feedforward, + 'enc_layers': args.enc_layers, + 'dec_layers': args.dec_layers, + 'nheads': args.nheads, + 'dropout': args.dropout, + 'pre_norm': args.pre_norm, + 'pretrain_backbone_path': args.pretrain_backbone_path, + } + elif args.policy_class == 'CNNMLP': + policy_config = {'lr': args.lr, + 'lr_backbone': args.lr_backbone, + 'backbone': args.backbone, + 'masks': args.masks, + 'weight_decay': args.weight_decay, + 'dilation': args.dilation, + 'position_embedding': args.position_embedding, + 'loss_function': args.loss_function, + 'chunk_size': 1, # 查询 + 'camera_names': camera_names, + 'use_depth_image': args.use_depth_image, + 'use_robot_base': args.use_robot_base, + 'hidden_dim': args.hidden_dim, + 'pretrain_backbone_path': args.pretrain_backbone_path, + } + elif args.policy_class == 'Diffusion': + policy_config = {'lr': args.lr, + 'lr_backbone': args.lr_backbone, + 'backbone': args.backbone, + 'masks': args.masks, + 'weight_decay': args.weight_decay, + 'dilation': args.dilation, + 'position_embedding': args.position_embedding, + 'loss_function': args.loss_function, + 'chunk_size': args.chunk_size, # 查询 + 'camera_names': camera_names, + 'use_depth_image': args.use_depth_image, + 'use_robot_base': args.use_robot_base, + 'observation_horizon': args.observation_horizon, + 'action_horizon': args.action_horizon, + 'num_inference_timesteps': args.num_inference_timesteps, + 'ema_power': args.ema_power, + 'hidden_dim': args.hidden_dim, + 'pretrain_backbone_path': args.pretrain_backbone_path, + } + else: + raise NotImplementedError + + config = { + 'num_epochs': args.num_epochs, + 'ckpt_dir': args.ckpt_dir, + 'policy_class': args.policy_class, + 'policy_config': policy_config, + 'seed': args.seed, + 'pretrain_ckpt_dir': args.pretrain_ckpt, + 'ckpt_save_interval': args.ckpt_save_interval, + 'plot_interval': args.plot_interval, + 'lr_decay_start_epoch': args.lr_decay_start_epoch, + 'min_lr': args.min_lr, + } + + # data Preprocess + train_dataloader, val_dataloader, stats, _ = load_data(dataset_dir, num_episodes, args.arm_delay_time, + args.use_depth_image, args.use_robot_base, camera_names, + args.batch_size, args.batch_size) + + # save dataset stats + if not os.path.isdir(args.ckpt_dir): + os.makedirs(args.ckpt_dir) + stats_path = os.path.join(args.ckpt_dir, args.ckpt_stats_name) + with open(stats_path, 'wb') as f: + pickle.dump(stats, f) + + best_ckpt_info = train_process(train_dataloader, val_dataloader, config, stats) + best_epoch, min_val_loss, best_state_dict = best_ckpt_info + + # save best checkpoint + ckpt_path = os.path.join(args.ckpt_dir, args.ckpt_name) + torch.save(best_state_dict, ckpt_path) + print(f'Best ckpt, val loss {min_val_loss:.6f} @ epoch{best_epoch}') + + +def make_policy(policy_class, policy_config, pretrain_ckpt_dir): + if policy_class == 'ACT': + policy = ACTPolicy(policy_config) + if len(pretrain_ckpt_dir) != 0: + state_dict = torch.load(pretrain_ckpt_dir) + + loading_status = policy.deserialize(state_dict) + if not loading_status: + print("ckpt path not exist") + elif policy_class == 'CNNMLP': + policy = CNNMLPPolicy(policy_config) + if len(pretrain_ckpt_dir) != 0: + loading_status = policy.deserialize(torch.load(pretrain_ckpt_dir)) + if not loading_status: + print("ckpt path not exist") + elif policy_class == 'Diffusion': + policy = DiffusionPolicy(policy_config) + if len(pretrain_ckpt_dir) != 0: + loading_status = policy.deserialize(torch.load(pretrain_ckpt_dir)) + if not loading_status: + print("ckpt path not exist") + else: + raise NotImplementedError + return policy + + +def make_optimizer(policy_class, policy): + if policy_class == 'ACT': + optimizer = policy.configure_optimizers() + elif policy_class == 'CNNMLP': + optimizer = policy.configure_optimizers() + elif policy_class == 'Diffusion': + optimizer = policy.configure_optimizers() + else: + raise NotImplementedError + return optimizer + + +def forward_pass(policy_config, data, policy): + image_data, image_depth_data, qpos_data, action_data, action_is_pad = data + (image_data, qpos_data, action_data, action_is_pad) = (image_data.cuda(), qpos_data.cuda(), + action_data.cuda(), action_is_pad.cuda()) + if policy_config['use_depth_image']: + image_depth_data = image_depth_data.cuda() + else: + image_depth_data = None + return policy(image_data, image_depth_data, qpos_data, action_data, action_is_pad) + + +def train_process(train_dataloader, val_dataloader, config, stats): + post_process = lambda a: a * stats['qpos_std'] + stats['qpos_mean'] + num_epochs = config['num_epochs'] + ckpt_dir = config['ckpt_dir'] + seed = config['seed'] + policy_class = config['policy_class'] + policy_config = config['policy_config'] + pretrain_ckpt_dir = config['pretrain_ckpt_dir'] + ckpt_save_interval = config.get('ckpt_save_interval', 100) + plot_interval = config.get('plot_interval', 100) + lr_decay_start_epoch = config.get('lr_decay_start_epoch', num_epochs) + min_lr = config.get('min_lr', 1e-6) + set_seed(seed) + + policy = make_policy(policy_class, policy_config, pretrain_ckpt_dir) + policy.cuda() + optimizer = make_optimizer(policy_class, policy) + + train_history = [] + validation_history = [] + min_val_loss = np.inf + best_ckpt_info = None + + original_lr = policy_config['lr'] + original_lr_backbone = policy_config['lr_backbone'] + len_train_loader = len(train_dataloader) + lr_decay_start_step = lr_decay_start_epoch * len_train_loader + total_steps = num_epochs * len_train_loader + total_decay_steps = total_steps - lr_decay_start_step + + for epoch in tqdm(range(num_epochs)): + print(f'\nEpoch {epoch}') + # validation + with torch.inference_mode(): + policy.eval() + epoch_dicts = [] + for batch_idx, data in enumerate(val_dataloader): + forward_dict, result = forward_pass(policy_config, data, policy) + # print("result:", post_process(result.cpu().detach().numpy())[0, :, 7:]) + epoch_dicts.append(forward_dict) + epoch_summary = compute_dict_mean(epoch_dicts) + validation_history.append(epoch_summary) + + epoch_val_loss = epoch_summary['loss'] + if epoch_val_loss < min_val_loss: + min_val_loss = epoch_val_loss + best_ckpt_info = (epoch, min_val_loss, deepcopy(policy.serialize())) + print(f'Val loss: {epoch_val_loss:.5f}') + summary_string = '' + for k, v in epoch_summary.items(): + summary_string += f'{k}: {v.item():.3f} ' + print(summary_string) + + # training + policy.train() + optimizer.zero_grad() + for batch_idx, data in enumerate(train_dataloader): + current_step = epoch * len_train_loader + batch_idx + if current_step >= lr_decay_start_step and total_decay_steps > 0: + decay_progress = (current_step - lr_decay_start_step) / total_decay_steps + decay_factor = 1.0 - decay_progress + + new_lr = max(original_lr * decay_factor, min_lr) + new_lr_backbone = max(original_lr_backbone * decay_factor, min_lr) + + optimizer.param_groups[0]['lr'] = new_lr + optimizer.param_groups[1]['lr'] = new_lr_backbone + + # debug + # print(optimizer.param_groups[0]['lr']) + forward_dict, result = forward_pass(policy_config, data, policy) + # print("result:", post_process(result.cpu().detach().numpy())[0, :, 7:]) + # backward + loss = forward_dict['loss'] + loss.backward() + optimizer.step() + optimizer.zero_grad() + train_history.append(detach_dict(forward_dict)) + epoch_summary = compute_dict_mean(train_history[(batch_idx+1)*epoch:(batch_idx+1)*(epoch+1)]) + epoch_train_loss = epoch_summary['loss'] + print(f'Train loss: {epoch_train_loss:.5f}') + summary_string = '' + for k, v in epoch_summary.items(): + summary_string += f'{k}: {v.item():.3f} ' + print(summary_string) + + if epoch % ckpt_save_interval == 0: + ckpt_path = os.path.join(ckpt_dir, f'policy_epoch_{epoch}_seed_{seed}.ckpt') + torch.save(policy.serialize(), ckpt_path) + + if epoch > 0 and epoch % plot_interval == 0: + plot_history(train_history, validation_history, epoch, ckpt_dir, seed) + + ckpt_path = os.path.join(ckpt_dir, f'policy_last.ckpt') + torch.save(policy.serialize(), ckpt_path) + + best_epoch, min_val_loss, best_state_dict = best_ckpt_info + ckpt_path = os.path.join(ckpt_dir, f'policy_epoch_{best_epoch}_seed_{seed}.ckpt') + torch.save(best_state_dict, ckpt_path) + print(f'Training finished:\nSeed {seed}, val loss {min_val_loss:.6f} at epoch {best_epoch}') + + # save training curves + plot_history(train_history, validation_history, num_epochs, ckpt_dir, seed) + + return best_ckpt_info + + +def plot_history(train_history, validation_history, num_epochs, ckpt_dir, seed): + # save training curves + for key in train_history[0]: + plot_path = os.path.join(ckpt_dir, f'train_val_{key}_seed_{seed}.png') + plt.figure() + train_values = [summary[key].item() for summary in train_history] + val_values = [summary[key].item() for summary in validation_history] + plt.plot(np.linspace(0, num_epochs-1, len(train_history)), train_values, label='train') + plt.plot(np.linspace(0, num_epochs-1, len(validation_history)), val_values, label='validation') + # plt.ylim([-0.1, 1]) + plt.tight_layout() + plt.legend() + plt.title(key) + plt.savefig(plot_path) + print(f'Saved plots to {ckpt_dir}') + + +def get_arguments(): + parser = argparse.ArgumentParser() + parser.add_argument('--dataset_dir', action='store', type=str, help='dataset_dir', default='./dataset', required=True) + parser.add_argument('--ckpt_dir', action='store', type=str, help='ckpt_dir', required=True) + parser.add_argument('--num_episodes', action='store', type=int, help='num_episodes', required=True) + + parser.add_argument('--pretrain_ckpt', action='store', type=str, help='pretrain_ckpt', default='', required=False) + parser.add_argument('--pretrain_backbone_path', action='store', type=str, help='pretrain_backbone_path', default='', required=False) + parser.add_argument('--task_name', action='store', type=str, help='task_name', default='aloha_mobile_dummy', required=False) + + parser.add_argument('--ckpt_name', action='store', type=str, help='ckpt_name', default='policy_best.ckpt', required=False) + parser.add_argument('--ckpt_stats_name', action='store', type=str, help='ckpt_stats_name', default='dataset_stats.pkl', required=False) + parser.add_argument('--policy_class', action='store', type=str, help='policy_class, capitalize, CNNMLP, ACT, Diffusion', default='ACT', required=False) + parser.add_argument('--batch_size', action='store', type=int, help='batch_size', default=32, required=False) + parser.add_argument('--seed', action='store', type=int, help='seed', default=0, required=False) + parser.add_argument('--num_epochs', action='store', type=int, help='num_epochs', default=3000, required=False) + parser.add_argument('--ckpt_save_interval', action='store', type=int, help='ckpt_save_interval', default=100, required=False) + parser.add_argument('--plot_interval', action='store', type=int, help='plot_interval', default=100, required=False) + + parser.add_argument('--lr', action='store', type=float, help='lr', default=4e-5, required=False) + parser.add_argument('--lr_decay_start_epoch', action='store', type=int, help='epoch to start LR decay', default=3000, required=False) + parser.add_argument('--min_lr', action='store', type=float, help='minimum learning rate for decay', default=1e-6, required=False) + parser.add_argument('--weight_decay', type=float, help='weight_decay', default=1e-4, required=False) + parser.add_argument('--dilation', action='store_true', + help="If true, we replace stride with dilation in the last convolutional block (DC5)", required=False) + parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'), + help="Type of positional embedding to use on top of the image features", required=False) + parser.add_argument('--masks', action='store_true', + help="Train segmentation head if the flag is provided") + + parser.add_argument('--state_dim', action='store', type=int, help='state_dim', default=14, required=False) + parser.add_argument('--lr_backbone', action='store', type=float, help='lr_backbone', default=4e-5, required=False) + parser.add_argument('--backbone', action='store', type=str, help='backbone', default='resnet18', required=False) + parser.add_argument('--loss_function', action='store', type=str, help='loss_function l1 l2 l1+l2', default='l1', required=False) + parser.add_argument('--enc_layers', action='store', type=int, help='enc_layers', default=4, required=False) + parser.add_argument('--dec_layers', action='store', type=int, help='dec_layers', default=7, required=False) + parser.add_argument('--nheads', action='store', type=int, help='nheads', default=8, required=False) + parser.add_argument('--dropout', default=0.1, type=float, help="Dropout applied in the transformer", required=False) + parser.add_argument('--pre_norm', action='store_true', required=False) + + # for ACT + parser.add_argument('--kl_weight', action='store', type=int, help='KL Weight', default=10, required=False) + parser.add_argument('--chunk_size', action='store', type=int, help='chunk_size', default=32, required=False) + parser.add_argument('--hidden_dim', action='store', type=int, help='hidden_dim', default=512, required=False) + parser.add_argument('--dim_feedforward', action='store', type=int, help='dim_feedforward', default=3200, required=False) + parser.add_argument('--temporal_agg', action='store', type=bool, help='temporal_agg', default=True, required=False) + + # for Diffusion + parser.add_argument('--observation_horizon', action='store', type=int, help='observation_horizon', default=1, required=False) + parser.add_argument('--action_horizon', action='store', type=int, help='action_horizon', default=8, required=False) + parser.add_argument('--num_inference_timesteps', action='store', type=int, help='num_inference_timesteps', default=10, required=False) + parser.add_argument('--ema_power', action='store', type=int, help='ema_power', default=0.75, required=False) + + parser.add_argument('--use_robot_base', action='store', type=bool, help='use_robot_base', default=False, required=False) + + parser.add_argument('--arm_delay_time', action='store', type=int, help='arm_delay_time', default=0, required=False) + + parser.add_argument('--use_depth_image', action='store', type=bool, help='use_depth_image', default=False, required=False) + + args = parser.parse_args() + return args + + +def main(): + args = get_arguments() + train(args) + +if __name__ == '__main__': + main() +# python act/train.py --dataset_dir ~/data --pretrain_ckpt policy_best.ckpt --ckpt_dir ~/train_dir/ --num_episodes 20 --batch_size 10 --num_epochs 2000 \ No newline at end of file diff --git a/aloha-devel/robomimic/__pycache__/macros.cpython-38.pyc b/aloha-devel/robomimic/__pycache__/macros.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb1386fcde9feeb06a96577c79e55e37f97a02e4 Binary files /dev/null and b/aloha-devel/robomimic/__pycache__/macros.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/algo/__init__.py b/aloha-devel/robomimic/algo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5fd1d53f90c78118a7ae9f8c7d6a55a4a4771554 --- /dev/null +++ b/aloha-devel/robomimic/algo/__init__.py @@ -0,0 +1,13 @@ +from robomimic.algo.algo import register_algo_factory_func, algo_name_to_factory_func, algo_factory, Algo, PolicyAlgo, ValueAlgo, PlannerAlgo, HierarchicalAlgo, RolloutPolicy + +# note: these imports are needed to register these classes in the global algo registry +from robomimic.algo.bc import BC, BC_Gaussian, BC_GMM, BC_VAE, BC_RNN, BC_RNN_GMM +from robomimic.algo.bcq import BCQ, BCQ_GMM, BCQ_Distributional +from robomimic.algo.cql import CQL +from robomimic.algo.iql import IQL +from robomimic.algo.gl import GL, GL_VAE, ValuePlanner +from robomimic.algo.hbc import HBC +from robomimic.algo.iris import IRIS +from robomimic.algo.td3_bc import TD3_BC +from robomimic.algo.diffusion_policy import DiffusionPolicyUNet +from robomimic.algo.act import ACT diff --git a/aloha-devel/robomimic/algo/__pycache__/__init__.cpython-38.pyc b/aloha-devel/robomimic/algo/__pycache__/__init__.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b24d5c92450e1c2312d6933c18c9b6044cb58407 Binary files /dev/null and b/aloha-devel/robomimic/algo/__pycache__/__init__.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/algo/__pycache__/bc.cpython-38.pyc b/aloha-devel/robomimic/algo/__pycache__/bc.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a7903bb7438b1589866a27c228e426d6b2f79a8 Binary files /dev/null and b/aloha-devel/robomimic/algo/__pycache__/bc.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/algo/act.py b/aloha-devel/robomimic/algo/act.py new file mode 100644 index 0000000000000000000000000000000000000000..8f35271f3fa3e04d6c0062a6d184d4bdbd54416c --- /dev/null +++ b/aloha-devel/robomimic/algo/act.py @@ -0,0 +1,247 @@ +""" +Implementation of Action Chunking with Transformers (ACT). +""" +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision.transforms as transforms + +import robomimic.utils.tensor_utils as TensorUtils + +from robomimic.algo import register_algo_factory_func, PolicyAlgo +from robomimic.algo.bc import BC_VAE + + +@register_algo_factory_func("act") +def algo_config_to_class(algo_config): + """ + Maps algo config to the BC algo class to instantiate, along with additional algo kwargs. + + Args: + algo_config (Config instance): algo config + + Returns: + algo_class: subclass of Algo + algo_kwargs (dict): dictionary of additional kwargs to pass to algorithm + """ + algo_class, algo_kwargs = ACT, {} + + return algo_class, algo_kwargs + + +class ACT(BC_VAE): + """ + BC training with a VAE policy. + """ + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + + self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + + self.nets = nn.ModuleDict() + self.chunk_size = self.global_config["train"]["seq_length"] + self.camera_keys = self.obs_config['modalities']['obs']['rgb'].copy() + self.proprio_keys = self.obs_config['modalities']['obs']['low_dim'].copy() + self.obs_keys = self.proprio_keys + self.camera_keys + + self.proprio_dim = 0 + for k in self.proprio_keys: + self.proprio_dim += self.obs_key_shapes[k][0] + + from act.detr.main import build_ACT_model_and_optimizer + policy_config = {'num_queries': self.chunk_size, + 'hidden_dim': self.algo_config.act.hidden_dim, + 'dim_feedforward': self.algo_config.act.dim_feedforward, + 'backbone': self.algo_config.act.backbone, + 'enc_layers': self.algo_config.act.enc_layers, + 'dec_layers': self.algo_config.act.dec_layers, + 'nheads': self.algo_config.act.nheads, + 'latent_dim': self.algo_config.act.latent_dim, + 'a_dim': self.ac_dim, + 'state_dim': self.proprio_dim, + 'camera_names': self.camera_keys + } + self.kl_weight = self.algo_config.act.kl_weight + model, optimizer = build_ACT_model_and_optimizer(policy_config) + self.nets["policy"] = model + self.nets = self.nets.float().to(self.device) + + self.temporal_agg = False + self.query_frequency = self.chunk_size # TODO maybe tune + + self._step_counter = 0 + self.a_hat_store = None + + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + + input_batch = dict() + input_batch["obs"] = {k: batch["obs"][k][:, 0, :] for k in batch["obs"] if k != 'pad_mask'} + input_batch["obs"]['pad_mask'] = batch["obs"]['pad_mask'] + input_batch["goal_obs"] = batch.get("goal_obs", None) # goals may not be present + input_batch["actions"] = batch["actions"][:, :, :] + # we move to device first before float conversion because image observation modalities will be uint8 - + # this minimizes the amount of data transferred to GPU + return TensorUtils.to_float(TensorUtils.to_device(input_batch, self.device)) + + def train_on_batch(self, batch, epoch, validate=False): + """ + Update from superclass to set categorical temperature, for categorcal VAEs. + """ + + return super(BC_VAE, self).train_on_batch(batch, epoch, validate=validate) + + def _forward_training(self, batch): + """ + Internal helper function for BC algo class. Compute forward pass + and return network outputs in @predictions dict. + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + Returns: + predictions (dict): dictionary containing network outputs + """ + + proprio = [batch["obs"][k] for k in self.proprio_keys] + proprio = torch.cat(proprio, axis=1) + qpos = proprio + + images = [] + for cam_name in self.camera_keys: + image = batch['obs'][cam_name] + image = self.normalize(image) + image = image.unsqueeze(axis=1) + images.append(image) + images = torch.cat(images, axis=1) + + env_state = torch.zeros([qpos.shape[0], 10]).cuda() # this is not used + + actions = batch['actions'] + is_pad = batch['obs']['pad_mask'] == 0 # from 1.0 or 0 to False and True + is_pad = is_pad.squeeze(dim=-1) + + a_hat, is_pad_hat, (mu, logvar) = self.nets["policy"](qpos, images, env_state, actions, is_pad) + total_kld, dim_wise_kld, mean_kld = self.kl_divergence(mu, logvar) + loss_dict = dict() + all_l1 = F.l1_loss(actions, a_hat, reduction='none') + l1 = (all_l1 * ~is_pad.unsqueeze(-1)).mean() + loss_dict['l1'] = l1 + loss_dict['kl'] = total_kld[0] + + + predictions = OrderedDict( + actions=actions, + kl_loss=loss_dict['kl'], + reconstruction_loss=loss_dict['l1'], + ) + + return predictions + + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + Returns: + action (torch.Tensor): action tensor + """ + assert not self.nets.training + + proprio = [obs_dict[k] for k in self.proprio_keys] + proprio = torch.cat(proprio, axis=1) + qpos = proprio + + images = [] + for cam_name in self.camera_keys: + image = obs_dict[cam_name] + image = self.normalize(image) + image = image.unsqueeze(axis=1) + images.append(image) + images = torch.cat(images, axis=1) + + env_state = torch.zeros([qpos.shape[0], 10]).cuda() # not used + + if self._step_counter % self.query_frequency == 0: + a_hat, is_pad_hat, (mu, logvar) = self.nets["policy"](qpos, images, env_state) + self.a_hat_store = a_hat + + action = self.a_hat_store[:, self._step_counter % self.query_frequency, :] + self._step_counter += 1 + return action + + + def reset(self): + """ + Reset algo state to prepare for environment rollouts. + """ + self._step_counter = 0 + + def _compute_losses(self, predictions, batch): + """ + Internal helper function for BC algo class. Compute losses based on + network outputs in @predictions dict, using reference labels in @batch. + Args: + predictions (dict): dictionary containing network outputs, from @_forward_training + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + Returns: + losses (dict): dictionary of losses computed over the batch + """ + + # total loss is sum of reconstruction and KL, weighted by beta + kl_loss = predictions["kl_loss"] + recons_loss = predictions["reconstruction_loss"] + action_loss = recons_loss + self.kl_weight * kl_loss + return OrderedDict( + recons_loss=recons_loss, + kl_loss=kl_loss, + action_loss=action_loss, + ) + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + Args: + info (dict): dictionary of info + Returns: + loss_log (dict): name -> summary statistic + """ + log = PolicyAlgo.log_info(self, info) + log["Loss"] = info["losses"]["action_loss"].item() + log["KL_Loss"] = info["losses"]["kl_loss"].item() + log["Reconstruction_Loss"] = info["losses"]["recons_loss"].item() + if "policy_grad_norms" in info: + log["Policy_Grad_Norms"] = info["policy_grad_norms"] + return log + + def kl_divergence(self, mu, logvar): + batch_size = mu.size(0) + assert batch_size != 0 + if mu.data.ndimension() == 4: + mu = mu.view(mu.size(0), mu.size(1)) + if logvar.data.ndimension() == 4: + logvar = logvar.view(logvar.size(0), logvar.size(1)) + + klds = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()) + total_kld = klds.sum(1).mean(0, True) + dimension_wise_kld = klds.mean(0) + mean_kld = klds.mean(1).mean(0, True) + + return total_kld, dimension_wise_kld, mean_kld + diff --git a/aloha-devel/robomimic/algo/algo.py b/aloha-devel/robomimic/algo/algo.py new file mode 100644 index 0000000000000000000000000000000000000000..4cd8cdebd243e05c0a959b76ffc572ae92a7860c --- /dev/null +++ b/aloha-devel/robomimic/algo/algo.py @@ -0,0 +1,674 @@ +""" +This file contains base classes that other algorithm classes subclass. +Each algorithm file also implements a algorithm factory function that +takes in an algorithm config (`config.algo`) and returns the particular +Algo subclass that should be instantiated, along with any extra kwargs. +These factory functions are registered into a global dictionary with the +@register_algo_factory_func function decorator. This makes it easy for +@algo_factory to instantiate the correct `Algo` subclass. +""" +import textwrap +from copy import deepcopy +from collections import OrderedDict + +import torch.nn as nn +import torch +import os +import numpy as np +import imageio + +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.torch_utils as TorchUtils +import robomimic.utils.obs_utils as ObsUtils +import robomimic.utils.action_utils as AcUtils +import robomimic.utils.vis_utils as VisUtils + +from torch.utils.data import DataLoader + +# mapping from algo name to factory functions that map algo configs to algo class names +REGISTERED_ALGO_FACTORY_FUNCS = OrderedDict() + + +def register_algo_factory_func(algo_name): + """ + Function decorator to register algo factory functions that map algo configs to algo class names. + Each algorithm implements such a function, and decorates it with this decorator. + + Args: + algo_name (str): the algorithm name to register the algorithm under + """ + def decorator(factory_func): + REGISTERED_ALGO_FACTORY_FUNCS[algo_name] = factory_func + return decorator + + +def algo_name_to_factory_func(algo_name): + """ + Uses registry to retrieve algo factory function from algo name. + + Args: + algo_name (str): the algorithm name + """ + return REGISTERED_ALGO_FACTORY_FUNCS[algo_name] + + +def algo_factory(algo_name, config, obs_key_shapes, ac_dim, device): + """ + Factory function for creating algorithms based on the algorithm name and config. + + Args: + algo_name (str): the algorithm name + + config (BaseConfig instance): config object + + obs_key_shapes (OrderedDict): dictionary that maps observation keys to shapes + + ac_dim (int): dimension of action space + + device (torch.Device): where the algo should live (i.e. cpu, gpu) + """ + + # @algo_name is included as an arg to be explicit, but make sure it matches the config + assert algo_name == config.algo_name + + # use algo factory func to get algo class and kwargs from algo config + factory_func = algo_name_to_factory_func(algo_name) + algo_cls, algo_kwargs = factory_func(config.algo) + + # create algo instance + return algo_cls( + algo_config=config.algo, + obs_config=config.observation, + global_config=config, + obs_key_shapes=obs_key_shapes, + ac_dim=ac_dim, + device=device, + **algo_kwargs + ) + + +class Algo(object): + """ + Base algorithm class that all other algorithms subclass. Defines several + functions that should be overriden by subclasses, in order to provide + a standard API to be used by training functions such as @run_epoch in + utils/train_utils.py. + """ + def __init__( + self, + algo_config, + obs_config, + global_config, + obs_key_shapes, + ac_dim, + device + ): + """ + Args: + algo_config (Config object): instance of Config corresponding to the algo section + of the config + + obs_config (Config object): instance of Config corresponding to the observation + section of the config + + global_config (Config object): global training config + + obs_key_shapes (OrderedDict): dictionary that maps observation keys to shapes + + ac_dim (int): dimension of action space + + device (torch.Device): where the algo should live (i.e. cpu, gpu) + """ + self.optim_params = deepcopy(algo_config.optim_params) + self.algo_config = algo_config + self.obs_config = obs_config + self.global_config = global_config + + self.ac_dim = ac_dim + self.device = device + self.obs_key_shapes = obs_key_shapes + + self.nets = nn.ModuleDict() + self._create_shapes(obs_config.modalities, obs_key_shapes) + self._create_networks() + self._create_optimizers() + assert isinstance(self.nets, nn.ModuleDict) + + def _create_shapes(self, obs_keys, obs_key_shapes): + """ + Create obs_shapes, goal_shapes, and subgoal_shapes dictionaries, to make it + easy for this algorithm object to keep track of observation key shapes. Each dictionary + maps observation key to shape. + + Args: + obs_keys (dict): dict of required observation keys for this training run (usually + specified by the obs config), e.g., {"obs": ["rgb", "proprio"], "goal": ["proprio"]} + obs_key_shapes (dict): dict of observation key shapes, e.g., {"rgb": [3, 224, 224]} + """ + # determine shapes + self.obs_shapes = OrderedDict() + self.goal_shapes = OrderedDict() + self.subgoal_shapes = OrderedDict() + + # We check across all modality groups (obs, goal, subgoal), and see if the inputted observation key exists + # across all modalitie specified in the config. If so, we store its corresponding shape internally + for k in obs_key_shapes: + if "obs" in self.obs_config.modalities and k in [obs_key for modality in self.obs_config.modalities.obs.values() for obs_key in modality]: + self.obs_shapes[k] = obs_key_shapes[k] + if "goal" in self.obs_config.modalities and k in [obs_key for modality in self.obs_config.modalities.goal.values() for obs_key in modality]: + self.goal_shapes[k] = obs_key_shapes[k] + if "subgoal" in self.obs_config.modalities and k in [obs_key for modality in self.obs_config.modalities.subgoal.values() for obs_key in modality]: + self.subgoal_shapes[k] = obs_key_shapes[k] + + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + @self.nets should be a ModuleDict. + """ + raise NotImplementedError + + def _create_optimizers(self): + """ + Creates optimizers using @self.optim_params and places them into @self.optimizers. + """ + self.optimizers = dict() + self.lr_schedulers = dict() + + for k in self.optim_params: + # only make optimizers for networks that have been created - @optim_params may have more + # settings for unused networks + if k in self.nets: + if isinstance(self.nets[k], nn.ModuleList): + self.optimizers[k] = [ + TorchUtils.optimizer_from_optim_params(net_optim_params=self.optim_params[k], net=self.nets[k][i]) + for i in range(len(self.nets[k])) + ] + self.lr_schedulers[k] = [ + TorchUtils.lr_scheduler_from_optim_params(net_optim_params=self.optim_params[k], net=self.nets[k][i], optimizer=self.optimizers[k][i]) + for i in range(len(self.nets[k])) + ] + else: + self.optimizers[k] = TorchUtils.optimizer_from_optim_params( + net_optim_params=self.optim_params[k], net=self.nets[k]) + self.lr_schedulers[k] = TorchUtils.lr_scheduler_from_optim_params( + net_optim_params=self.optim_params[k], net=self.nets[k], optimizer=self.optimizers[k]) + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + return batch + + def postprocess_batch_for_training(self, batch, obs_normalization_stats): + """ + Does some operations (like channel swap, uint8 to float conversion, normalization) + after @process_batch_for_training is called, in order to ensure these operations + take place on GPU. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader. Assumed to be on the device where + training will occur (after @process_batch_for_training + is called) + + obs_normalization_stats (dict or None): if provided, this should map observation + keys to dicts with a "mean" and "std" of shape (1, ...) where ... is the + default shape for the observation. + + Returns: + batch (dict): postproceesed batch + """ + obs_keys = ["obs", "next_obs", "goal_obs"] + for k in obs_keys: + if k in batch and batch[k] is not None: + batch[k] = ObsUtils.process_obs_dict(batch[k]) + if obs_normalization_stats is not None: + batch[k] = ObsUtils.normalize_dict(batch[k], obs_normalization_stats=obs_normalization_stats) + return batch + + def train_on_batch(self, batch, epoch, validate=False): + """ + Training on a single batch of data. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + assert validate or self.nets.training + return OrderedDict() + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss log (dict): name -> summary statistic + """ + log = OrderedDict() + + # record current optimizer learning rates + for k in self.optimizers: + for i, param_group in enumerate(self.optimizers[k].param_groups): + log["Optimizer/{}{}_lr".format(k, i)] = param_group["lr"] + + return log + + def on_epoch_end(self, epoch): + """ + Called at the end of each epoch. + """ + + # LR scheduling updates + for k in self.lr_schedulers: + if self.lr_schedulers[k] is not None: + self.lr_schedulers[k].step() + + def set_eval(self): + """ + Prepare networks for evaluation. + """ + self.nets.eval() + + def set_train(self): + """ + Prepare networks for training. + """ + self.nets.train() + + def serialize(self): + """ + Get dictionary of current model parameters. + """ + return self.nets.state_dict() + + def deserialize(self, model_dict): + """ + Load model from a checkpoint. + + Args: + model_dict (dict): a dictionary saved by self.serialize() that contains + the same keys as @self.network_classes + """ + self.nets.load_state_dict(model_dict) + + def __repr__(self): + """ + Pretty print algorithm and network description. + """ + return "{} (\n".format(self.__class__.__name__) + \ + textwrap.indent(self.nets.__repr__(), ' ') + "\n)" + + def reset(self): + """ + Reset algo state to prepare for environment rollouts. + """ + pass + + +class PolicyAlgo(Algo): + """ + Base class for all algorithms that can be used as policies. + """ + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + action (torch.Tensor): action tensor + """ + raise NotImplementedError + + def compute_traj_pred_actual_actions(self, traj, return_images=False): + """ + traj is an R2D2Dataset object representing one trajectory + This function is slow (>1s per trajectory) because there is no batching + and instead loops through all timesteps one by one + TODO: documentation + """ + if return_images: + image_keys = [item for item in traj.__getitem__(0)['obs'].keys() if "image" in item] + images = {key: [] for key in image_keys} + else: + images = None + + dataloader = DataLoader( + dataset=traj, + sampler=None, + batch_size=1, + shuffle=False, + num_workers=1, + drop_last=True, + ) + + self.reset() + actual_actions = [] + predicted_actions = [] + + # loop through each timestep + for batch in iter(dataloader): + batch = self.process_batch_for_training(batch) + + if return_images: + for image_key in image_keys: + im = batch["obs"][image_key][0][-1] + im = TensorUtils.to_numpy(im).astype(np.uint32) + images[image_key].append(im) + + batch = self.postprocess_batch_for_training(batch, obs_normalization_stats=None) # ignore obs_normalization for now + + model_output = self.get_action(batch["obs"]) + + actual_action = TensorUtils.to_numpy( + batch["actions"][0][0] + ) + predicted_action = TensorUtils.to_numpy( + model_output[0] + ) + + actual_actions.append(actual_action) + predicted_actions.append(predicted_action) + + actual_actions = np.array(actual_actions) + predicted_actions = np.array(predicted_actions) + return actual_actions, predicted_actions, images + + def compute_mse_visualize(self, trainset, validset, num_samples, savedir=None): + """If savedir is not None, then also visualize the model predictions and save them to savedir""" + visualize = savedir is not None + + # set model into eval mode + self.set_eval() + random_state = np.random.RandomState(0) + train_indices = random_state.choice( + len(trainset.datasets), + min(len(trainset.datasets), num_samples) + ).astype(int) + training_sampled_data = [trainset.datasets[idx] for idx in train_indices] + + if validset is not None: + valid_indices = random_state.choice( + len(validset.datasets), + min(len(validset.datasets), num_samples) + ).astype(int) + validation_sampled_data = [validset.datasets[idx] for idx in valid_indices] + + inference_datasets_mapping = {"Train": training_sampled_data, "Valid": validation_sampled_data} + else: + inference_datasets_mapping = {"Train": training_sampled_data} + + # extract action name for visualization + action_keys = self.global_config.train.action_keys + training_sample=training_sampled_data[0][0] + modified_action_keys = [element.replace("action/", "") for element in action_keys] + action_names = [] + + for i, action_key in enumerate(action_keys): + if isinstance(training_sample[action_key][0], np.ndarray): + action_names.extend([f'{modified_action_keys[i]}_{j+1}' for j in range(len(training_sample[action_key][0]))]) + else: + action_names.append(modified_action_keys[i]) + + if visualize: + print("Saving model prediction plots to {}".format(savedir)) + + mse_log = {} + vis_log = {} + # loop through training and validation sets + for inference_key in inference_datasets_mapping: + actual_actions_all_traj = [] # (NxT, D) + predicted_actions_all_traj = [] # (NxT, D) + + # loop through each trajectory + traj_num = 1 + for d in inference_datasets_mapping[inference_key]: + actual_actions, predicted_actions, images = self.compute_traj_pred_actual_actions(d, return_images=visualize) + actual_actions_all_traj.append(actual_actions) + predicted_actions_all_traj.append(predicted_actions) + if visualize: + traj_key = "{}_traj_{}".format(inference_key.lower(), traj_num) + save_path = os.path.join(savedir, traj_key + ".png") + VisUtils.make_model_prediction_plot( + hdf5_path=d.hdf5_path, + save_path=save_path, + images=images, + action_names=action_names, + actual_actions=actual_actions, + predicted_actions=predicted_actions, + ) + vis_log[traj_key] = imageio.imread(save_path) + traj_num += 1 + + actual_actions_all_traj = np.concatenate(actual_actions_all_traj, axis=0) + predicted_actions_all_traj = np.concatenate(predicted_actions_all_traj, axis=0) + accuracy_thresholds = np.logspace(-3,-5, num=3).tolist() + mse = torch.nn.functional.mse_loss( + torch.tensor(predicted_actions_all_traj), + torch.tensor(actual_actions_all_traj), + reduction='none' + ) # (NxT, D) + mse_log[f'{inference_key}/action_mse_error'] = mse.mean().item() # average MSE across all timesteps averaged across all action dimensions (D,) + + # compute percentage of timesteps that have MSE less than the accuracy thresholds + for accuracy_threshold in accuracy_thresholds: + mse_log[f'{inference_key}/action_accuracy@{accuracy_threshold}'] = (torch.less(mse,accuracy_threshold).float().mean().item()) + + return mse_log, vis_log + + +class ValueAlgo(Algo): + """ + Base class for all algorithms that can learn a value function. + """ + def get_state_value(self, obs_dict, goal_dict=None): + """ + Get state value outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + value (torch.Tensor): value tensor + """ + raise NotImplementedError + + def get_state_action_value(self, obs_dict, actions, goal_dict=None): + """ + Get state-action value outputs. + + Args: + obs_dict (dict): current observation + actions (torch.Tensor): action + goal_dict (dict): (optional) goal + + Returns: + value (torch.Tensor): value tensor + """ + raise NotImplementedError + + +class PlannerAlgo(Algo): + """ + Base class for all algorithms that can be used for planning subgoals + conditioned on current observations and potential goal observations. + """ + def get_subgoal_predictions(self, obs_dict, goal_dict=None): + """ + Get predicted subgoal outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + subgoal prediction (dict): name -> Tensor [batch_size, ...] + """ + raise NotImplementedError + + def sample_subgoals(self, obs_dict, goal_dict, num_samples=1): + """ + For planners that rely on sampling subgoals. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + subgoals (dict): name -> Tensor [batch_size, num_samples, ...] + """ + raise NotImplementedError + + +class HierarchicalAlgo(Algo): + """ + Base class for all hierarchical algorithms that consist of (1) subgoal planning + and (2) subgoal-conditioned policy learning. + """ + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + action (torch.Tensor): action tensor + """ + raise NotImplementedError + + def get_subgoal_predictions(self, obs_dict, goal_dict=None): + """ + Get subgoal predictions from high-level subgoal planner. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + subgoal (dict): predicted subgoal + """ + raise NotImplementedError + + @property + def current_subgoal(self): + """ + Get the current subgoal for conditioning the low-level policy + + Returns: + current subgoal (dict): predicted subgoal + """ + raise NotImplementedError + + +class RolloutPolicy(object): + """ + Wraps @Algo object to make it easy to run policies in a rollout loop. + """ + def __init__(self, policy, obs_normalization_stats=None, action_normalization_stats=None): + """ + Args: + policy (Algo instance): @Algo object to wrap to prepare for rollouts + + obs_normalization_stats (dict): optionally pass a dictionary for observation + normalization. This should map observation keys to dicts + with a "mean" and "std" of shape (1, ...) where ... is the default + shape for the observation. + """ + self.policy = policy + self.obs_normalization_stats = obs_normalization_stats + self.action_normalization_stats = action_normalization_stats + + def start_episode(self): + """ + Prepare the policy to start a new rollout. + """ + self.policy.set_eval() + self.policy.reset() + + def _prepare_observation(self, ob, batched=False): + """ + Prepare raw observation dict from environment for policy. + + Args: + ob (dict): single observation dictionary from environment (no batch dimension, + and np.array values for each key) + + batched (bool): whether the input is already batched + """ + if self.obs_normalization_stats is not None: + ob = ObsUtils.normalize_dict(ob, obs_normalization_stats=self.obs_normalization_stats) + ob = TensorUtils.to_tensor(ob) + if not batched: + ob = TensorUtils.to_batch(ob) + ob = TensorUtils.to_device(ob, self.policy.device) + ob = TensorUtils.to_float(ob) + return ob + + def __repr__(self): + """Pretty print network description""" + return self.policy.__repr__() + + def __call__(self, ob, goal=None, batched=False): + """ + Produce action from raw observation dict (and maybe goal dict) from environment. + + Args: + ob (dict): single observation dictionary from environment (no batch dimension, + and np.array values for each key) + goal (dict): goal observation + batched (bool): whether the input is already batched + """ + ob = self._prepare_observation(ob, batched=batched) + if goal is not None: + goal = self._prepare_observation(goal, batched=batched) + ac = self.policy.get_action(obs_dict=ob, goal_dict=goal) + if not batched: + ac = ac[0] + ac = TensorUtils.to_numpy(ac) + if self.action_normalization_stats is not None: + action_keys = self.policy.global_config.train.action_keys + action_shapes = {k: self.action_normalization_stats[k]["offset"].shape[1:] for k in self.action_normalization_stats} + ac_dict = AcUtils.vector_to_action_dict(ac, action_shapes=action_shapes, action_keys=action_keys) + ac_dict = ObsUtils.unnormalize_dict(ac_dict, normalization_stats=self.action_normalization_stats) + action_config = self.policy.global_config.train.action_config + for key, value in ac_dict.items(): + this_format = action_config[key].get("format", None) + if this_format == "rot_6d": + rot_6d = torch.from_numpy(value).unsqueeze(0) + conversion_format = action_config[key].get("convert_at_runtime", "rot_axis_angle") + if conversion_format == "rot_axis_angle": + rot = TorchUtils.rot_6d_to_axis_angle(rot_6d=rot_6d).squeeze().numpy() + elif conversion_format == "rot_euler": + rot = TorchUtils.rot_6d_to_euler_angles(rot_6d=rot_6d, convention="XYZ").squeeze().numpy() + else: + raise ValueError + ac_dict[key] = rot + ac = AcUtils.action_dict_to_vector(ac_dict, action_keys=action_keys) + return ac diff --git a/aloha-devel/robomimic/algo/bc.py b/aloha-devel/robomimic/algo/bc.py new file mode 100644 index 0000000000000000000000000000000000000000..9f4cac04cde6a32c4ad7335c0fb1b359075b1a08 --- /dev/null +++ b/aloha-devel/robomimic/algo/bc.py @@ -0,0 +1,899 @@ +""" +Implementation of Behavioral Cloning (BC). +""" +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.distributions as D + +import robomimic.models.base_nets as BaseNets +import robomimic.models.obs_nets as ObsNets +import robomimic.models.policy_nets as PolicyNets +import robomimic.models.vae_nets as VAENets +import robomimic.utils.loss_utils as LossUtils +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.torch_utils as TorchUtils +import robomimic.utils.obs_utils as ObsUtils + +from robomimic.algo import register_algo_factory_func, PolicyAlgo + + +@register_algo_factory_func("bc") +def algo_config_to_class(algo_config): + """ + Maps algo config to the BC algo class to instantiate, along with additional algo kwargs. + + Args: + algo_config (Config instance): algo config + + Returns: + algo_class: subclass of Algo + algo_kwargs (dict): dictionary of additional kwargs to pass to algorithm + """ + + # note: we need the check below because some configs import BCConfig and exclude + # some of these options + gaussian_enabled = ("gaussian" in algo_config and algo_config.gaussian.enabled) + gmm_enabled = ("gmm" in algo_config and algo_config.gmm.enabled) + vae_enabled = ("vae" in algo_config and algo_config.vae.enabled) + + rnn_enabled = algo_config.rnn.enabled + transformer_enabled = algo_config.transformer.enabled + + if gaussian_enabled: + if rnn_enabled: + raise NotImplementedError + elif transformer_enabled: + raise NotImplementedError + else: + algo_class, algo_kwargs = BC_Gaussian, {} + elif gmm_enabled: + if rnn_enabled: + algo_class, algo_kwargs = BC_RNN_GMM, {} + elif transformer_enabled: + algo_class, algo_kwargs = BC_Transformer_GMM, {} + else: + algo_class, algo_kwargs = BC_GMM, {} + elif vae_enabled: + if rnn_enabled: + raise NotImplementedError + elif transformer_enabled: + raise NotImplementedError + else: + algo_class, algo_kwargs = BC_VAE, {} + else: + if rnn_enabled: + algo_class, algo_kwargs = BC_RNN, {} + elif transformer_enabled: + algo_class, algo_kwargs = BC_Transformer, {} + else: + algo_class, algo_kwargs = BC, {} + + return algo_class, algo_kwargs + + +class BC(PolicyAlgo): + """ + Normal BC training. + """ + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + self.nets = nn.ModuleDict() + self.nets["policy"] = PolicyNets.ActorNetwork( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.actor_layer_dims, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + self.nets = self.nets.float().to(self.device) + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + input_batch = dict() + input_batch["obs"] = {k: batch["obs"][k][:, 0, :] for k in batch["obs"]} + input_batch["goal_obs"] = batch.get("goal_obs", None) # goals may not be present + input_batch["actions"] = batch["actions"][:, 0, :] + # we move to device first before float conversion because image observation modalities will be uint8 - + # this minimizes the amount of data transferred to GPU + return TensorUtils.to_float(TensorUtils.to_device(input_batch, self.device)) + + + def train_on_batch(self, batch, epoch, validate=False): + """ + Training on a single batch of data. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + with TorchUtils.maybe_no_grad(no_grad=validate): + info = super(BC, self).train_on_batch(batch, epoch, validate=validate) + predictions = self._forward_training(batch) + losses = self._compute_losses(predictions, batch) + + info["predictions"] = TensorUtils.detach(predictions) + info["losses"] = TensorUtils.detach(losses) + + if not validate: + step_info = self._train_step(losses) + info.update(step_info) + + return info + + def _forward_training(self, batch): + """ + Internal helper function for BC algo class. Compute forward pass + and return network outputs in @predictions dict. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + Returns: + predictions (dict): dictionary containing network outputs + """ + predictions = OrderedDict() + actions = self.nets["policy"](obs_dict=batch["obs"], goal_dict=batch["goal_obs"]) + predictions["actions"] = actions + return predictions + + def _compute_losses(self, predictions, batch): + """ + Internal helper function for BC algo class. Compute losses based on + network outputs in @predictions dict, using reference labels in @batch. + + Args: + predictions (dict): dictionary containing network outputs, from @_forward_training + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + Returns: + losses (dict): dictionary of losses computed over the batch + """ + losses = OrderedDict() + a_target = batch["actions"] + actions = predictions["actions"] + losses["l2_loss"] = nn.MSELoss()(actions, a_target) + losses["l1_loss"] = nn.SmoothL1Loss()(actions, a_target) + # cosine direction loss on eef delta position + losses["cos_loss"] = LossUtils.cosine_loss(actions[..., :3], a_target[..., :3]) + + action_losses = [ + self.algo_config.loss.l2_weight * losses["l2_loss"], + self.algo_config.loss.l1_weight * losses["l1_loss"], + self.algo_config.loss.cos_weight * losses["cos_loss"], + ] + action_loss = sum(action_losses) + losses["action_loss"] = action_loss + return losses + + def _train_step(self, losses): + """ + Internal helper function for BC algo class. Perform backpropagation on the + loss tensors in @losses to update networks. + + Args: + losses (dict): dictionary of losses computed over the batch, from @_compute_losses + """ + + # gradient step + info = OrderedDict() + policy_grad_norms = TorchUtils.backprop_for_loss( + net=self.nets["policy"], + optim=self.optimizers["policy"], + loss=losses["action_loss"], + max_grad_norm=self.global_config.train.max_grad_norm, + ) + info["policy_grad_norms"] = policy_grad_norms + return info + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + log = super(BC, self).log_info(info) + log["Loss"] = info["losses"]["action_loss"].item() + if "l2_loss" in info["losses"]: + log["L2_Loss"] = info["losses"]["l2_loss"].item() + if "l1_loss" in info["losses"]: + log["L1_Loss"] = info["losses"]["l1_loss"].item() + if "cos_loss" in info["losses"]: + log["Cosine_Loss"] = info["losses"]["cos_loss"].item() + if "policy_grad_norms" in info: + log["Policy_Grad_Norms"] = info["policy_grad_norms"] + return log + + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + action (torch.Tensor): action tensor + """ + assert not self.nets.training + return self.nets["policy"](obs_dict, goal_dict=goal_dict) + + +class BC_Gaussian(BC): + """ + BC training with a Gaussian policy. + """ + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + assert self.algo_config.gaussian.enabled + + self.nets = nn.ModuleDict() + self.nets["policy"] = PolicyNets.GaussianActorNetwork( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.actor_layer_dims, + fixed_std=self.algo_config.gaussian.fixed_std, + init_std=self.algo_config.gaussian.init_std, + std_limits=(self.algo_config.gaussian.min_std, 7.5), + std_activation=self.algo_config.gaussian.std_activation, + low_noise_eval=self.algo_config.gaussian.low_noise_eval, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + + self.nets = self.nets.float().to(self.device) + + def _forward_training(self, batch): + """ + Internal helper function for BC algo class. Compute forward pass + and return network outputs in @predictions dict. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + Returns: + predictions (dict): dictionary containing network outputs + """ + dists = self.nets["policy"].forward_train( + obs_dict=batch["obs"], + goal_dict=batch["goal_obs"], + ) + + # make sure that this is a batch of multivariate action distributions, so that + # the log probability computation will be correct + assert len(dists.batch_shape) == 1 + log_probs = dists.log_prob(batch["actions"]) + + predictions = OrderedDict( + log_probs=log_probs, + ) + return predictions + + def _compute_losses(self, predictions, batch): + """ + Internal helper function for BC algo class. Compute losses based on + network outputs in @predictions dict, using reference labels in @batch. + + Args: + predictions (dict): dictionary containing network outputs, from @_forward_training + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + Returns: + losses (dict): dictionary of losses computed over the batch + """ + + # loss is just negative log-likelihood of action targets + action_loss = -predictions["log_probs"].mean() + return OrderedDict( + log_probs=-action_loss, + action_loss=action_loss, + ) + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + log = PolicyAlgo.log_info(self, info) + log["Loss"] = info["losses"]["action_loss"].item() + log["Log_Likelihood"] = info["losses"]["log_probs"].item() + if "policy_grad_norms" in info: + log["Policy_Grad_Norms"] = info["policy_grad_norms"] + return log + + +class BC_GMM(BC_Gaussian): + """ + BC training with a Gaussian Mixture Model policy. + """ + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + assert self.algo_config.gmm.enabled + + self.nets = nn.ModuleDict() + self.nets["policy"] = PolicyNets.GMMActorNetwork( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.actor_layer_dims, + num_modes=self.algo_config.gmm.num_modes, + min_std=self.algo_config.gmm.min_std, + std_activation=self.algo_config.gmm.std_activation, + low_noise_eval=self.algo_config.gmm.low_noise_eval, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + + self.nets = self.nets.float().to(self.device) + + +class BC_VAE(BC): + """ + BC training with a VAE policy. + """ + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + self.nets = nn.ModuleDict() + self.nets["policy"] = PolicyNets.VAEActor( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + device=self.device, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + **VAENets.vae_args_from_config(self.algo_config.vae), + ) + + self.nets = self.nets.float().to(self.device) + + def train_on_batch(self, batch, epoch, validate=False): + """ + Update from superclass to set categorical temperature, for categorical VAEs. + """ + if self.algo_config.vae.prior.use_categorical: + temperature = self.algo_config.vae.prior.categorical_init_temp - epoch * self.algo_config.vae.prior.categorical_temp_anneal_step + temperature = max(temperature, self.algo_config.vae.prior.categorical_min_temp) + self.nets["policy"].set_gumbel_temperature(temperature) + return super(BC_VAE, self).train_on_batch(batch, epoch, validate=validate) + + def _forward_training(self, batch): + """ + Internal helper function for BC algo class. Compute forward pass + and return network outputs in @predictions dict. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + Returns: + predictions (dict): dictionary containing network outputs + """ + vae_inputs = dict( + actions=batch["actions"], + obs_dict=batch["obs"], + goal_dict=batch["goal_obs"], + freeze_encoder=batch.get("freeze_encoder", False), + ) + + vae_outputs = self.nets["policy"].forward_train(**vae_inputs) + predictions = OrderedDict( + actions=vae_outputs["decoder_outputs"], + kl_loss=vae_outputs["kl_loss"], + reconstruction_loss=vae_outputs["reconstruction_loss"], + encoder_z=vae_outputs["encoder_z"], + ) + if not self.algo_config.vae.prior.use_categorical: + with torch.no_grad(): + encoder_variance = torch.exp(vae_outputs["encoder_params"]["logvar"]) + predictions["encoder_variance"] = encoder_variance + return predictions + + def _compute_losses(self, predictions, batch): + """ + Internal helper function for BC algo class. Compute losses based on + network outputs in @predictions dict, using reference labels in @batch. + + Args: + predictions (dict): dictionary containing network outputs, from @_forward_training + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + Returns: + losses (dict): dictionary of losses computed over the batch + """ + + # total loss is sum of reconstruction and KL, weighted by beta + kl_loss = predictions["kl_loss"] + recons_loss = predictions["reconstruction_loss"] + action_loss = recons_loss + self.algo_config.vae.kl_weight * kl_loss + return OrderedDict( + recons_loss=recons_loss, + kl_loss=kl_loss, + action_loss=action_loss, + ) + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + log = PolicyAlgo.log_info(self, info) + log["Loss"] = info["losses"]["action_loss"].item() + log["KL_Loss"] = info["losses"]["kl_loss"].item() + log["Reconstruction_Loss"] = info["losses"]["recons_loss"].item() + if self.algo_config.vae.prior.use_categorical: + log["Gumbel_Temperature"] = self.nets["policy"].get_gumbel_temperature() + else: + log["Encoder_Variance"] = info["predictions"]["encoder_variance"].mean().item() + if "policy_grad_norms" in info: + log["Policy_Grad_Norms"] = info["policy_grad_norms"] + return log + + +class BC_RNN(BC): + """ + BC training with an RNN policy. + """ + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + self.nets = nn.ModuleDict() + self.nets["policy"] = PolicyNets.RNNActorNetwork( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.actor_layer_dims, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + **BaseNets.rnn_args_from_config(self.algo_config.rnn), + ) + + self._rnn_hidden_state = None + self._rnn_horizon = self.algo_config.rnn.horizon + self._rnn_counter = 0 + self._rnn_is_open_loop = self.algo_config.rnn.get("open_loop", False) + + self.nets = self.nets.float().to(self.device) + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + input_batch = dict() + input_batch["obs"] = batch["obs"] + input_batch["goal_obs"] = batch.get("goal_obs", None) # goals may not be present + input_batch["actions"] = batch["actions"] + + if self._rnn_is_open_loop: + # replace the observation sequence with one that only consists of the first observation. + # This way, all actions are predicted "open-loop" after the first observation, based + # on the rnn hidden state. + n_steps = batch["actions"].shape[1] + obs_seq_start = TensorUtils.index_at_time(batch["obs"], ind=0) + input_batch["obs"] = TensorUtils.unsqueeze_expand_at(obs_seq_start, size=n_steps, dim=1) + + # we move to device first before float conversion because image observation modalities will be uint8 - + # this minimizes the amount of data transferred to GPU + return TensorUtils.to_float(TensorUtils.to_device(input_batch, self.device)) + + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + action (torch.Tensor): action tensor + """ + assert not self.nets.training + + if self._rnn_hidden_state is None or self._rnn_counter % self._rnn_horizon == 0: + batch_size = list(obs_dict.values())[0].shape[0] + self._rnn_hidden_state = self.nets["policy"].get_rnn_init_state(batch_size=batch_size, device=self.device) + + if self._rnn_is_open_loop: + # remember the initial observation, and use it instead of the current observation + # for open-loop action sequence prediction + self._open_loop_obs = TensorUtils.clone(TensorUtils.detach(obs_dict)) + + obs_to_use = obs_dict + if self._rnn_is_open_loop: + # replace current obs with last recorded obs + obs_to_use = self._open_loop_obs + + self._rnn_counter += 1 + action, self._rnn_hidden_state = self.nets["policy"].forward_step( + obs_to_use, goal_dict=goal_dict, rnn_state=self._rnn_hidden_state) + return action + + def reset(self): + """ + Reset algo state to prepare for environment rollouts. + """ + self._rnn_hidden_state = None + self._rnn_counter = 0 + + +class BC_RNN_GMM(BC_RNN): + """ + BC training with an RNN GMM policy. + """ + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + assert self.algo_config.gmm.enabled + assert self.algo_config.rnn.enabled + + self.nets = nn.ModuleDict() + self.nets["policy"] = PolicyNets.RNNGMMActorNetwork( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.actor_layer_dims, + num_modes=self.algo_config.gmm.num_modes, + min_std=self.algo_config.gmm.min_std, + std_activation=self.algo_config.gmm.std_activation, + low_noise_eval=self.algo_config.gmm.low_noise_eval, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + **BaseNets.rnn_args_from_config(self.algo_config.rnn), + ) + + self._rnn_hidden_state = None + self._rnn_horizon = self.algo_config.rnn.horizon + self._rnn_counter = 0 + self._rnn_is_open_loop = self.algo_config.rnn.get("open_loop", False) + + self.nets = self.nets.float().to(self.device) + + def _forward_training(self, batch): + """ + Internal helper function for BC algo class. Compute forward pass + and return network outputs in @predictions dict. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + Returns: + predictions (dict): dictionary containing network outputs + """ + dists = self.nets["policy"].forward_train( + obs_dict=batch["obs"], + goal_dict=batch["goal_obs"], + ) + + # make sure that this is a batch of multivariate action distributions, so that + # the log probability computation will be correct + assert len(dists.batch_shape) == 2 # [B, T] + log_probs = dists.log_prob(batch["actions"]) + + predictions = OrderedDict( + log_probs=log_probs, + ) + return predictions + + def _compute_losses(self, predictions, batch): + """ + Internal helper function for BC algo class. Compute losses based on + network outputs in @predictions dict, using reference labels in @batch. + + Args: + predictions (dict): dictionary containing network outputs, from @_forward_training + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + Returns: + losses (dict): dictionary of losses computed over the batch + """ + + # loss is just negative log-likelihood of action targets + action_loss = -predictions["log_probs"].mean() + return OrderedDict( + log_probs=-action_loss, + action_loss=action_loss, + ) + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + log = PolicyAlgo.log_info(self, info) + log["Loss"] = info["losses"]["action_loss"].item() + log["Log_Likelihood"] = info["losses"]["log_probs"].item() + if "policy_grad_norms" in info: + log["Policy_Grad_Norms"] = info["policy_grad_norms"] + return log + + +class BC_Transformer(BC): + """ + BC training with a Transformer policy. + """ + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + assert self.algo_config.transformer.enabled + + self.nets = nn.ModuleDict() + self.nets["policy"] = PolicyNets.TransformerActorNetwork( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + **BaseNets.transformer_args_from_config(self.algo_config.transformer), + ) + self._set_params_from_config() + self.nets = self.nets.float().to(self.device) + + def _set_params_from_config(self): + """ + Read specific config variables we need for training / eval. + Called by @_create_networks method + """ + self.context_length = self.algo_config.transformer.context_length + self.supervise_all_steps = self.algo_config.transformer.supervise_all_steps + self.pred_future_acs = self.algo_config.transformer.pred_future_acs + if self.pred_future_acs: + assert self.supervise_all_steps is True + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + input_batch = dict() + h = self.context_length + input_batch["obs"] = {k: batch["obs"][k][:, :h, :] for k in batch["obs"]} + input_batch["goal_obs"] = batch.get("goal_obs", None) # goals may not be present + + if self.supervise_all_steps: + # supervision on entire sequence (instead of just current timestep) + if self.pred_future_acs: + ac_start = h - 1 + else: + ac_start = 0 + input_batch["actions"] = batch["actions"][:, ac_start:ac_start+h, :] + else: + # just use current timestep + input_batch["actions"] = batch["actions"][:, h-1, :] + + if self.pred_future_acs: + assert input_batch["actions"].shape[1] == h + + input_batch = TensorUtils.to_device(TensorUtils.to_float(input_batch), self.device) + return input_batch + + def _forward_training(self, batch, epoch=None): + """ + Internal helper function for BC_Transformer algo class. Compute forward pass + and return network outputs in @predictions dict. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + Returns: + predictions (dict): dictionary containing network outputs + """ + # ensure that transformer context length is consistent with temporal dimension of observations + TensorUtils.assert_size_at_dim( + batch["obs"], + size=(self.context_length), + dim=1, + msg="Error: expect temporal dimension of obs batch to match transformer context length {}".format(self.context_length), + ) + + predictions = OrderedDict() + predictions["actions"] = self.nets["policy"](obs_dict=batch["obs"], actions=None, goal_dict=batch["goal_obs"]) + if not self.supervise_all_steps: + # only supervise final timestep + predictions["actions"] = predictions["actions"][:, -1, :] + return predictions + + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + Returns: + action (torch.Tensor): action tensor + """ + assert not self.nets.training + + output = self.nets["policy"](obs_dict, actions=None, goal_dict=goal_dict) + + if self.supervise_all_steps: + if self.algo_config.transformer.pred_future_acs: + output = output[:, 0, :] + else: + output = output[:, -1, :] + else: + output = output[:, -1, :] + + return output + + + +class BC_Transformer_GMM(BC_Transformer): + """ + BC training with a Transformer GMM policy. + """ + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + assert self.algo_config.gmm.enabled + assert self.algo_config.transformer.enabled + + if self.algo_config.language_conditioned: + self.obs_shapes["lang_emb"] = [768] # clip is 768-dim embedding + + self.nets = nn.ModuleDict() + self.nets["policy"] = PolicyNets.TransformerGMMActorNetwork( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + num_modes=self.algo_config.gmm.num_modes, + min_std=self.algo_config.gmm.min_std, + std_activation=self.algo_config.gmm.std_activation, + low_noise_eval=self.algo_config.gmm.low_noise_eval, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + **BaseNets.transformer_args_from_config(self.algo_config.transformer), + ) + self._set_params_from_config() + self.nets = self.nets.float().to(self.device) + + def _forward_training(self, batch, epoch=None): + """ + Modify from super class to support GMM training. + """ + # ensure that transformer context length is consistent with temporal dimension of observations + TensorUtils.assert_size_at_dim( + batch["obs"], + size=(self.context_length), + dim=1, + msg="Error: expect temporal dimension of obs batch to match transformer context length {}".format(self.context_length), + ) + + dists = self.nets["policy"].forward_train( + obs_dict=batch["obs"], + actions=None, + goal_dict=batch["goal_obs"], + low_noise_eval=False, + ) + + # make sure that this is a batch of multivariate action distributions, so that + # the log probability computation will be correct + assert len(dists.batch_shape) == 2 # [B, T] + + if not self.supervise_all_steps: + # only use final timestep prediction by making a new distribution with only final timestep. + # This essentially does `dists = dists[:, -1]` + component_distribution = D.Normal( + loc=dists.component_distribution.base_dist.loc[:, -1], + scale=dists.component_distribution.base_dist.scale[:, -1], + ) + component_distribution = D.Independent(component_distribution, 1) + mixture_distribution = D.Categorical(logits=dists.mixture_distribution.logits[:, -1]) + dists = D.MixtureSameFamily( + mixture_distribution=mixture_distribution, + component_distribution=component_distribution, + ) + + log_probs = dists.log_prob(batch["actions"]) + + predictions = OrderedDict( + log_probs=log_probs, + ) + return predictions + + def _compute_losses(self, predictions, batch): + """ + Internal helper function for BC_Transformer_GMM algo class. Compute losses based on + network outputs in @predictions dict, using reference labels in @batch. + Args: + predictions (dict): dictionary containing network outputs, from @_forward_training + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + Returns: + losses (dict): dictionary of losses computed over the batch + """ + + # loss is just negative log-likelihood of action targets + action_loss = -predictions["log_probs"].mean() + return OrderedDict( + log_probs=-action_loss, + action_loss=action_loss, + ) + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + Args: + info (dict): dictionary of info + Returns: + loss_log (dict): name -> summary statistic + """ + log = PolicyAlgo.log_info(self, info) + log["Loss"] = info["losses"]["action_loss"].item() + log["Log_Likelihood"] = info["losses"]["log_probs"].item() + if "policy_grad_norms" in info: + log["Policy_Grad_Norms"] = info["policy_grad_norms"] + return log \ No newline at end of file diff --git a/aloha-devel/robomimic/algo/bcq.py b/aloha-devel/robomimic/algo/bcq.py new file mode 100644 index 0000000000000000000000000000000000000000..5843ccb5bd594c596a8dc138eab863bb3f5e3550 --- /dev/null +++ b/aloha-devel/robomimic/algo/bcq.py @@ -0,0 +1,1022 @@ +""" +Batch-Constrained Q-Learning (BCQ), with support for more general +generative action models (the original paper uses a cVAE). +(Paper - https://arxiv.org/abs/1812.02900). +""" +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import robomimic.models.obs_nets as ObsNets +import robomimic.models.policy_nets as PolicyNets +import robomimic.models.value_nets as ValueNets +import robomimic.models.vae_nets as VAENets +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.torch_utils as TorchUtils +import robomimic.utils.obs_utils as ObsUtils +import robomimic.utils.loss_utils as LossUtils + +from robomimic.algo import register_algo_factory_func, PolicyAlgo, ValueAlgo + + +@register_algo_factory_func("bcq") +def algo_config_to_class(algo_config): + """ + Maps algo config to the BCQ algo class to instantiate, along with additional algo kwargs. + + Args: + algo_config (Config instance): algo config + + Returns: + algo_class: subclass of Algo + algo_kwargs (dict): dictionary of additional kwargs to pass to algorithm + """ + if algo_config.critic.distributional.enabled: + return BCQ_Distributional, {} + if algo_config.action_sampler.gmm.enabled: + return BCQ_GMM, {} + assert algo_config.action_sampler.vae.enabled + return BCQ, {} + + +class BCQ(PolicyAlgo, ValueAlgo): + """ + Default BCQ training, based on https://arxiv.org/abs/1812.02900 and + https://github.com/sfujim/BCQ + """ + def __init__(self, **kwargs): + PolicyAlgo.__init__(self, **kwargs) + + # save the discount factor - it may be overriden later + self.set_discount(self.algo_config.discount) + + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + self.nets = nn.ModuleDict() + + self._create_critics() + self._create_action_sampler() + if self.algo_config.actor.enabled: + self._create_actor() + + # sync target networks at beginning of training + with torch.no_grad(): + for critic_ind in range(len(self.nets["critic"])): + TorchUtils.hard_update( + source=self.nets["critic"][critic_ind], + target=self.nets["critic_target"][critic_ind], + ) + + if self.algo_config.actor.enabled: + TorchUtils.hard_update( + source=self.nets["actor"], + target=self.nets["actor_target"], + ) + + self.nets = self.nets.float().to(self.device) + + def _create_critics(self): + """ + Called in @_create_networks to make critic networks. + """ + critic_class = ValueNets.ActionValueNetwork + critic_args = dict( + obs_shapes=self.obs_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.critic.layer_dims, + value_bounds=self.algo_config.critic.value_bounds, + goal_shapes=self.goal_shapes, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + + # Q network ensemble and target ensemble + self.nets["critic"] = nn.ModuleList() + self.nets["critic_target"] = nn.ModuleList() + for _ in range(self.algo_config.critic.ensemble.n): + critic = critic_class(**critic_args) + self.nets["critic"].append(critic) + + critic_target = critic_class(**critic_args) + self.nets["critic_target"].append(critic_target) + + def _create_action_sampler(self): + """ + Called in @_create_networks to make action sampler network. + """ + + # VAE network for approximate sampling from batch dataset + assert self.algo_config.action_sampler.vae.enabled + self.nets["action_sampler"] = PolicyNets.VAEActor( + obs_shapes=self.obs_shapes, + ac_dim=self.ac_dim, + device=self.device, + goal_shapes=self.goal_shapes, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + **VAENets.vae_args_from_config(self.algo_config.action_sampler.vae), + ) + + def _create_actor(self): + """ + Called in @_create_networks to make actor network. + """ + assert self.algo_config.actor.enabled + actor_class = PolicyNets.PerturbationActorNetwork + actor_args = dict( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.actor.layer_dims, + perturbation_scale=self.algo_config.actor.perturbation_scale, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + + self.nets["actor"] = actor_class(**actor_args) + self.nets["actor_target"] = actor_class(**actor_args) + + def _check_epoch(self, net_name, epoch): + """ + Helper function to check whether backprop should happen this epoch. + + Args: + net_name (str): name of network in @self.nets and @self.optim_params + epoch (int): epoch number + """ + epoch_start_check = (self.optim_params[net_name]["start_epoch"] == -1) or (epoch >= self.optim_params[net_name]["start_epoch"]) + epoch_end_check = (self.optim_params[net_name]["end_epoch"] == -1) or (epoch < self.optim_params[net_name]["end_epoch"]) + return (epoch_start_check and epoch_end_check) + + def set_discount(self, discount): + """ + Useful function to modify discount factor if necessary (e.g. for n-step returns). + """ + self.discount = discount + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + input_batch = dict() + + # n-step returns (default is 1) + n_step = self.algo_config.n_step + assert batch["actions"].shape[1] >= n_step + + # remove temporal batches for all + input_batch["obs"] = {k: batch["obs"][k][:, 0, :] for k in batch["obs"]} + input_batch["next_obs"] = {k: batch["next_obs"][k][:, n_step - 1, :] for k in batch["next_obs"]} + input_batch["goal_obs"] = batch.get("goal_obs", None) # goals may not be present + input_batch["actions"] = batch["actions"][:, 0, :] + + # note: ensure scalar signals (rewards, done) retain last dimension of 1 to be compatible with model outputs + + # single timestep reward is discounted sum of intermediate rewards in sequence + reward_seq = batch["rewards"][:, :n_step] + discounts = torch.pow(self.algo_config.discount, torch.arange(n_step).float()).unsqueeze(0) + input_batch["rewards"] = (reward_seq * discounts).sum(dim=1).unsqueeze(1) + + # discount rate will be gamma^N for computing n-step returns + new_discount = (self.algo_config.discount ** n_step) + self.set_discount(new_discount) + + # consider this n-step seqeunce done if any intermediate dones are present + done_seq = batch["dones"][:, :n_step] + input_batch["dones"] = (done_seq.sum(dim=1) > 0).float().unsqueeze(1) + + if self.algo_config.infinite_horizon: + # scale terminal rewards by 1 / (1 - gamma) for infinite horizon MDPs + done_inds = input_batch["dones"].round().long().nonzero(as_tuple=False)[:, 0] + if done_inds.shape[0] > 0: + input_batch["rewards"][done_inds] = input_batch["rewards"][done_inds] * (1. / (1. - self.discount)) + + # we move to device first before float conversion because image observation modalities will be uint8 - + # this minimizes the amount of data transferred to GPU + return TensorUtils.to_float(TensorUtils.to_device(input_batch, self.device)) + + def _train_action_sampler_on_batch(self, batch, epoch, no_backprop=False): + """ + A modular helper function that can be overridden in case + subclasses would like to modify training behavior for the + action sampler. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + no_backprop (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + outputs (dict): dictionary of outputs to use during critic training + (for computing target values) + """ + info = OrderedDict() + if self.algo_config.action_sampler.vae.prior.use_categorical: + temperature = self.algo_config.action_sampler.vae.prior.categorical_init_temp - epoch * self.algo_config.action_sampler.vae.prior.categorical_temp_anneal_step + temperature = max(temperature, self.algo_config.action_sampler.vae.prior.categorical_min_temp) + self.nets["action_sampler"].set_gumbel_temperature(temperature) + + vae_inputs = dict( + actions=batch["actions"], + obs_dict=batch["obs"], + goal_dict=batch["goal_obs"], + ) + + # maybe freeze encoder weights + if (self.algo_config.action_sampler.freeze_encoder_epoch != -1) and (epoch >= self.algo_config.action_sampler.freeze_encoder_epoch): + vae_inputs["freeze_encoder"] = True + + # VAE forward + vae_outputs = self.nets["action_sampler"].forward_train(**vae_inputs) + recons_loss = vae_outputs["reconstruction_loss"] + kl_loss = vae_outputs["kl_loss"] + vae_loss = recons_loss + self.algo_config.action_sampler.vae.kl_weight * kl_loss + info["action_sampler/loss"] = vae_loss + info["action_sampler/recons_loss"] = recons_loss + info["action_sampler/kl_loss"] = kl_loss + if not self.algo_config.action_sampler.vae.prior.use_categorical: + with torch.no_grad(): + encoder_variance = torch.exp(vae_outputs["encoder_params"]["logvar"]).mean() + info["action_sampler/encoder_variance"] = encoder_variance + outputs = TensorUtils.detach(vae_outputs) + + # VAE gradient step + if not no_backprop: + vae_grad_norms = TorchUtils.backprop_for_loss( + net=self.nets["action_sampler"], + optim=self.optimizers["action_sampler"], + loss=vae_loss, + ) + info["action_sampler/grad_norms"] = vae_grad_norms + return info, outputs + + def _train_critic_on_batch(self, batch, action_sampler_outputs, epoch, no_backprop=False): + """ + A modular helper function that can be overridden in case + subclasses would like to modify training behavior for the + critics. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + action_sampler_outputs (dict): dictionary of outputs from the action sampler. Used + to form target values for training the critic + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + no_backprop (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + critic_outputs (dict): dictionary of critic outputs - useful for + logging purposes + """ + info = OrderedDict() + + # batch variables + s_batch = batch["obs"] + a_batch = batch["actions"] + r_batch = batch["rewards"] + ns_batch = batch["next_obs"] + goal_s_batch = batch["goal_obs"] + + # 1 if not done, 0 otherwise + done_mask_batch = 1. - batch["dones"] + info["done_masks"] = done_mask_batch + + # Bellman backup for Q-targets + q_targets = self._get_target_values( + next_states=ns_batch, + goal_states=goal_s_batch, + rewards=r_batch, + dones=done_mask_batch, + action_sampler_outputs=action_sampler_outputs, + ) + info["critic/q_targets"] = q_targets + + # Train all critics using this set of targets for regression + critic_outputs = [] + for critic_ind, critic in enumerate(self.nets["critic"]): + critic_loss, critic_output = self._compute_critic_loss( + critic=critic, + states=s_batch, + actions=a_batch, + goal_states=goal_s_batch, + q_targets=q_targets, + ) + info["critic/critic{}_loss".format(critic_ind + 1)] = critic_loss + critic_outputs.append(critic_output) + + if not no_backprop: + critic_grad_norms = TorchUtils.backprop_for_loss( + net=self.nets["critic"][critic_ind], + optim=self.optimizers["critic"][critic_ind], + loss=critic_loss, + max_grad_norm=self.algo_config.critic.max_gradient_norm, + ) + info["critic/critic{}_grad_norms".format(critic_ind + 1)] = critic_grad_norms + + return info, critic_outputs + + def _train_actor_on_batch(self, batch, action_sampler_outputs, critic_outputs, epoch, no_backprop=False): + """ + A modular helper function that can be overridden in case + subclasses would like to modify training behavior for the + perturbation actor. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + action_sampler_outputs (dict): dictionary of outputs from the action sampler. Currently + unused, although more sophisticated models may use it. + + critic_outputs (dict): dictionary of outputs from the critic. Currently + unused, although more sophisticated models may use it. + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + no_backprop (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + assert self.algo_config.actor.enabled + + info = OrderedDict() + + # Actor loss (update with DDPG loss) + s_batch = batch["obs"] + goal_s_batch = batch["goal_obs"] + + # sample some actions from action sampler and perturb them, then improve perturbations + # where improvement is measured by the critic + sampled_actions = self.nets["action_sampler"](s_batch, goal_s_batch).detach() # don't backprop into samples + perturbed_actions = self.nets["actor"](s_batch, sampled_actions, goal_s_batch) + actor_loss = -(self.nets["critic"][0](s_batch, perturbed_actions, goal_s_batch)).mean() + info["actor/loss"] = actor_loss + + if not no_backprop: + actor_grad_norms = TorchUtils.backprop_for_loss( + net=self.nets["actor"], + optim=self.optimizers["actor"], + loss=actor_loss, + ) + info["actor/grad_norms"] = actor_grad_norms + + return info + + def _get_target_values(self, next_states, goal_states, rewards, dones, action_sampler_outputs=None): + """ + Helper function to get target values for training Q-function with TD-loss. + + Args: + next_states (dict): batch of next observations + goal_states (dict): if not None, batch of goal observations + rewards (torch.Tensor): batch of rewards - should be shape (B, 1) + dones (torch.Tensor): batch of done signals - should be shape (B, 1) + action_sampler_outputs (dict): dictionary of outputs from the action sampler. Currently + unused, although more sophisticated models may use it. + + Returns: + q_targets (torch.Tensor): target Q-values to use for TD loss + """ + + with torch.no_grad(): + # we need to stack the observations with redundancy @num_action_samples here, then decode + # to get all sampled actions. for example, if we generate 2 samples per observation and + # the batch size is 3, then ob_tiled = [ob1; ob1; ob2; ob2; ob3; ob3] + next_states_tiled = ObsUtils.repeat_and_stack_observation(next_states, n=self.algo_config.critic.num_action_samples) + goal_states_tiled = None + if len(self.goal_shapes) > 0: + goal_states_tiled = ObsUtils.repeat_and_stack_observation(goal_states, n=self.algo_config.critic.num_action_samples) + + # sample action proposals + next_sampled_actions = self._sample_actions_for_value_maximization( + states_tiled=next_states_tiled, + goal_states_tiled=goal_states_tiled, + for_target_update=True, + ) + + q_targets = self._get_target_values_from_sampled_actions( + next_states_tiled=next_states_tiled, + next_sampled_actions=next_sampled_actions, + goal_states_tiled=goal_states_tiled, + rewards=rewards, + dones=dones, + ) + + return q_targets + + def _sample_actions_for_value_maximization(self, states_tiled, goal_states_tiled, for_target_update): + """ + Helper function to sample actions for maximization (the "batch-constrained" part of + batch-constrained q-learning). + + Args: + states_tiled (dict): observations to use for sampling actions. Assumes that tiling + has already occurred - so that if the batch size is B, and N samples are + desired for each observation in the batch, the leading dimension for each + observation in the dict is B * N + + goal_states_tiled (dict): if not None, goal observations + + for_target_update (bool): if True, actions are being sampled for use in training the + critic - which means the target actor network should be used + + Returns: + sampled_actions (torch.Tensor): actions sampled from the action sampler, and maybe + perturbed by the actor network + """ + + with torch.no_grad(): + sampled_actions = self.nets["action_sampler"](states_tiled, goal_states_tiled) + if self.algo_config.actor.enabled: + actor = self.nets["actor"] + if for_target_update: + actor = self.nets["actor_target"] + # perturb the actions with the policy + sampled_actions = actor(states_tiled, sampled_actions, goal_states_tiled) + + return sampled_actions + + def _get_target_values_from_sampled_actions(self, next_states_tiled, next_sampled_actions, goal_states_tiled, rewards, dones): + """ + Helper function to get target values for training Q-function with TD-loss. The function + assumes that action candidates to maximize over have already been computed, and that + the input states have been tiled (repeated) to be compatible with the sampled actions. + + Args: + next_states_tiled (dict): next observations to use for sampling actions. Assumes that + tiling has already occurred - so that if the batch size is B, and N samples are + desired for each observation in the batch, the leading dimension for each + observation in the dict is B * N + + next_sampled_actions (torch.Tensor): actions sampled from the action sampler. This function + will maximize the critic over these action candidates (using the TD3 trick) + + goal_states_tiled (dict): if not None, goal observations + + rewards (torch.Tensor): batch of rewards - should be shape (B, 1) + + dones (torch.Tensor): batch of done signals - should be shape (B, 1) + + Returns: + q_targets (torch.Tensor): target Q-values to use for TD loss + """ + with torch.no_grad(): + # feed tiled observations and sampled actions into the critics and then + # reshape to get all Q-values in second dimension per observation in batch. + all_value_targets = self.nets["critic_target"][0](next_states_tiled, next_sampled_actions, goal_states_tiled).reshape( + -1, self.algo_config.critic.num_action_samples) + max_value_targets = all_value_targets + min_value_targets = all_value_targets + + # TD3 trick to combine max and min over all Q-ensemble estimates into single target estimates + for critic_target in self.nets["critic_target"][1:]: + all_value_targets = critic_target(next_states_tiled, next_sampled_actions, goal_states_tiled).reshape( + -1, self.algo_config.critic.num_action_samples) + max_value_targets = torch.max(max_value_targets, all_value_targets) + min_value_targets = torch.min(min_value_targets, all_value_targets) + all_value_targets = self.algo_config.critic.ensemble.weight * min_value_targets + \ + (1. - self.algo_config.critic.ensemble.weight) * max_value_targets + + # take maximum over all sampled action values per observation and compute targets + value_targets = torch.max(all_value_targets, dim=1, keepdim=True)[0] + q_targets = rewards + dones * self.discount * value_targets + + return q_targets + + def _compute_critic_loss(self, critic, states, actions, goal_states, q_targets): + """ + Helper function to compute loss between estimated Q-values and target Q-values. + It should also return outputs needed for downstream training (for training the + actor). + + Args: + critic (torch.nn.Module): critic network + states (dict): batch of observations + actions (torch.Tensor): batch of actions + goal_states (dict): if not None, batch of goal observations + q_targets (torch.Tensor): batch of target q-values for the TD loss + + Returns: + critic_loss (torch.Tensor): critic loss + critic_output (dict): additional outputs from the critic. This function + returns None, but subclasses may want to provide some information + here. + """ + q_estimated = critic(states, actions, goal_states) + if self.algo_config.critic.use_huber: + critic_loss = nn.SmoothL1Loss()(q_estimated, q_targets) + else: + critic_loss = nn.MSELoss()(q_estimated, q_targets) + return critic_loss, None + + def train_on_batch(self, batch, epoch, validate=False): + """ + Training on a single batch of data. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + with TorchUtils.maybe_no_grad(no_grad=validate): + info = PolicyAlgo.train_on_batch(self, batch, epoch, validate=validate) + + # Action Sampler training + no_action_sampler_backprop = validate or (not self._check_epoch(net_name="action_sampler", epoch=epoch)) + with TorchUtils.maybe_no_grad(no_grad=no_action_sampler_backprop): + action_sampler_info, action_sampler_outputs = self._train_action_sampler_on_batch( + batch=batch, + epoch=epoch, + no_backprop=no_action_sampler_backprop, + ) + info.update(action_sampler_info) + + # make sure action sampler is in eval mode for models like GMM which may require low-noise + # samples when sampling actions. + self.nets["action_sampler"].eval() + + # Critic training + no_critic_backprop = validate or (not self._check_epoch(net_name="critic", epoch=epoch)) + with TorchUtils.maybe_no_grad(no_grad=no_critic_backprop): + critic_info, critic_outputs = self._train_critic_on_batch( + batch=batch, + action_sampler_outputs=action_sampler_outputs, + epoch=epoch, + no_backprop=no_critic_backprop, + ) + info.update(critic_info) + + if self.algo_config.actor.enabled: + # Actor training + no_actor_backprop = validate or (not self._check_epoch(net_name="actor", epoch=epoch)) + with TorchUtils.maybe_no_grad(no_grad=no_actor_backprop): + actor_info = self._train_actor_on_batch( + batch=batch, + action_sampler_outputs=action_sampler_outputs, + critic_outputs=critic_outputs, + epoch=epoch, + no_backprop=no_actor_backprop, + ) + info.update(actor_info) + + if not validate: + # restore to train mode if necessary + self.nets["action_sampler"].train() + + # update the target critic networks (only when critic has gradient update) + if not no_critic_backprop: + with torch.no_grad(): + for critic_ind in range(len(self.nets["critic"])): + TorchUtils.soft_update( + source=self.nets["critic"][critic_ind], + target=self.nets["critic_target"][critic_ind], + tau=self.algo_config.target_tau, + ) + + # update target actor network (only when actor has gradient update) + if self.algo_config.actor.enabled and (not no_actor_backprop): + with torch.no_grad(): + TorchUtils.soft_update( + source=self.nets["actor"], + target=self.nets["actor_target"], + tau=self.algo_config.target_tau, + ) + + return info + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + loss_log = OrderedDict() + + # record current optimizer learning rates + for k in self.optimizers: + keys = [k] + optims = [self.optimizers[k]] + if k == "critic": + # account for critic having one optimizer per ensemble member + keys = ["{}{}".format(k, critic_ind) for critic_ind in range(len(self.nets["critic"]))] + optims = self.optimizers[k] + for kp, optimizer in zip(keys, optims): + for i, param_group in enumerate(optimizer.param_groups): + loss_log["Optimizer/{}{}_lr".format(kp, i)] = param_group["lr"] + + # extract relevant logs for action sampler, critic, and actor + loss_log["Loss"] = 0. + for loss_logger in [self._log_action_sampler_info, self._log_critic_info, self._log_actor_info]: + this_log = loss_logger(info) + if "Loss" in this_log: + # manually merge total loss + loss_log["Loss"] += this_log["Loss"] + del this_log["Loss"] + loss_log.update(this_log) + + return loss_log + + def _log_action_sampler_info(self, info): + """ + Helper function to extract action sampler-relevant information for logging. + """ + loss_log = OrderedDict() + loss_log["Action_Sampler/Loss"] = info["action_sampler/loss"].item() + loss_log["Action_Sampler/Reconsruction_Loss"] = info["action_sampler/recons_loss"].item() + loss_log["Action_Sampler/KL_Loss"] = info["action_sampler/kl_loss"].item() + if self.algo_config.action_sampler.vae.prior.use_categorical: + loss_log["Action_Sampler/Gumbel_Temperature"] = self.nets["action_sampler"].get_gumbel_temperature() + else: + loss_log["Action_Sampler/Encoder_Variance"] = info["action_sampler/encoder_variance"].item() + if "action_sampler/grad_norms" in info: + loss_log["Action_Sampler/Grad_Norms"] = info["action_sampler/grad_norms"] + loss_log["Loss"] = loss_log["Action_Sampler/Loss"] + return loss_log + + def _log_critic_info(self, info): + """ + Helper function to extract critic-relevant information for logging. + """ + loss_log = OrderedDict() + if "done_masks" in info: + loss_log["Critic/Done_Mask_Percentage"] = 100. * torch.mean(info["done_masks"]).item() + if "critic/q_targets" in info: + loss_log["Critic/Q_Targets"] = info["critic/q_targets"].mean().item() + loss_log["Loss"] = 0. + for critic_ind in range(len(self.nets["critic"])): + loss_log["Critic/Critic{}_Loss".format(critic_ind + 1)] = info["critic/critic{}_loss".format(critic_ind + 1)].item() + if "critic/critic{}_grad_norms".format(critic_ind + 1) in info: + loss_log["Critic/Critic{}_Grad_Norms".format(critic_ind + 1)] = info["critic/critic{}_grad_norms".format(critic_ind + 1)] + loss_log["Loss"] += loss_log["Critic/Critic{}_Loss".format(critic_ind + 1)] + return loss_log + + def _log_actor_info(self, info): + """ + Helper function to extract actor-relevant information for logging. + """ + loss_log = OrderedDict() + if self.algo_config.actor.enabled: + loss_log["Actor/Loss"] = info["actor/loss"].item() + if "actor/grad_norms" in info: + loss_log["Actor/Grad_Norms"] = info["actor/grad_norms"] + loss_log["Loss"] = loss_log["Actor/Loss"] + return loss_log + + def set_train(self): + """ + Prepare networks for evaluation. Update from super class to make sure + target networks stay in evaluation mode all the time. + """ + self.nets.train() + + # target networks always in eval + for critic_ind in range(len(self.nets["critic_target"])): + self.nets["critic_target"][critic_ind].eval() + + if self.algo_config.actor.enabled: + self.nets["actor_target"].eval() + + def on_epoch_end(self, epoch): + """ + Called at the end of each epoch. + """ + + # LR scheduling updates + for lr_sc in self.lr_schedulers["critic"]: + if lr_sc is not None: + lr_sc.step() + + if self.lr_schedulers["action_sampler"] is not None: + self.lr_schedulers["action_sampler"].step() + + if self.algo_config.actor.enabled and self.lr_schedulers["actor"] is not None: + self.lr_schedulers["actor"].step() + + def _get_best_value(self, obs_dict, goal_dict=None): + """ + Internal helper function for getting the best value for a given state and + the corresponding best action. Meant to be used at test-time. Key differences + between this and retrieving target values at train-time are that (1) only a + single critic is used for the value estimate and (2) the critic and actor + are used instead of the target critic and target actor. + + Args: + obs_dict (dict): batch of current observations + goal_dict (dict): (optional) goal + + Returns: + best_value (torch.Tensor): best values + best_action (torch.Tensor): best actions + """ + assert not self.nets.training + + random_key = list(obs_dict.keys())[0] + batch_size = obs_dict[random_key].shape[0] + + # number of action proposals from action sampler + num_action_samples = self.algo_config.critic.num_action_samples_rollout + + # we need to stack the observations with redundancy @num_action_samples here, then decode + # to get all sampled actions. for example, if we generate 2 samples per observation and + # the batch size is 3, then ob_tiled = [ob1; ob1; ob2; ob2; ob3; ob3] + ob_tiled = ObsUtils.repeat_and_stack_observation(obs_dict, n=num_action_samples) + goal_tiled = None + if len(self.goal_shapes) > 0: + goal_tiled = ObsUtils.repeat_and_stack_observation(goal_dict, n=num_action_samples) + + sampled_actions = self._sample_actions_for_value_maximization( + states_tiled=ob_tiled, + goal_states_tiled=goal_tiled, + for_target_update=False, + ) + + # feed tiled observations and perturbed sampled actions into the critic and then + # reshape to get all Q-values in second dimension per observation in batch. + # finally, just take a maximum across that second dimension to take the best sampled action + all_critic_values = self.nets["critic"][0](ob_tiled, sampled_actions, goal_tiled).reshape(-1, num_action_samples) + best_action_index = torch.argmax(all_critic_values, dim=1) + + all_actions = sampled_actions.reshape(batch_size, num_action_samples, -1) + best_action = all_actions[torch.arange(all_actions.shape[0]), best_action_index] + best_value = all_critic_values[torch.arange(all_critic_values.shape[0]), best_action_index].unsqueeze(1) + + return best_value, best_action + + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + action (torch.Tensor): action tensor + """ + assert not self.nets.training + + _, best_action = self._get_best_value(obs_dict=obs_dict, goal_dict=goal_dict) + return best_action + + def get_state_value(self, obs_dict, goal_dict=None): + """ + Get state value outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + value (torch.Tensor): value tensor + """ + assert not self.nets.training + + best_value, _ = self._get_best_value(obs_dict=obs_dict, goal_dict=goal_dict) + return best_value + + def get_state_action_value(self, obs_dict, actions, goal_dict=None): + """ + Get state-action value outputs. + + Args: + obs_dict (dict): current observation + actions (torch.Tensor): action + goal_dict (dict): (optional) goal + + Returns: + value (torch.Tensor): value tensor + """ + assert not self.nets.training + + return self.nets["critic"][0](obs_dict, actions, goal_dict) + + +class BCQ_GMM(BCQ): + """ + A simple modification to BCQ that replaces the VAE used to sample action proposals from the + batch with a GMM. + """ + def _create_action_sampler(self): + """ + Called in @_create_networks to make action sampler network. + """ + assert self.algo_config.action_sampler.gmm.enabled + + # GMM network for approximate sampling from batch dataset + self.nets["action_sampler"] = PolicyNets.GMMActorNetwork( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.action_sampler.actor_layer_dims, + num_modes=self.algo_config.action_sampler.gmm.num_modes, + min_std=self.algo_config.action_sampler.gmm.min_std, + std_activation=self.algo_config.action_sampler.gmm.std_activation, + low_noise_eval=self.algo_config.action_sampler.gmm.low_noise_eval, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + + def _train_action_sampler_on_batch(self, batch, epoch, no_backprop=False): + """ + Modify this helper function from superclass to train GMM action sampler + with maximum likelihood. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + no_backprop (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + outputs (dict): dictionary of outputs to use during critic training + (for computing target values) + """ + info = OrderedDict() + + # GMM forward + dists = self.nets["action_sampler"].forward_train( + obs_dict=batch["obs"], + goal_dict=batch["goal_obs"], + ) + + # make sure that this is a batch of multivariate action distributions, so that + # the log probability computation will be correct + assert len(dists.batch_shape) == 1 + log_probs = dists.log_prob(batch["actions"]) + loss = -log_probs.mean() + info["action_sampler/loss"] = loss + + # GMM gradient step + if not no_backprop: + gmm_grad_norms = TorchUtils.backprop_for_loss( + net=self.nets["action_sampler"], + optim=self.optimizers["action_sampler"], + loss=loss, + ) + info["action_sampler/grad_norms"] = gmm_grad_norms + return info, None + + def _log_action_sampler_info(self, info): + """ + Update from superclass for GMM (no KL loss). + """ + loss_log = OrderedDict() + loss_log["Action_Sampler/Loss"] = info["action_sampler/loss"].item() + if "action_sampler/grad_norms" in info: + loss_log["Action_Sampler/Grad_Norms"] = info["action_sampler/grad_norms"] + loss_log["Loss"] = loss_log["Action_Sampler/Loss"] + return loss_log + + +class BCQ_Distributional(BCQ): + """ + BCQ with distributional critics. Distributional critics output categorical + distributions over a discrete set of values instead of expected returns. + Some parts of this implementation were adapted from ACME (https://github.com/deepmind/acme). + """ + def _create_critics(self): + """ + Called in @_create_networks to make critic networks. + """ + assert self.algo_config.critic.distributional.enabled + critic_class = ValueNets.DistributionalActionValueNetwork + critic_args = dict( + obs_shapes=self.obs_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.critic.layer_dims, + value_bounds=self.algo_config.critic.value_bounds, + num_atoms=self.algo_config.critic.distributional.num_atoms, + goal_shapes=self.goal_shapes, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + + # Q network ensemble and target ensemble + self.nets["critic"] = nn.ModuleList() + self.nets["critic_target"] = nn.ModuleList() + + # NOTE: ensemble value in config is ignored, and only 1 critic is used. + critic = critic_class(**critic_args) + self.nets["critic"].append(critic) + + critic_target = critic_class(**critic_args) + self.nets["critic_target"].append(critic_target) + + def _get_target_values_from_sampled_actions(self, next_states_tiled, next_sampled_actions, goal_states_tiled, rewards, dones): + """ + Helper function to get target values for training Q-function with TD-loss. Update from superclass + to account for distributional value functions. + + Args: + next_states_tiled (dict): next observations to use for sampling actions. Assumes that + tiling has already occurred - so that if the batch size is B, and N samples are + desired for each observation in the batch, the leading dimension for each + observation in the dict is B * N + + next_sampled_actions (torch.Tensor): actions sampled from the action sampler. This function + will maximize the critic over these action candidates (using the TD3 trick) + + goal_states_tiled (dict): if not None, goal observations + + rewards (torch.Tensor): batch of rewards - should be shape (B, 1) + + dones (torch.Tensor): batch of done signals - should be shape (B, 1) + + Returns: + target_categorical_probabilities (torch.Tensor): target categorical probabilities + to use in the bellman backup + """ + + with torch.no_grad(): + # compute expected returns of the sampled actions and maximize to find the best action + all_vds = self.nets["critic_target"][0].forward_train(next_states_tiled, next_sampled_actions, goal_states_tiled) + expected_values = all_vds.mean().reshape(-1, self.algo_config.critic.num_action_samples) + best_action_index = torch.argmax(expected_values, dim=1) + all_actions = next_sampled_actions.reshape(-1, self.algo_config.critic.num_action_samples, self.ac_dim) + best_action = all_actions[torch.arange(all_actions.shape[0]), best_action_index] + + # get the corresponding probabilities for the categorical distributions corresponding to the best actions + all_vd_probs = all_vds.probs.reshape(-1, self.algo_config.critic.num_action_samples, self.algo_config.critic.distributional.num_atoms) + target_vd_probs = all_vd_probs[torch.arange(all_vd_probs.shape[0]), best_action_index] + + # bellman backup to get a new grid of values - then project onto the canonical atoms to obtain a + # target set of categorical probabilities over the atoms + atom_value_grid = all_vds.values + target_value_grid = rewards + dones * self.discount * atom_value_grid + target_categorical_probabilities = LossUtils.project_values_onto_atoms( + values=target_value_grid, + probabilities=target_vd_probs, + atoms=atom_value_grid, + ) + + return target_categorical_probabilities + + def _compute_critic_loss(self, critic, states, actions, goal_states, q_targets): + """ + Overrides super class to compute a distributional loss. Since values are + categorical distributions, this is just computing a cross-entropy + loss between the two distributions. + + NOTE: q_targets is expected to be a batch of normalized probability vectors that correspond to + the target categorical distributions over the value atoms. + + Args: + critic (torch.nn.Module): critic network + states (dict): batch of observations + actions (torch.Tensor): batch of actions + goal_states (dict): if not None, batch of goal observations + q_targets (torch.Tensor): batch of target q-values for the TD loss + + Returns: + critic_loss (torch.Tensor): critic loss + critic_output (dict): additional outputs from the critic. This function + returns None, but subclasses may want to provide some information + here. + """ + + # this should be the equivalent of softmax with logits from tf + vd = critic.forward_train(states, actions, goal_states) + log_probs = F.log_softmax(vd.logits, dim=-1) + critic_loss = nn.KLDivLoss(reduction='batchmean')(log_probs, q_targets) + return critic_loss, None diff --git a/aloha-devel/robomimic/algo/cql.py b/aloha-devel/robomimic/algo/cql.py new file mode 100644 index 0000000000000000000000000000000000000000..0c24d50abd91426a4d96e91896c958b7df1ada0a --- /dev/null +++ b/aloha-devel/robomimic/algo/cql.py @@ -0,0 +1,668 @@ +""" +Implementation of Conservative Q-Learning (CQL). +Based off of https://github.com/aviralkumar2907/CQL. +(Paper - https://arxiv.org/abs/2006.04779). +""" +import numpy as np +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.optim as optim + +import robomimic.models.base_nets as BaseNets +import robomimic.models.obs_nets as ObsNets +import robomimic.models.policy_nets as PolicyNets +import robomimic.models.value_nets as ValueNets +import robomimic.utils.obs_utils as ObsUtils +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.torch_utils as TorchUtils +from robomimic.algo import register_algo_factory_func, ValueAlgo, PolicyAlgo + + +@register_algo_factory_func("cql") +def algo_config_to_class(algo_config): + """ + Maps algo config to the CQL algo class to instantiate, along with additional algo kwargs. + + Args: + algo_config (Config instance): algo config + + Returns: + algo_class: subclass of Algo + algo_kwargs (dict): dictionary of additional kwargs to pass to algorithm + """ + return CQL, {} + + +class CQL(PolicyAlgo, ValueAlgo): + """ + CQL-extension of SAC for the off-policy, offline setting. See https://arxiv.org/abs/2006.04779 + """ + def __init__(self, **kwargs): + # Store entropy / cql settings first since the super init call requires them + self.automatic_entropy_tuning = kwargs["algo_config"].actor.target_entropy is not None + self.automatic_cql_tuning = kwargs["algo_config"].critic.target_q_gap is not None and \ + kwargs["algo_config"].critic.target_q_gap >= 0.0 + + # Run super init first + super().__init__(**kwargs) + + # Reward settings + self.n_step = self.algo_config.n_step + self.discount = self.algo_config.discount ** self.n_step + + # Now also store additional SAC- and CQL-specific stuff from the config + self._num_batch_steps = 0 + self.bc_start_steps = self.algo_config.actor.bc_start_steps + self.deterministic_backup = self.algo_config.critic.deterministic_backup + self.td_loss_fcn = nn.SmoothL1Loss() if self.algo_config.critic.use_huber else nn.MSELoss() + + # Entropy settings + self.target_entropy = -np.prod(self.ac_dim) if self.algo_config.actor.target_entropy in {None, "default"} else\ + self.algo_config.actor.target_entropy + + # CQL settings + self.min_q_weight = self.algo_config.critic.min_q_weight + self.target_q_gap = self.algo_config.critic.target_q_gap if self.automatic_cql_tuning else 0.0 + + @property + def log_entropy_weight(self): + return self.nets["log_entropy_weight"]() if self.automatic_entropy_tuning else\ + torch.zeros(1, requires_grad=False, device=self.device) + + @property + def log_cql_weight(self): + return self.nets["log_cql_weight"]() if self.automatic_cql_tuning else\ + torch.log(torch.tensor(self.algo_config.critic.cql_weight, requires_grad=False, device=self.device)) + + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + + Networks for this algo: critic (potentially ensemble), policy + """ + + # Create nets + self.nets = nn.ModuleDict() + + # Assemble args to pass to actor + actor_args = dict(self.algo_config.actor.net.common) + + # Add network-specific args and define network class + if self.algo_config.actor.net.type == "gaussian": + actor_cls = PolicyNets.GaussianActorNetwork + actor_args.update(dict(self.algo_config.actor.net.gaussian)) + else: + # Unsupported actor type! + raise ValueError(f"Unsupported actor requested. " + f"Requested: {self.algo_config.actor.net.type}, " + f"valid options are: {['gaussian']}") + + # Policy + self.nets["actor"] = actor_cls( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.actor.layer_dims, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + **actor_args, + ) + + # Critics + self.nets["critic"] = nn.ModuleList() + self.nets["critic_target"] = nn.ModuleList() + for _ in range(self.algo_config.critic.ensemble.n): + for net_list in (self.nets["critic"], self.nets["critic_target"]): + critic = ValueNets.ActionValueNetwork( + obs_shapes=self.obs_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.critic.layer_dims, + value_bounds=self.algo_config.critic.value_bounds, + goal_shapes=self.goal_shapes, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + net_list.append(critic) + + # Entropy (if automatically tuning) + if self.automatic_entropy_tuning: + self.nets["log_entropy_weight"] = BaseNets.Parameter(torch.zeros(1)) + + # CQL (if automatically tuning) + if self.automatic_cql_tuning: + self.nets["log_cql_weight"] = BaseNets.Parameter(torch.zeros(1)) + + # Send networks to appropriate device + self.nets = self.nets.float().to(self.device) + + # sync target networks at beginning of training + with torch.no_grad(): + for critic, critic_target in zip(self.nets["critic"], self.nets["critic_target"]): + TorchUtils.hard_update( + source=critic, + target=critic_target, + ) + + def _create_optimizers(self): + """ + Creates optimizers using @self.optim_params and places them into @self.optimizers. + + Overrides base method since we might need to create aditional optimizers for the entropy + and cql weight parameters (by default, the base class only creates optimizers for all + entries in @self.nets that have corresponding entries in `self.optim_params` but these + parameters do not). + """ + + # Create actor and critic optimizers via super method + super()._create_optimizers() + + # We still need to potentially create additional optimizers based on algo settings + + # entropy (if automatically tuning) + if self.automatic_entropy_tuning: + self.optimizers["entropy"] = optim.Adam( + params=self.nets["log_entropy_weight"].parameters(), + lr=self.optim_params["actor"]["learning_rate"]["initial"], + weight_decay=0.0, + ) + + # cql (if automatically tuning) + if self.automatic_cql_tuning: + self.optimizers["cql"] = optim.Adam( + params=self.nets["log_cql_weight"].parameters(), + lr=self.optim_params["critic"]["learning_rate"]["initial"], + weight_decay=0.0, + ) + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out relevant info and prepare the batch for training. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + input_batch = dict() + + # Make sure the trajectory of actions received is greater than our step horizon + assert batch["actions"].shape[1] >= self.n_step + + # remove temporal batches for all + input_batch["obs"] = {k: batch["obs"][k][:, 0, :] for k in batch["obs"]} + input_batch["next_obs"] = {k: batch["next_obs"][k][:, self.n_step - 1, :] for k in batch["next_obs"]} + input_batch["goal_obs"] = batch.get("goal_obs", None) # goals may not be present + input_batch["actions"] = batch["actions"][:, 0, :] + + # note: ensure scalar signals (rewards, done) retain last dimension of 1 to be compatible with model outputs + + # single timestep reward is discounted sum of intermediate rewards in sequence + reward_seq = batch["rewards"][:, :self.n_step] + discounts = torch.pow(self.algo_config.discount, torch.arange(self.n_step).float()).unsqueeze(0) + input_batch["rewards"] = (reward_seq * discounts).sum(dim=1).unsqueeze(1) + + # consider this n-step seqeunce done if any intermediate dones are present + done_seq = batch["dones"][:, :self.n_step] + input_batch["dones"] = (done_seq.sum(dim=1) > 0).float().unsqueeze(1) + + # we move to device first before float conversion because image observation modalities will be uint8 - + # this minimizes the amount of data transferred to GPU + return TensorUtils.to_float(TensorUtils.to_device(input_batch, self.device)) + + def train_on_batch(self, batch, epoch, validate=False): + """ + Training on a single batch of data. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + info = OrderedDict() + + # Set the correct context for this training step + with TorchUtils.maybe_no_grad(no_grad=validate): + # Always run super call first + super_info = super().train_on_batch(batch, epoch, validate=validate) + # Train actor + actor_info = self._train_policy_on_batch(batch, epoch, validate) + # Train critic(s) + critic_info = self._train_critic_on_batch(batch, epoch, validate) + # Update info + info.update(super_info) + info.update(actor_info) + info.update(critic_info) + + # Return stats + return info + + def _train_policy_on_batch(self, batch, epoch, validate=False): + """ + Training policy on a single batch of data. + + Loss is the ExpValue over sampled states of the (weighted) logprob of a sampled action + under the current policy minus the Q value of associated with the (s, a) combo + + Intuitively, this tries to improve the odds of sampling actions with high Q values while simultaneously + penalizing high probability actions. + + Since we're in the continuous setting, we monte carlo sample. + + Concretely: + Loss = Average[ entropy_weight * logprob(f(eps; s) | s) - Q(s, f(eps; s) ] + + where we use the reparameterization trick with Gaussian function f(*) to parameterize + actions as a function of the sampled noise param eps given input state s + + Additionally, we update the (log) entropy weight parameter if we're tuning that as well. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + info = OrderedDict() + + # Sample actions from policy and get log probs + dist = self.nets["actor"].forward_train(obs_dict=batch["obs"], goal_dict=batch["goal_obs"]) + actions, log_prob = self._get_actions_and_log_prob(dist=dist) + + # Calculate alpha + entropy_weight_loss = -(self.log_entropy_weight * (log_prob + self.target_entropy).detach()).mean() if\ + self.automatic_entropy_tuning else 0.0 + entropy_weight = self.log_entropy_weight.exp() + + # Get predicted Q-values for all state, action pairs + pred_qs = [critic(obs_dict=batch["obs"], acts=actions, goal_dict=batch["goal_obs"]) + for critic in self.nets["critic"]] + # We take the minimum for stability + pred_qs, _ = torch.cat(pred_qs, dim=1).min(dim=1, keepdim=True) + + # Use BC if we're in the beginning of training, otherwise calculate policy loss normally + baseline = dist.log_prob(batch["actions"]).unsqueeze(dim=-1) if\ + self._num_batch_steps < self.bc_start_steps else pred_qs + policy_loss = (entropy_weight * log_prob - baseline).mean() + + # Add info + info["entropy_weight"] = entropy_weight.item() + info["entropy_weight_loss"] = entropy_weight_loss.item() if \ + self.automatic_entropy_tuning else entropy_weight_loss + info["actor/loss"] = policy_loss + + # Take a training step if we're not validating + if not validate: + # Update batch step + self._num_batch_steps += 1 + if self.automatic_entropy_tuning: + # Alpha + self.optimizers["entropy"].zero_grad() + entropy_weight_loss.backward() + self.optimizers["entropy"].step() + info["entropy_grad_norms"] = self.log_entropy_weight.grad.data.norm(2).pow(2).item() + + # Policy + actor_grad_norms = TorchUtils.backprop_for_loss( + net=self.nets["actor"], + optim=self.optimizers["actor"], + loss=policy_loss, + max_grad_norm=self.algo_config.actor.max_gradient_norm, + ) + # Add info + info["actor/grad_norms"] = actor_grad_norms + + # Return stats + return info + + def _train_critic_on_batch(self, batch, epoch, validate=False): + """ + Training critic(s) on a single batch of data. + + For a given batch of (s, a, r, s') tuples and n sampled actions (a_, a'_ corresponding to actions + sampled from the learned policy at states s and s', respectively; a~ corresponding to uniformly random + sampled actions): + + Loss = CQL_loss + SAC_loss + + Since we're in the continuous setting, we monte carlo sample for all ExpValues, which become Averages instead + + SAC_loss is the standard single-step TD error, corresponding to the following: + + SAC_loss = 0.5 * Average[ (Q(s,a) - (r + Average over a'_ [ Q(s', a'_) ]))^2 ] + + The CQL_loss corresponds to a weighted secondary objective, corresponding to the (ExpValue of Q values over + sampled states and sampled actions from the LEARNED policy) minus the (ExpValue of Q values over + sampled states and sampled actions from the DATASET policy) plus a regularizer as a function + of the learned policy. + + Intuitively, this tries to penalize Q-values arbitrarily resulting from the learned policy (which may produce + out-of-distribution (s,a) pairs) while preserving (known) Q-values taken from the dataset policy. + + As we are using SAC, we choose our regularizer to correspond to the negative KL divergence between our + learned policy and a uniform distribution such that the first term in the CQL loss corresponds to the + soft maximum over all Q values at any state s. + + For stability, we importance sample actions over random actions and from the current policy at s, s'. + + Moreover, if we want to tune the cql_weight automatically, we include the threshold value target_q_gap + to penalize Q values that are overly-optimistic by the given threshold. + + In this case, the CQL_loss is as follows: + + CQL_loss = cql_weight * (Average [log (Average over a` in {a~, a_, a_'}: exp(Q(s,a`) - logprob(a`)) - Average [Q(s,a)]] - target_q_gap) + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + info = OrderedDict() + B, A = batch["actions"].shape + N = self.algo_config.critic.num_random_actions + + # Get predicted Q-values from taken actions + q_preds = [critic(obs_dict=batch["obs"], acts=batch["actions"], goal_dict=batch["goal_obs"]) + for critic in self.nets["critic"]] + + # Sample actions at the current and next step + curr_dist = self.nets["actor"].forward_train(obs_dict=batch["obs"], goal_dict=batch["goal_obs"]) + next_dist = self.nets["actor"].forward_train(obs_dict=batch["next_obs"], goal_dict=batch["goal_obs"]) + next_actions, next_log_prob = self._get_actions_and_log_prob(dist=next_dist) + + # Don't capture gradients here, since the critic target network doesn't get trained (only soft updated) + with torch.no_grad(): + # We take the max over all samples if the number of action samples is > 1 + if self.algo_config.critic.num_action_samples > 1: + # Generate the target q values, using the backup from the next state + temp_actions = next_dist.rsample(sample_shape=(self.algo_config.critic.num_action_samples,)).permute(1, 0, 2) + target_qs = [self._get_qs_from_actions( + obs_dict=batch["next_obs"], actions=temp_actions, goal_dict=batch["goal_obs"], q_net=critic) + .max(dim=1, keepdim=True)[0] for critic in self.nets["critic_target"]] + else: + target_qs = [critic(obs_dict=batch["next_obs"], acts=next_actions, goal_dict=batch["goal_obs"]) + for critic in self.nets["critic_target"]] + # Take the minimum over all critics + target_qs, _ = torch.cat(target_qs, dim=1).min(dim=1, keepdim=True) + # If only sampled once from each critic and not using a deterministic backup, subtract the logprob as well + if self.algo_config.critic.num_action_samples == 1 and not self.deterministic_backup: + target_qs = target_qs - self.log_entropy_weight.exp() * next_log_prob + + # Calculate the q target values + done_mask_batch = 1. - batch["dones"] + info["done_masks"] = done_mask_batch + q_target = batch["rewards"] + done_mask_batch * self.discount * target_qs + + # Calculate CQL stuff + cql_random_actions = torch.FloatTensor(N, B, A).uniform_(-1., 1.).to(self.device) # shape (N, B, A) + cql_random_log_prob = np.log(0.5 ** A) + cql_curr_actions, cql_curr_log_prob = self._get_actions_and_log_prob(dist=curr_dist, sample_shape=(N,)) # shape (N, B, A) and (N, B, 1) + cql_next_actions, cql_next_log_prob = self._get_actions_and_log_prob(dist=next_dist, sample_shape=(N,)) # shape (N, B, A) and (N, B, 1) + cql_curr_log_prob = cql_curr_log_prob.squeeze(dim=-1).permute(1, 0).detach() # shape (B, N) + cql_next_log_prob = cql_next_log_prob.squeeze(dim=-1).permute(1, 0).detach() # shape (B, N) + q_cats = [] # Each entry shape will be (B, N) + + for critic, q_pred in zip(self.nets["critic"], q_preds): + # Compose Q values over all sampled actions (importance sampled) + q_rand = self._get_qs_from_actions(obs_dict=batch["obs"], actions=cql_random_actions.permute(1, 0, 2), goal_dict=batch["goal_obs"], q_net=critic) + q_curr = self._get_qs_from_actions(obs_dict=batch["obs"], actions=cql_curr_actions.permute(1, 0, 2), goal_dict=batch["goal_obs"], q_net=critic) + q_next = self._get_qs_from_actions(obs_dict=batch["obs"], actions=cql_next_actions.permute(1, 0, 2), goal_dict=batch["goal_obs"], q_net=critic) + q_cat = torch.cat([ + q_rand - cql_random_log_prob, + q_next - cql_next_log_prob, + q_curr - cql_curr_log_prob, + ], dim=1) # shape (B, 3 * N) + q_cats.append(q_cat) + + # Calculate the losses for all critics + cql_losses = [] + critic_losses = [] + cql_weight = torch.clamp(self.log_cql_weight.exp(), min=0.0, max=1000000.0) + info["critic/cql_weight"] = cql_weight.item() + for i, (q_pred, q_cat) in enumerate(zip(q_preds, q_cats)): + # Calculate td error loss + td_loss = self.td_loss_fcn(q_pred, q_target) + # Calculate cql loss + cql_loss = cql_weight * (self.min_q_weight * (torch.logsumexp(q_cat, dim=1).mean() - q_pred.mean()) - + self.target_q_gap) + cql_losses.append(cql_loss) + # Calculate total loss + loss = td_loss + cql_loss + critic_losses.append(loss) + info[f"critic/critic{i+1}_loss"] = loss + + # Run gradient descent if we're not validating + if not validate: + # Train CQL weight if tuning automatically + if self.automatic_cql_tuning: + cql_weight_loss = -torch.stack(cql_losses).mean() + info[ + "critic/cql_weight_loss"] = cql_weight_loss.item() # Make sure to not store computation graph since we retain graph after backward() call + self.optimizers["cql"].zero_grad() + cql_weight_loss.backward(retain_graph=True) + self.optimizers["cql"].step() + info["critic/cql_grad_norms"] = self.log_cql_weight.grad.data.norm(2).pow(2).item() + + # Train critics + for i, (critic_loss, critic, critic_target, optimizer) in enumerate(zip( + critic_losses, self.nets["critic"], self.nets["critic_target"], self.optimizers["critic"] + )): + retain_graph = (i < (len(critic_losses) - 1)) + critic_grad_norms = TorchUtils.backprop_for_loss( + net=critic, + optim=optimizer, + loss=critic_loss, + max_grad_norm=self.algo_config.critic.max_gradient_norm, + retain_graph=retain_graph, + ) + info[f"critic/critic{i+1}_grad_norms"] = critic_grad_norms + with torch.no_grad(): + TorchUtils.soft_update(source=critic, target=critic_target, tau=self.algo_config.target_tau) + + # Return stats + return info + + def _get_actions_and_log_prob(self, dist, sample_shape=torch.Size()): + """ + Helper method to sample actions and compute corresponding log probabilities + + Args: + dist (Distribution): Distribution to sample from + sample_shape (torch.Size or tuple): Shape of output when sampling (number of samples) + + Returns: + 2-tuple: + - (tensor) sampled actions (..., B, ..., A) + - (tensor) corresponding log probabilities (..., B, ..., 1) + """ + # Process networks with tanh differently than normal distributions + if self.algo_config.actor.net.common.use_tanh: + actions, actions_pre_tanh = dist.rsample(sample_shape=sample_shape, return_pretanh_value=True) + log_prob = dist.log_prob(actions, pre_tanh_value=actions_pre_tanh).unsqueeze(dim=-1) + else: + actions = dist.rsample(sample_shape=sample_shape) + log_prob = dist.log_prob(actions) + + return actions, log_prob + + @staticmethod + def _get_qs_from_actions(obs_dict, actions, goal_dict, q_net): + """ + Helper function for grabbing Q values given a single state and multiple (N) sampled actions. + + Args: + obs_dict (dict): Observation dict from batch + actions (tensor): Torch tensor, with dim1 assumed to be the extra sampled dimension + goal_dict (dict): Goal dict from batch + q_net (nn.Module): Q net to pass the observations and actions + + Returns: + tensor: (B, N) corresponding Q values + """ + # Get the number of sampled actions + B, N, D = actions.shape + + # Repeat obs and goals in the batch dimension + obs_dict_stacked = ObsUtils.repeat_and_stack_observation(obs_dict, N) + goal_dict_stacked = ObsUtils.repeat_and_stack_observation(goal_dict, N) + + # Pass the obs and (flattened) actions through to get the Q values + qs = q_net(obs_dict=obs_dict_stacked, acts=actions.reshape(-1, D), goal_dict=goal_dict_stacked) + + # Unflatten output + qs = qs.reshape(B, N) + + return qs + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + loss_log = OrderedDict() + + # record current optimizer learning rates + for k in self.optimizers: + keys = [k] + optims = [self.optimizers[k]] + if k == "critic": + # account for critic having one optimizer per ensemble member + keys = ["{}{}".format(k, critic_ind) for critic_ind in range(len(self.nets["critic"]))] + optims = self.optimizers[k] + for kp, optimizer in zip(keys, optims): + for i, param_group in enumerate(optimizer.param_groups): + loss_log["Optimizer/{}{}_lr".format(kp, i)] = param_group["lr"] + + # extract relevant logs for critic, and actor + loss_log["Loss"] = 0. + for loss_logger in [self._log_critic_info, self._log_actor_info]: + this_log = loss_logger(info) + if "Loss" in this_log: + # manually merge total loss + loss_log["Loss"] += this_log["Loss"] + del this_log["Loss"] + loss_log.update(this_log) + + return loss_log + + def _log_critic_info(self, info): + """ + Helper function to extract critic-relevant information for logging. + """ + loss_log = OrderedDict() + if "done_masks" in info: + loss_log["Critic/Done_Mask_Percentage"] = 100. * torch.mean(info["done_masks"]).item() + if "critic/q_targets" in info: + loss_log["Critic/Q_Targets"] = info["critic/q_targets"].mean().item() + loss_log["Loss"] = 0. + for critic_ind in range(len(self.nets["critic"])): + loss_log["Critic/Critic{}_Loss".format(critic_ind + 1)] = info["critic/critic{}_loss".format(critic_ind + 1)].item() + if "critic/critic{}_grad_norms".format(critic_ind + 1) in info: + loss_log["Critic/Critic{}_Grad_Norms".format(critic_ind + 1)] = info["critic/critic{}_grad_norms".format(critic_ind + 1)] + loss_log["Loss"] += loss_log["Critic/Critic{}_Loss".format(critic_ind + 1)] + if "critic/cql_weight_loss" in info: + loss_log["Critic/CQL_Weight"] = info["critic/cql_weight"] + loss_log["Critic/CQL_Weight_Loss"] = info["critic/cql_weight_loss"] + loss_log["Critic/CQL_Grad_Norms"] = info["critic/cql_grad_norms"] + return loss_log + + def _log_actor_info(self, info): + """ + Helper function to extract actor-relevant information for logging. + """ + loss_log = OrderedDict() + loss_log["Actor/Loss"] = info["actor/loss"].item() + if "actor/grad_norms" in info: + loss_log["Actor/Grad_Norms"] = info["actor/grad_norms"] + loss_log["Loss"] = loss_log["Actor/Loss"] + loss_log["Entropy_Weight_Loss"] = info["entropy_weight_loss"] + loss_log["Entropy_Weight"] = info["entropy_weight"] + if "entropy_grad_norms" in info: + loss_log["Entropy_Grad_Norms"] = info["entropy_grad_norms"] + return loss_log + + def set_train(self): + """ + Prepare networks for evaluation. Update from super class to make sure + target networks stay in evaluation mode all the time. + """ + self.nets.train() + + # target networks always in eval + for critic in self.nets["critic_target"]: + critic.eval() + + def on_epoch_end(self, epoch): + """ + Called at the end of each epoch. + """ + + # LR scheduling updates + for lr_sc in self.lr_schedulers["critic"]: + if lr_sc is not None: + lr_sc.step() + + if self.lr_schedulers["actor"] is not None: + self.lr_schedulers["actor"].step() + + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + action (torch.Tensor): action tensor + """ + assert not self.nets.training + + return self.nets["actor"](obs_dict=obs_dict, goal_dict=goal_dict) + + def get_state_action_value(self, obs_dict, actions, goal_dict=None): + """ + Get state-action value outputs. + + Args: + obs_dict (dict): current observation + actions (torch.Tensor): action + goal_dict (dict): (optional) goal + + Returns: + value (torch.Tensor): value tensor + """ + assert not self.nets.training + + return self.nets["critic"][0](obs_dict, actions, goal_dict) diff --git a/aloha-devel/robomimic/algo/diffusion_policy.py b/aloha-devel/robomimic/algo/diffusion_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..574fa792ebd28c50bbcddd68afdb7bd15369cf1f --- /dev/null +++ b/aloha-devel/robomimic/algo/diffusion_policy.py @@ -0,0 +1,700 @@ +""" +Implementation of Diffusion Policy https://diffusion-policy.cs.columbia.edu/ by Cheng Chi +""" +from typing import Callable, Union +import math +from collections import OrderedDict, deque +from packaging.version import parse as parse_version +import random +import torch +import torch.nn as nn +import torch.nn.functional as F +# requires diffusers==0.11.1 +from diffusers.schedulers.scheduling_ddpm import DDPMScheduler +from diffusers.schedulers.scheduling_ddim import DDIMScheduler +from diffusers.training_utils import EMAModel + +import robomimic.models.obs_nets as ObsNets +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.torch_utils as TorchUtils +import robomimic.utils.obs_utils as ObsUtils + +from robomimic.algo import register_algo_factory_func, PolicyAlgo + +import random +import robomimic.utils.torch_utils as TorchUtils +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.obs_utils as ObsUtils + +@register_algo_factory_func("diffusion_policy") +def algo_config_to_class(algo_config): + """ + Maps algo config to the BC algo class to instantiate, along with additional algo kwargs. + + Args: + algo_config (Config instance): algo config + + Returns: + algo_class: subclass of Algo + algo_kwargs (dict): dictionary of additional kwargs to pass to algorithm + """ + + if algo_config.unet.enabled: + return DiffusionPolicyUNet, {} + elif algo_config.transformer.enabled: + raise NotImplementedError() + else: + raise RuntimeError() + +class DiffusionPolicyUNet(PolicyAlgo): + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + if self.algo_config.language_conditioned: + self.obs_shapes["lang_emb"] = [768] # clip is 768-dim embedding + + # set up different observation groups for @MIMO_MLP + observation_group_shapes = OrderedDict() + observation_group_shapes["obs"] = OrderedDict(self.obs_shapes) + encoder_kwargs = ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder) + + obs_encoder = ObsNets.ObservationGroupEncoder( + observation_group_shapes=observation_group_shapes, + encoder_kwargs=encoder_kwargs, + ) + # IMPORTANT! + # replace all BatchNorm with GroupNorm to work with EMA + # performance will tank if you forget to do this! + obs_encoder = replace_bn_with_gn(obs_encoder) + + obs_dim = obs_encoder.output_shape()[0] + + # create network object + noise_pred_net = ConditionalUnet1D( + input_dim=self.ac_dim, + global_cond_dim=obs_dim*self.algo_config.horizon.observation_horizon + ) + + # the final arch has 2 parts + nets = nn.ModuleDict({ + 'policy': nn.ModuleDict({ + 'obs_encoder': obs_encoder, + 'noise_pred_net': noise_pred_net + }) + }) + + nets = nets.float().to(self.device) + + # setup noise scheduler + noise_scheduler = None + if self.algo_config.ddpm.enabled: + noise_scheduler = DDPMScheduler( + num_train_timesteps=self.algo_config.ddpm.num_train_timesteps, + beta_schedule=self.algo_config.ddpm.beta_schedule, + clip_sample=self.algo_config.ddpm.clip_sample, + prediction_type=self.algo_config.ddpm.prediction_type + ) + elif self.algo_config.ddim.enabled: + noise_scheduler = DDIMScheduler( + num_train_timesteps=self.algo_config.ddim.num_train_timesteps, + beta_schedule=self.algo_config.ddim.beta_schedule, + clip_sample=self.algo_config.ddim.clip_sample, + set_alpha_to_one=self.algo_config.ddim.set_alpha_to_one, + steps_offset=self.algo_config.ddim.steps_offset, + prediction_type=self.algo_config.ddim.prediction_type + ) + else: + raise RuntimeError() + + # setup EMA + ema = None + if self.algo_config.ema.enabled: + ema = EMAModel(model=nets, power=self.algo_config.ema.power) + + # set attrs + self.nets = nets + self.noise_scheduler = noise_scheduler + self.ema = ema + self.action_check_done = False + self.obs_queue = None + self.action_queue = None + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + To = self.algo_config.horizon.observation_horizon + Ta = self.algo_config.horizon.action_horizon + Tp = self.algo_config.horizon.prediction_horizon + + input_batch = dict() + input_batch["obs"] = {k: batch["obs"][k][:, :To, :] for k in batch["obs"]} + input_batch["goal_obs"] = batch.get("goal_obs", None) # goals may not be present + input_batch["actions"] = batch["actions"][:, :Tp, :] + + # check if actions are normalized to [-1,1] + if not self.action_check_done: + actions = input_batch["actions"] + in_range = (-1 <= actions) & (actions <= 1) + all_in_range = torch.all(in_range).item() + if not all_in_range: + raise ValueError('"actions" must be in range [-1,1] for Diffusion Policy! Check if hdf5_normalize_action is enabled.') + self.action_check_done = True + + return TensorUtils.to_device(TensorUtils.to_float(input_batch), self.device) + + def train_on_batch(self, batch, epoch, validate=False): + """ + Training on a single batch of data. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + To = self.algo_config.horizon.observation_horizon + Ta = self.algo_config.horizon.action_horizon + Tp = self.algo_config.horizon.prediction_horizon + action_dim = self.ac_dim + B = batch['actions'].shape[0] + + + with TorchUtils.maybe_no_grad(no_grad=validate): + info = super(DiffusionPolicyUNet, self).train_on_batch(batch, epoch, validate=validate) + actions = batch['actions'] + + # encode obs + inputs = { + 'obs': batch["obs"], + 'goal': batch["goal_obs"] + } + for k in self.obs_shapes: + # first two dimensions should be [B, T] for inputs + assert inputs['obs'][k].ndim - 2 == len(self.obs_shapes[k]) + + obs_features = TensorUtils.time_distributed(inputs, self.nets['policy']['obs_encoder'], inputs_as_kwargs=True) + assert obs_features.ndim == 3 # [B, T, D] + + obs_cond = obs_features.flatten(start_dim=1) + + # sample noise to add to actions + noise = torch.randn(actions.shape, device=self.device) + + # sample a diffusion iteration for each data point + timesteps = torch.randint( + 0, self.noise_scheduler.config.num_train_timesteps, + (B,), device=self.device + ).long() + + # add noise to the clean actions according to the noise magnitude at each diffusion iteration + # (this is the forward diffusion process) + noisy_actions = self.noise_scheduler.add_noise( + actions, noise, timesteps) + + # predict the noise residual + noise_pred = self.nets['policy']['noise_pred_net']( + noisy_actions, timesteps, global_cond=obs_cond) + + # L2 loss + loss = F.mse_loss(noise_pred, noise) + + # logging + losses = { + 'l2_loss': loss + } + info["losses"] = TensorUtils.detach(losses) + + if not validate: + # gradient step + policy_grad_norms = TorchUtils.backprop_for_loss( + net=self.nets, + optim=self.optimizers["policy"], + loss=loss, + ) + + # update Exponential Moving Average of the model weights + if self.ema is not None: + self.ema.step(self.nets) + + step_info = { + 'policy_grad_norms': policy_grad_norms + } + info.update(step_info) + + return info + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + log = super(DiffusionPolicyUNet, self).log_info(info) + log["Loss"] = info["losses"]["l2_loss"].item() + if "policy_grad_norms" in info: + log["Policy_Grad_Norms"] = info["policy_grad_norms"] + return log + + def reset(self): + """ + Reset algo state to prepare for environment rollouts. + """ + # setup inference queues + To = self.algo_config.horizon.observation_horizon + Ta = self.algo_config.horizon.action_horizon + obs_queue = deque(maxlen=To) + action_queue = deque(maxlen=Ta) + self.obs_queue = obs_queue + self.action_queue = action_queue + + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + + Args: + obs_dict (dict): current observation [1, Do] + goal_dict (dict): (optional) goal + + Returns: + action (torch.Tensor): action tensor [1, Da] + """ + # obs_dict: key: [1,D] + To = self.algo_config.horizon.observation_horizon + Ta = self.algo_config.horizon.action_horizon + + # TODO: obs_queue already handled by frame_stack + # make sure we have at least To observations in obs_queue + # if not enough, repeat + # if already full, append one to the obs_queue + # n_repeats = max(To - len(self.obs_queue), 1) + # self.obs_queue.extend([obs_dict] * n_repeats) + + if len(self.action_queue) == 0: + # no actions left, run inference + # turn obs_queue into dict of tensors (concat at T dim) + # import pdb; pdb.set_trace() + # obs_dict_list = TensorUtils.list_of_flat_dict_to_dict_of_list(list(self.obs_queue)) + # obs_dict_tensor = dict((k, torch.cat(v, dim=0).unsqueeze(0)) for k,v in obs_dict_list.items()) + + # run inference + # [1,T,Da] + action_sequence = self._get_action_trajectory(obs_dict=obs_dict) + + # put actions into the queue + self.action_queue.extend(action_sequence[0]) + + # has action, execute from left to right + # [Da] + action = self.action_queue.popleft() + + # [1,Da] + action = action.unsqueeze(0) + return action + + def _get_action_trajectory(self, obs_dict, goal_dict=None): + assert not self.nets.training + To = self.algo_config.horizon.observation_horizon + Ta = self.algo_config.horizon.action_horizon + Tp = self.algo_config.horizon.prediction_horizon + action_dim = self.ac_dim + if self.algo_config.ddpm.enabled is True: + num_inference_timesteps = self.algo_config.ddpm.num_inference_timesteps + elif self.algo_config.ddim.enabled is True: + num_inference_timesteps = self.algo_config.ddim.num_inference_timesteps + else: + raise ValueError + + # select network + nets = self.nets + if self.ema is not None: + nets = self.ema.averaged_model + + # encode obs + inputs = { + 'obs': obs_dict, + 'goal': goal_dict + } + for k in self.obs_shapes: + # first two dimensions should be [B, T] for inputs + assert inputs['obs'][k].ndim - 2 == len(self.obs_shapes[k]) + obs_features = TensorUtils.time_distributed(inputs, nets['policy']['obs_encoder'], inputs_as_kwargs=True) + assert obs_features.ndim == 3 # [B, T, D] + B = obs_features.shape[0] + + # reshape observation to (B,obs_horizon*obs_dim) + obs_cond = obs_features.flatten(start_dim=1) + + # initialize action from Guassian noise + noisy_action = torch.randn( + (B, Tp, action_dim), device=self.device) + naction = noisy_action + + # init scheduler + self.noise_scheduler.set_timesteps(num_inference_timesteps) + + for k in self.noise_scheduler.timesteps: + # predict noise + noise_pred = nets['policy']['noise_pred_net']( + sample=naction, + timestep=k, + global_cond=obs_cond + ) + + # inverse diffusion step (remove noise) + naction = self.noise_scheduler.step( + model_output=noise_pred, + timestep=k, + sample=naction + ).prev_sample + + # process action using Ta + start = To - 1 + end = start + Ta + action = naction[:,start:end] + return action + + def serialize(self): + """ + Get dictionary of current model parameters. + """ + return { + "nets": self.nets.state_dict(), + "ema": self.ema.averaged_model.state_dict() if self.ema is not None else None, + } + + def deserialize(self, model_dict): + """ + Load model from a checkpoint. + + Args: + model_dict (dict): a dictionary saved by self.serialize() that contains + the same keys as @self.network_classes + """ + self.nets.load_state_dict(model_dict["nets"]) + if model_dict.get("ema", None) is not None: + self.ema.averaged_model.load_state_dict(model_dict["ema"]) + + + + + +# =================== Vision Encoder Utils ===================== +def replace_submodules( + root_module: nn.Module, + predicate: Callable[[nn.Module], bool], + func: Callable[[nn.Module], nn.Module]) -> nn.Module: + """ + Replace all submodules selected by the predicate with + the output of func. + + predicate: Return true if the module is to be replaced. + func: Return new module to use. + """ + if predicate(root_module): + return func(root_module) + + if parse_version(torch.__version__) < parse_version('1.9.0'): + raise ImportError('This function requires pytorch >= 1.9.0') + + bn_list = [k.split('.') for k, m + in root_module.named_modules(remove_duplicate=True) + if predicate(m)] + for *parent, k in bn_list: + parent_module = root_module + if len(parent) > 0: + parent_module = root_module.get_submodule('.'.join(parent)) + if isinstance(parent_module, nn.Sequential): + src_module = parent_module[int(k)] + else: + src_module = getattr(parent_module, k) + tgt_module = func(src_module) + if isinstance(parent_module, nn.Sequential): + parent_module[int(k)] = tgt_module + else: + setattr(parent_module, k, tgt_module) + # verify that all modules are replaced + bn_list = [k.split('.') for k, m + in root_module.named_modules(remove_duplicate=True) + if predicate(m)] + assert len(bn_list) == 0 + return root_module + +def replace_bn_with_gn( + root_module: nn.Module, + features_per_group: int=16) -> nn.Module: + """ + Relace all BatchNorm layers with GroupNorm. + """ + replace_submodules( + root_module=root_module, + predicate=lambda x: isinstance(x, nn.BatchNorm2d), + func=lambda x: nn.GroupNorm( + num_groups=x.num_features//features_per_group, + num_channels=x.num_features) + ) + return root_module + +# =================== UNet for Diffusion ============== + +class SinusoidalPosEmb(nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x): + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device) * -emb) + emb = x[:, None] * emb[None, :] + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class Downsample1d(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = nn.Conv1d(dim, dim, 3, 2, 1) + + def forward(self, x): + return self.conv(x) + +class Upsample1d(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1) + + def forward(self, x): + return self.conv(x) + + +class Conv1dBlock(nn.Module): + ''' + Conv1d --> GroupNorm --> Mish + ''' + + def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8): + super().__init__() + + self.block = nn.Sequential( + nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2), + nn.GroupNorm(n_groups, out_channels), + nn.Mish(), + ) + + def forward(self, x): + return self.block(x) + + +class ConditionalResidualBlock1D(nn.Module): + def __init__(self, + in_channels, + out_channels, + cond_dim, + kernel_size=3, + n_groups=8): + super().__init__() + + self.blocks = nn.ModuleList([ + Conv1dBlock(in_channels, out_channels, kernel_size, n_groups=n_groups), + Conv1dBlock(out_channels, out_channels, kernel_size, n_groups=n_groups), + ]) + + # FiLM modulation https://arxiv.org/abs/1709.07871 + # predicts per-channel scale and bias + cond_channels = out_channels * 2 + self.out_channels = out_channels + self.cond_encoder = nn.Sequential( + nn.Mish(), + nn.Linear(cond_dim, cond_channels), + nn.Unflatten(-1, (-1, 1)) + ) + + # make sure dimensions compatible + self.residual_conv = nn.Conv1d(in_channels, out_channels, 1) \ + if in_channels != out_channels else nn.Identity() + + def forward(self, x, cond): + ''' + x : [ batch_size x in_channels x horizon ] + cond : [ batch_size x cond_dim] + + returns: + out : [ batch_size x out_channels x horizon ] + ''' + out = self.blocks[0](x) + embed = self.cond_encoder(cond) + + embed = embed.reshape( + embed.shape[0], 2, self.out_channels, 1) + scale = embed[:,0,...] + bias = embed[:,1,...] + out = scale * out + bias + + out = self.blocks[1](out) + out = out + self.residual_conv(x) + return out + + +class ConditionalUnet1D(nn.Module): + def __init__(self, + input_dim, + global_cond_dim, + diffusion_step_embed_dim=256, + down_dims=[256,512,1024], + kernel_size=5, + n_groups=8 + ): + """ + input_dim: Dim of actions. + global_cond_dim: Dim of global conditioning applied with FiLM + in addition to diffusion step embedding. This is usually obs_horizon * obs_dim + diffusion_step_embed_dim: Size of positional encoding for diffusion iteration k + down_dims: Channel size for each UNet level. + The length of this array determines numebr of levels. + kernel_size: Conv kernel size + n_groups: Number of groups for GroupNorm + """ + + super().__init__() + all_dims = [input_dim] + list(down_dims) + start_dim = down_dims[0] + + dsed = diffusion_step_embed_dim + diffusion_step_encoder = nn.Sequential( + SinusoidalPosEmb(dsed), + nn.Linear(dsed, dsed * 4), + nn.Mish(), + nn.Linear(dsed * 4, dsed), + ) + cond_dim = dsed + global_cond_dim + + in_out = list(zip(all_dims[:-1], all_dims[1:])) + mid_dim = all_dims[-1] + self.mid_modules = nn.ModuleList([ + ConditionalResidualBlock1D( + mid_dim, mid_dim, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups + ), + ConditionalResidualBlock1D( + mid_dim, mid_dim, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups + ), + ]) + + down_modules = nn.ModuleList([]) + for ind, (dim_in, dim_out) in enumerate(in_out): + is_last = ind >= (len(in_out) - 1) + down_modules.append(nn.ModuleList([ + ConditionalResidualBlock1D( + dim_in, dim_out, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups), + ConditionalResidualBlock1D( + dim_out, dim_out, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups), + Downsample1d(dim_out) if not is_last else nn.Identity() + ])) + + up_modules = nn.ModuleList([]) + for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): + is_last = ind >= (len(in_out) - 1) + up_modules.append(nn.ModuleList([ + ConditionalResidualBlock1D( + dim_out*2, dim_in, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups), + ConditionalResidualBlock1D( + dim_in, dim_in, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups), + Upsample1d(dim_in) if not is_last else nn.Identity() + ])) + + final_conv = nn.Sequential( + Conv1dBlock(start_dim, start_dim, kernel_size=kernel_size), + nn.Conv1d(start_dim, input_dim, 1), + ) + + self.diffusion_step_encoder = diffusion_step_encoder + self.up_modules = up_modules + self.down_modules = down_modules + self.final_conv = final_conv + + print("number of parameters: {:e}".format( + sum(p.numel() for p in self.parameters())) + ) + + def forward(self, + sample: torch.Tensor, + timestep: Union[torch.Tensor, float, int], + global_cond=None): + """ + x: (B,T,input_dim) + timestep: (B,) or int, diffusion step + global_cond: (B,global_cond_dim) + output: (B,T,input_dim) + """ + # (B,T,C) + sample = sample.moveaxis(-1,-2) + # (B,C,T) + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device) + elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + global_feature = self.diffusion_step_encoder(timesteps) + + if global_cond is not None: + global_feature = torch.cat([ + global_feature, global_cond + ], axis=-1) + + x = sample + h = [] + for idx, (resnet, resnet2, downsample) in enumerate(self.down_modules): + x = resnet(x, global_feature) + x = resnet2(x, global_feature) + h.append(x) + x = downsample(x) + + for mid_module in self.mid_modules: + x = mid_module(x, global_feature) + + for idx, (resnet, resnet2, upsample) in enumerate(self.up_modules): + x = torch.cat((x, h.pop()), dim=1) + x = resnet(x, global_feature) + x = resnet2(x, global_feature) + x = upsample(x) + + x = self.final_conv(x) + + # (B,C,T) + x = x.moveaxis(-1,-2) + # (B,T,C) + return x diff --git a/aloha-devel/robomimic/algo/gl.py b/aloha-devel/robomimic/algo/gl.py new file mode 100644 index 0000000000000000000000000000000000000000..24ae800892ee0866f9b4df3d94ff49eb1cd8d112 --- /dev/null +++ b/aloha-devel/robomimic/algo/gl.py @@ -0,0 +1,775 @@ +""" +Subgoal prediction models, used in HBC / IRIS. +""" +import numpy as np +from collections import OrderedDict +from copy import deepcopy + +import torch +import torch.nn as nn + +import robomimic.models.obs_nets as ObsNets +import robomimic.models.vae_nets as VAENets +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.torch_utils as TorchUtils +import robomimic.utils.obs_utils as ObsUtils + +from robomimic.algo import register_algo_factory_func, PlannerAlgo, ValueAlgo + + +@register_algo_factory_func("gl") +def algo_config_to_class(algo_config): + """ + Maps algo config to the GL algo class to instantiate, along with additional algo kwargs. + + Args: + algo_config (Config instance): algo config + + Returns: + algo_class: subclass of Algo + algo_kwargs (dict): dictionary of additional kwargs to pass to algorithm + """ + if algo_config.vae.enabled: + return GL_VAE, {} + return GL, {} + + +class GL(PlannerAlgo): + """ + Implements goal prediction component for HBC and IRIS. + """ + def __init__( + self, + algo_config, + obs_config, + global_config, + obs_key_shapes, + ac_dim, + device + ): + """ + Args: + algo_config (Config object): instance of Config corresponding to the algo section + of the config + + obs_config (Config object): instance of Config corresponding to the observation + section of the config + + global_config (Config object): global training config + + obs_key_shapes (OrderedDict): dictionary that maps observation keys to shapes + + ac_dim (int): dimension of action space + + device (torch.Device): where the algo should live (i.e. cpu, gpu) + """ + + self._subgoal_horizon = algo_config.subgoal_horizon + super(GL, self).__init__( + algo_config=algo_config, + obs_config=obs_config, + global_config=global_config, + obs_key_shapes=obs_key_shapes, + ac_dim=ac_dim, + device=device + ) + + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + self.nets = nn.ModuleDict() + + obs_group_shapes = OrderedDict() + obs_group_shapes["obs"] = OrderedDict(self.obs_shapes) + if len(self.goal_shapes) > 0: + obs_group_shapes["goal"] = OrderedDict(self.goal_shapes) + + # deterministic goal prediction network + self.nets["goal_network"] = ObsNets.MIMO_MLP( + input_obs_group_shapes=obs_group_shapes, + output_shapes=self.subgoal_shapes, + layer_dims=self.algo_config.ae.planner_layer_dims, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + + self.nets = self.nets.float().to(self.device) + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + input_batch = dict() + + # remove temporal batches for all except scalar signals (to be compatible with model outputs) + input_batch["obs"] = { k: batch["obs"][k][:, 0, :] for k in batch["obs"] } + # extract multi-horizon subgoal target + input_batch["subgoals"] = {k: batch["next_obs"][k][:, self._subgoal_horizon - 1, :] for k in batch["next_obs"]} + input_batch["target_subgoals"] = input_batch["subgoals"] + input_batch["goal_obs"] = batch.get("goal_obs", None) # goals may not be present + + # we move to device first before float conversion because image observation modalities will be uint8 - + # this minimizes the amount of data transferred to GPU + return TensorUtils.to_float(TensorUtils.to_device(input_batch, self.device)) + + def get_actor_goal_for_training_from_processed_batch(self, processed_batch, **kwargs): + """ + Retrieve subgoals from processed batch to use for training the actor. Subclasses + can modify this function to change the subgoals. + + Args: + processed_batch (dict): processed batch from @process_batch_for_training + + Returns: + actor_subgoals (dict): subgoal observations to condition actor on + """ + return processed_batch["target_subgoals"] + + def train_on_batch(self, batch, epoch, validate=False): + """ + Training on a single batch of data. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + with TorchUtils.maybe_no_grad(no_grad=validate): + info = super(GL, self).train_on_batch(batch, epoch, validate=validate) + + # predict subgoal observations with goal network + pred_subgoals = self.nets["goal_network"](obs=batch["obs"], goal=batch["goal_obs"]) + + # compute loss as L2 error for each observation key + losses = OrderedDict() + target_subgoals = batch["target_subgoals"] # targets for network prediction + goal_loss = 0. + for k in pred_subgoals: + assert pred_subgoals[k].shape == target_subgoals[k].shape, "mismatch in predicted and target subgoals!" + mode_loss = nn.MSELoss()(pred_subgoals[k], target_subgoals[k]) + goal_loss += mode_loss + losses["goal_{}_loss".format(k)] = mode_loss + losses["goal_loss"] = goal_loss + info.update(TensorUtils.detach(losses)) + + if not validate: + # gradient step + goal_grad_norms = TorchUtils.backprop_for_loss( + net=self.nets["goal_network"], + optim=self.optimizers["goal_network"], + loss=losses["goal_loss"], + ) + info["goal_grad_norms"] = goal_grad_norms + + return info + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + loss_log = super(GL, self).log_info(info) + + loss_log["Loss"] = info["goal_loss"].item() + for k in info: + if k.endswith("_loss"): + loss_log[k] = info[k].item() + if "goal_grad_norms" in info: + loss_log["Grad_Norms"] = info["goal_grad_norms"] + + return loss_log + + def get_subgoal_predictions(self, obs_dict, goal_dict=None): + """ + Takes a batch of observations and predicts a batch of subgoals. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + subgoal prediction (dict): name -> Tensor [batch_size, ...] + """ + return self.nets["goal_network"](obs=obs_dict, goal=goal_dict) + + def sample_subgoals(self, obs_dict, goal_dict=None, num_samples=1): + """ + Sample @num_samples subgoals from the network per observation. + Since this class implements a deterministic subgoal prediction, + this function returns identical subgoals for each input observation. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + subgoals (dict): name -> Tensor [batch_size, num_samples, ...] + """ + + # stack observations to get all samples in one forward pass + obs_tiled = ObsUtils.repeat_and_stack_observation(obs_dict, n=num_samples) + goal_tiled = None + if goal_dict is not None: + goal_tiled = ObsUtils.repeat_and_stack_observation(goal_dict, n=num_samples) + + # [batch_size * num_samples, ...] + goals = self.get_subgoal_predictions(obs_dict=obs_tiled, goal_dict=goal_tiled) + # reshape to [batch_size, num_samples, ...] + return TensorUtils.reshape_dimensions(goals, begin_axis=0, end_axis=0, target_dims=(-1, num_samples)) + + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. Assumes one input observation (first dimension should be 1). + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + action (torch.Tensor): action tensor + """ + raise Exception("Rollouts are not supported by GL") + + +class GL_VAE(GL): + """ + Implements goal prediction via VAE. + """ + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + self.nets = nn.ModuleDict() + + self.nets["goal_network"] = VAENets.VAE( + input_shapes=self.subgoal_shapes, + output_shapes=self.subgoal_shapes, + condition_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + device=self.device, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + **VAENets.vae_args_from_config(self.algo_config.vae), + ) + + self.nets = self.nets.float().to(self.device) + + def get_actor_goal_for_training_from_processed_batch( + self, + processed_batch, + use_latent_subgoals=False, + use_prior_correction=False, + num_prior_samples=100, + **kwargs, + ): + """ + Modify from superclass to support a @use_latent_subgoals option. + The VAE can optionally return latent subgoals by passing the subgoal + observations in the batch through the encoder. + + Args: + processed_batch (dict): processed batch from @process_batch_for_training + + use_latent_subgoals (bool): if True, condition the actor on latent subgoals + by using the VAE encoder to encode subgoal observations at train-time, + and using the VAE prior to generate latent subgoals at test-time + + use_prior_correction (bool): if True, use a "prior correction" trick to + choose a latent subgoal sampled from the prior that is close to the + latent from the VAE encoder (posterior). This can help with issues at + test-time where the encoder latent distribution might not match + the prior latent distribution. + + num_prior_samples (int): number of VAE prior samples to take and choose among, + if @use_prior_correction is true + + Returns: + actor_subgoals (dict): subgoal observations to condition actor on + """ + + if not use_latent_subgoals: + return processed_batch["target_subgoals"] + + # batch variables + obs = processed_batch["obs"] + subgoals = processed_batch["subgoals"] # full subgoal observations + target_subgoals = processed_batch["target_subgoals"] # targets for network prediction + goal_obs = processed_batch["goal_obs"] + + with torch.no_grad(): + # run VAE forward pass to get samples from posterior for the current observation and subgoal + vae_outputs = self.nets["goal_network"]( + inputs=subgoals, # encoder takes full subgoals + outputs=target_subgoals, # reconstruct target subgoals + goals=goal_obs, + conditions=obs, # condition on observations + ) + posterior_z = vae_outputs["encoder_z"] + latent_subgoals = posterior_z + + if use_prior_correction: + # instead of treating posterior samples as latent subgoals, sample latents from + # the prior and choose the closest one as the latent subgoal + + random_key = list(obs.keys())[0] + batch_size = obs[random_key].shape[0] + + # for each batch member, get @num_prior_samples samples from the prior + obs_tiled = ObsUtils.repeat_and_stack_observation(obs, n=num_prior_samples) + goal_tiled = None + if len(self.goal_shapes) > 0: + goal_tiled = ObsUtils.repeat_and_stack_observation(goal_obs, n=num_prior_samples) + + prior_z_samples = self.nets["goal_network"].sample_prior( + conditions=obs_tiled, + goals=goal_tiled, + ) + + # choose prior samples that are closest to the sampled posterior latents + # note: every posterior sample in the batch has @num_prior_samples corresponding prior samples + + # reshape prior samples to (batch_size, num_samples, latent_dim) + prior_z_samples = prior_z_samples.reshape(batch_size, num_prior_samples, -1) + + # reshape posterior latents to (batch_size, 1, latent_dim) + posterior_z_expanded = posterior_z.unsqueeze(1) + + # compute distances with broadcasting so that each posterior sample + # has distances to all of its prior samples + distances = (prior_z_samples - posterior_z_expanded).pow(2).sum(dim=2) + + # then gather the closest prior sample for each posterior sample + neighbors = torch.argmin(distances, dim=1) + latent_subgoals = prior_z_samples[torch.arange(batch_size).long(), neighbors] + + return { "latent_subgoal" : latent_subgoals } + + def train_on_batch(self, batch, epoch, validate=False): + """ + Training on a single batch of data. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + with TorchUtils.maybe_no_grad(no_grad=validate): + info = super(GL, self).train_on_batch(batch, epoch, validate=validate) + + if self.algo_config.vae.prior.use_categorical: + temperature = self.algo_config.vae.prior.categorical_init_temp - epoch * self.algo_config.vae.prior.categorical_temp_anneal_step + temperature = max(temperature, self.algo_config.vae.prior.categorical_min_temp) + self.nets["goal_network"].set_gumbel_temperature(temperature) + + # batch variables + obs = batch["obs"] + subgoals = batch["subgoals"] # full subgoal observations + target_subgoals = batch["target_subgoals"] # targets for network prediction + goal_obs = batch["goal_obs"] + + vae_outputs = self.nets["goal_network"]( + inputs=subgoals, # encoder takes full subgoals + outputs=target_subgoals, # reconstruct target subgoals + goals=goal_obs, + conditions=obs, # condition on observations + ) + recons_loss = vae_outputs["reconstruction_loss"] + kl_loss = vae_outputs["kl_loss"] + goal_loss = recons_loss + self.algo_config.vae.kl_weight * kl_loss + info["recons_loss"] = recons_loss + info["kl_loss"] = kl_loss + info["goal_loss"] = goal_loss + + if not self.algo_config.vae.prior.use_categorical: + with torch.no_grad(): + info["encoder_variance"] = torch.exp(vae_outputs["encoder_params"]["logvar"]) + + # VAE gradient step + if not validate: + goal_grad_norms = TorchUtils.backprop_for_loss( + net=self.nets["goal_network"], + optim=self.optimizers["goal_network"], + loss=goal_loss, + ) + info["goal_grad_norms"] = goal_grad_norms + + return info + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + loss_log = super(GL_VAE, self).log_info(info) + loss_log["Reconstruction_Loss"] = info["recons_loss"].item() + loss_log["KL_Loss"] = info["kl_loss"].item() + if self.algo_config.vae.prior.use_categorical: + loss_log["Gumbel_Temperature"] = self.nets["goal_network"].get_gumbel_temperature() + else: + loss_log["Encoder_Variance"] = info["encoder_variance"].mean().item() + return loss_log + + def get_subgoal_predictions(self, obs_dict, goal_dict=None): + """ + Takes a batch of observations and predicts a batch of subgoals. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + subgoal prediction (dict): name -> Tensor [batch_size, ...] + """ + + if self.global_config.algo.latent_subgoal.enabled: + # latent subgoals from sampling prior + latent_subgoals = self.nets["goal_network"].sample_prior( + conditions=obs_dict, + goals=goal_dict, + ) + + return OrderedDict(latent_subgoal=latent_subgoals) + + # sample a single goal from the VAE + goals = self.sample_subgoals(obs_dict=obs_dict, goal_dict=goal_dict, num_samples=1) + return { k : goals[k][:, 0, ...] for k in goals } + + def sample_subgoals(self, obs_dict, goal_dict=None, num_samples=1): + """ + Sample @num_samples subgoals from the VAE per observation. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + subgoals (dict): name -> Tensor [batch_size, num_samples, ...] + """ + + # stack observations to get all samples in one forward pass + obs_tiled = ObsUtils.repeat_and_stack_observation(obs_dict, n=num_samples) + goal_tiled = None + if goal_dict is not None: + goal_tiled = ObsUtils.repeat_and_stack_observation(goal_dict, n=num_samples) + + # VAE decode expects number of samples explicitly + mod = list(obs_tiled.keys())[0] + n = obs_tiled[mod].shape[0] + # [batch_size * num_samples, ...] + goals = self.nets["goal_network"].decode(n=n, conditions=obs_tiled, goals=goal_tiled) + # reshape to [batch_size, num_samples, ...] + return TensorUtils.reshape_dimensions(goals, begin_axis=0, end_axis=0, target_dims=(-1, num_samples)) + + +class ValuePlanner(PlannerAlgo, ValueAlgo): + """ + Base class for all algorithms that are used for planning subgoals + based on (1) a @PlannerAlgo that is used to sample candidate subgoals + and (2) a @ValueAlgo that is used to select one of the subgoals. + """ + def __init__( + self, + planner_algo_class, + value_algo_class, + algo_config, + obs_config, + global_config, + obs_key_shapes, + ac_dim, + device, + + ): + """ + Args: + planner_algo_class (Algo class): algo class for the planner + + value_algo_class (Algo class): algo class for the value network + + algo_config (Config object): instance of Config corresponding to the algo section + of the config + + obs_config (Config object): instance of Config corresponding to the observation + section of the config + + global_config (Config object); global config + + obs_key_shapes (OrderedDict): dictionary that maps input/output observation keys to shapes + + ac_dim (int): action dimension + + device: torch device + """ + self.algo_config = algo_config + self.obs_config = obs_config + self.global_config = global_config + + self.ac_dim = ac_dim + self.device = device + + self.planner = planner_algo_class( + algo_config=algo_config.planner, + obs_config=obs_config.planner, + global_config=global_config, + obs_key_shapes=obs_key_shapes, + ac_dim=ac_dim, + device=device + ) + + self.value_net = value_algo_class( + algo_config=algo_config.value, + obs_config=obs_config.value, + global_config=global_config, + obs_key_shapes=obs_key_shapes, + ac_dim=ac_dim, + device=device + ) + + self.subgoal_shapes = self.planner.subgoal_shapes + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + input_batch = dict() + + input_batch["planner"] = self.planner.process_batch_for_training(batch) + input_batch["value_net"] = self.value_net.process_batch_for_training(batch) + + # we move to device first before float conversion because image observation modalities will be uint8 - + # this minimizes the amount of data transferred to GPU + return TensorUtils.to_float(TensorUtils.to_device(input_batch, self.device)) + + def train_on_batch(self, batch, epoch, validate=False): + """ + Training on a single batch of data. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + if validate: + assert not self.planner.nets.training + assert not self.value_net.nets.training + + info = dict(planner=dict(), value_net=dict()) + + # train planner + info["planner"].update(self.planner.train_on_batch(batch["planner"], epoch, validate=validate)) + + # train value network + info["value_net"].update(self.value_net.train_on_batch(batch["value_net"], epoch, validate=validate)) + + return info + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + loss = 0. + + # planner + planner_log = self.planner.log_info(info["planner"]) + planner_log = dict(("Planner/" + k, v) for k, v in planner_log.items()) + loss += planner_log["Planner/Loss"] + + # value network + value_net_log = self.value_net.log_info(info["value_net"]) + value_net_log = dict(("ValueNetwork/" + k, v) for k, v in value_net_log.items()) + loss += value_net_log["ValueNetwork/Loss"] + planner_log.update(value_net_log) + + planner_log["Loss"] = loss + return planner_log + + def on_epoch_end(self, epoch): + """ + Called at the end of each epoch. + """ + self.planner.on_epoch_end(epoch) + self.value_net.on_epoch_end(epoch) + + def set_eval(self): + """ + Prepare networks for evaluation. + """ + self.planner.set_eval() + self.value_net.set_eval() + + def set_train(self): + """ + Prepare networks for training. + """ + self.planner.set_train() + self.value_net.set_train() + + def serialize(self): + """ + Get dictionary of current model parameters. + """ + return dict( + planner=self.planner.serialize(), + value_net=self.value_net.serialize(), + ) + + def deserialize(self, model_dict): + """ + Load model from a checkpoint. + + Args: + model_dict (dict): a dictionary saved by self.serialize() that contains + the same keys as @self.network_classes + """ + self.planner.deserialize(model_dict["planner"]) + self.value_net.deserialize(model_dict["value_net"]) + + def reset(self): + """ + Reset algo state to prepare for environment rollouts. + """ + self.planner.reset() + self.value_net.reset() + + def __repr__(self): + """ + Pretty print algorithm and network description. + """ + msg = str(self.__class__.__name__) + import textwrap + return msg + "Planner:\n" + textwrap.indent(self.planner.__repr__(), ' ') + \ + "\n\nValue Network:\n" + textwrap.indent(self.value_net.__repr__(), ' ') + + def get_subgoal_predictions(self, obs_dict, goal_dict=None): + """ + Takes a batch of observations and predicts a batch of subgoals. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + subgoal prediction (dict): name -> Tensor [batch_size, ...] + """ + + num_samples = self.algo_config.num_samples + + # sample subgoals from the planner (shape: [batch_size, num_samples, ...]) + subgoals = self.sample_subgoals(obs_dict=obs_dict, goal_dict=goal_dict, num_samples=num_samples) + + # stack subgoals to get all values in one forward pass (shape [batch_size * num_samples, ...]) + k = list(obs_dict.keys())[0] + bsize = obs_dict[k].shape[0] + subgoals_tiled = TensorUtils.reshape_dimensions(subgoals, begin_axis=0, end_axis=1, target_dims=(bsize * num_samples,)) + + # also repeat goals if necessary + goal_tiled = None + if len(self.planner.goal_shapes) > 0: + goal_tiled = ObsUtils.repeat_and_stack_observation(goal_dict, n=num_samples) + + # evaluate the value of each subgoal + subgoal_values = self.value_net.get_state_value(obs_dict=subgoals_tiled, goal_dict=goal_tiled).reshape(-1, num_samples) + + # pick the best subgoal + best_index = torch.argmax(subgoal_values, dim=1) + best_subgoal = {k: subgoals[k][torch.arange(bsize), best_index] for k in subgoals} + return best_subgoal + + def sample_subgoals(self, obs_dict, goal_dict, num_samples=1): + """ + Sample @num_samples subgoals from the planner algo per observation. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + subgoals (dict): name -> Tensor [batch_size, num_samples, ...] + """ + return self.planner.sample_subgoals(obs_dict=obs_dict, goal_dict=goal_dict, num_samples=num_samples) + + def get_state_value(self, obs_dict, goal_dict=None): + """ + Get state value outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + value (torch.Tensor): value tensor + """ + return self.value_net.get_state_value(obs_dict=obs_dict, goal_dict=goal_dict) + + def get_state_action_value(self, obs_dict, actions, goal_dict=None): + """ + Get state-action value outputs. + + Args: + obs_dict (dict): current observation + actions (torch.Tensor): action + goal_dict (dict): (optional) goal + + Returns: + value (torch.Tensor): value tensor + """ + return self.value_net.get_state_action_value(obs_dict=obs_dict, actions=actions, goal_dict=goal_dict) diff --git a/aloha-devel/robomimic/algo/hbc.py b/aloha-devel/robomimic/algo/hbc.py new file mode 100644 index 0000000000000000000000000000000000000000..543b1fbcf4ced11b9628d506b1972f1123a357b6 --- /dev/null +++ b/aloha-devel/robomimic/algo/hbc.py @@ -0,0 +1,344 @@ +""" +Implementation of Hierarchical Behavioral Cloning, where +a planner model outputs subgoals (future observations), and +an actor model is conditioned on the subgoals to try and +reach them. Largely based on the Generalization Through Imitation (GTI) +paper (see https://arxiv.org/abs/2003.06085). +""" +import textwrap +import numpy as np +from collections import OrderedDict +from copy import deepcopy + +import torch + +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.obs_utils as ObsUtils +from robomimic.config.config import Config +from robomimic.algo import register_algo_factory_func, algo_name_to_factory_func, HierarchicalAlgo, GL_VAE + + +@register_algo_factory_func("hbc") +def algo_config_to_class(algo_config): + """ + Maps algo config to the HBC algo class to instantiate, along with additional algo kwargs. + + Args: + algo_config (Config instance): algo config + + Returns: + algo_class: subclass of Algo + algo_kwargs (dict): dictionary of additional kwargs to pass to algorithm + """ + pol_cls, _ = algo_name_to_factory_func("bc")(algo_config.actor) + plan_cls, _ = algo_name_to_factory_func("gl")(algo_config.planner) + return HBC, dict(policy_algo_class=pol_cls, planner_algo_class=plan_cls) + + +class HBC(HierarchicalAlgo): + """ + Default HBC training, largely based on https://arxiv.org/abs/2003.06085 + """ + def __init__( + self, + planner_algo_class, + policy_algo_class, + algo_config, + obs_config, + global_config, + obs_key_shapes, + ac_dim, + device, + ): + """ + Args: + planner_algo_class (Algo class): algo class for the planner + + policy_algo_class (Algo class): algo class for the policy + + algo_config (Config object): instance of Config corresponding to the algo section + of the config + + obs_config (Config object): instance of Config corresponding to the observation + section of the config + + global_config (Config object): global training config + + obs_key_shapes (dict): dictionary that maps input/output observation keys to shapes + + ac_dim (int): action dimension + + device: torch device + """ + self.algo_config = algo_config + self.obs_config = obs_config + self.global_config = global_config + + self.ac_dim = ac_dim + self.device = device + + self._subgoal_step_count = 0 # current step count for deciding when to update subgoal + self._current_subgoal = None # latest subgoal + self._subgoal_update_interval = self.algo_config.subgoal_update_interval # subgoal update frequency + self._subgoal_horizon = self.algo_config.planner.subgoal_horizon + self._actor_horizon = self.algo_config.actor.rnn.horizon + + self._algo_mode = self.algo_config.mode + assert self._algo_mode in ["separate", "cascade"] + + self.planner = planner_algo_class( + algo_config=algo_config.planner, + obs_config=obs_config.planner, + global_config=global_config, + obs_key_shapes=obs_key_shapes, + ac_dim=ac_dim, + device=device + ) + + # goal-conditional actor follows goals set by the planner + self.actor_goal_shapes = self.planner.subgoal_shapes + if self.algo_config.latent_subgoal.enabled: + assert planner_algo_class == GL_VAE # only VAE supported for now + self.actor_goal_shapes = OrderedDict(latent_subgoal=(self.planner.algo_config.vae.latent_dim,)) + + # only for the actor: override goal modalities and shapes to match the subgoal set by the planner + actor_obs_key_shapes = deepcopy(obs_key_shapes) + # make sure we are not modifying existing observation key shapes + for k in self.actor_goal_shapes: + if k in actor_obs_key_shapes: + assert actor_obs_key_shapes[k] == self.actor_goal_shapes[k] + actor_obs_key_shapes.update(self.actor_goal_shapes) + + goal_obs_keys = {obs_modality: [] for obs_modality in ObsUtils.OBS_MODALITY_CLASSES.keys()} + for k in self.actor_goal_shapes.keys(): + goal_obs_keys[ObsUtils.OBS_KEYS_TO_MODALITIES[k]].append(k) + + actor_obs_config = deepcopy(obs_config.actor) + with actor_obs_config.unlocked(): + actor_obs_config["goal"] = Config(**goal_obs_keys) + + self.actor = policy_algo_class( + algo_config=algo_config.actor, + obs_config=actor_obs_config, + global_config=global_config, + obs_key_shapes=actor_obs_key_shapes, + ac_dim=ac_dim, + device=device, + ) + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + input_batch = dict() + + input_batch["planner"] = self.planner.process_batch_for_training(batch) + input_batch["actor"] = self.actor.process_batch_for_training(batch) + + if self.algo_config.actor_use_random_subgoals: + # optionally use randomly sampled step between [1, seq_length] as policy goal + policy_subgoal_indices = torch.randint( + low=0, high=self.global_config.train.seq_length, size=(batch["actions"].shape[0],)) + goal_obs = TensorUtils.gather_sequence(batch["next_obs"], policy_subgoal_indices) + goal_obs = TensorUtils.to_float(TensorUtils.to_device(goal_obs, self.device)) + input_batch["actor"]["goal_obs"] = \ + self.planner.get_actor_goal_for_training_from_processed_batch( + goal_obs, + use_latent_subgoals=self.algo_config.latent_subgoal.enabled, + use_prior_correction=self.algo_config.latent_subgoal.prior_correction.enabled, + num_prior_samples=self.algo_config.latent_subgoal.prior_correction.num_samples, + ) + else: + # otherwise, use planner subgoal target as goal for the policy + input_batch["actor"]["goal_obs"] = \ + self.planner.get_actor_goal_for_training_from_processed_batch( + input_batch["planner"], + use_latent_subgoals=self.algo_config.latent_subgoal.enabled, + use_prior_correction=self.algo_config.latent_subgoal.prior_correction.enabled, + num_prior_samples=self.algo_config.latent_subgoal.prior_correction.num_samples, + ) + + # we move to device first before float conversion because image observation modalities will be uint8 - + # this minimizes the amount of data transferred to GPU + return TensorUtils.to_float(TensorUtils.to_device(input_batch, self.device)) + + def train_on_batch(self, batch, epoch, validate=False): + """ + Training on a single batch of data. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + info = dict(planner=dict(), actor=dict()) + # train planner + info["planner"].update(self.planner.train_on_batch(batch["planner"], epoch, validate=validate)) + + # train actor + if self._algo_mode == "separate": + # train low-level actor by getting subgoals from the dataset + info["actor"].update(self.actor.train_on_batch(batch["actor"], epoch, validate=validate)) + + elif self._algo_mode == "cascade": + # get predictions from the planner + with torch.no_grad(): + batch["actor"]["goal_obs"] = self.planner.get_subgoal_predictions( + obs_dict=batch["planner"]["obs"], goal_dict=batch["planner"]["goal_obs"]) + + # train actor with the predicted goal + info["actor"].update(self.actor.train_on_batch(batch["actor"], epoch, validate=validate)) + + else: + raise NotImplementedError("algo mode {} is not implemented".format(self._algo_mode)) + + return info + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + planner_log = dict() + actor_log = dict() + loss = 0. + + planner_log = self.planner.log_info(info["planner"]) + planner_log = dict(("Planner/" + k, v) for k, v in planner_log.items()) + loss += planner_log["Planner/Loss"] + + actor_log = self.actor.log_info(info["actor"]) + actor_log = dict(("Actor/" + k, v) for k, v in actor_log.items()) + loss += actor_log["Actor/Loss"] + + planner_log.update(actor_log) + planner_log["Loss"] = loss + return planner_log + + def on_epoch_end(self, epoch): + """ + Called at the end of each epoch. + """ + self.planner.on_epoch_end(epoch) + self.actor.on_epoch_end(epoch) + + def set_eval(self): + """ + Prepare networks for evaluation. + """ + self.planner.set_eval() + self.actor.set_eval() + + def set_train(self): + """ + Prepare networks for training. + """ + self.planner.set_train() + self.actor.set_train() + + def serialize(self): + """ + Get dictionary of current model parameters. + """ + return dict( + planner=self.planner.serialize(), + actor=self.actor.serialize(), + ) + + def deserialize(self, model_dict): + """ + Load model from a checkpoint. + + Args: + model_dict (dict): a dictionary saved by self.serialize() that contains + the same keys as @self.network_classes + """ + self.actor.deserialize(model_dict["actor"]) + self.planner.deserialize(model_dict["planner"]) + + @property + def current_subgoal(self): + """ + Return the current subgoal (at rollout time) with shape (batch, ...) + """ + return { k : self._current_subgoal[k].clone() for k in self._current_subgoal } + + @current_subgoal.setter + def current_subgoal(self, sg): + """ + Sets the current subgoal being used by the actor. + """ + for k, v in sg.items(): + if not self.algo_config.latent_subgoal.enabled: + # subgoal should only match subgoal shapes if not using latent subgoals + assert list(v.shape[1:]) == list(self.planner.subgoal_shapes[k]) + # subgoal shapes should always match actor goal shapes + assert list(v.shape[1:]) == list(self.actor_goal_shapes[k]) + self._current_subgoal = { k : sg[k].clone() for k in sg } + + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + action (torch.Tensor): action tensor + """ + if self._current_subgoal is None or self._subgoal_step_count % self._subgoal_update_interval == 0: + # update current subgoal + self.current_subgoal = self.planner.get_subgoal_predictions(obs_dict=obs_dict, goal_dict=goal_dict) + + action = self.actor.get_action(obs_dict=obs_dict, goal_dict=self.current_subgoal) + self._subgoal_step_count += 1 + return action + + def reset(self): + """ + Reset algo state to prepare for environment rollouts. + """ + self._current_subgoal = None + self._subgoal_step_count = 0 + self.planner.reset() + self.actor.reset() + + def __repr__(self): + """ + Pretty print algorithm and network description. + """ + msg = str(self.__class__.__name__) + msg += "(subgoal_horizon={}, actor_horizon={}, subgoal_update_interval={}, mode={}, " \ + "actor_use_random_subgoals={})\n".format( + self._subgoal_horizon, + self._actor_horizon, + self._subgoal_update_interval, + self._algo_mode, + self.algo_config.actor_use_random_subgoals + ) + return msg + "Planner:\n" + textwrap.indent(self.planner.__repr__(), ' ') + \ + "\n\nPolicy:\n" + textwrap.indent(self.actor.__repr__(), ' ') diff --git a/aloha-devel/robomimic/algo/iql.py b/aloha-devel/robomimic/algo/iql.py new file mode 100644 index 0000000000000000000000000000000000000000..bde522b2292e6140b5ce4e3120ad0c83e4064fff --- /dev/null +++ b/aloha-devel/robomimic/algo/iql.py @@ -0,0 +1,428 @@ +""" +Implementation of Implicit Q-Learning (IQL). +Based off of https://github.com/rail-berkeley/rlkit/blob/master/rlkit/torch/sac/iql_trainer.py. +(Paper - https://arxiv.org/abs/2110.06169). +""" +import numpy as np +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import robomimic.models.policy_nets as PolicyNets +import robomimic.models.value_nets as ValueNets +import robomimic.utils.obs_utils as ObsUtils +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.torch_utils as TorchUtils +from robomimic.algo import register_algo_factory_func, ValueAlgo, PolicyAlgo + + +@register_algo_factory_func("iql") +def algo_config_to_class(algo_config): + """ + Maps algo config to the IQL algo class to instantiate, along with additional algo kwargs. + + Args: + algo_config (Config instance): algo config + + Returns: + algo_class: subclass of Algo + algo_kwargs (dict): dictionary of additional kwargs to pass to algorithm + """ + return IQL, {} + + +class IQL(PolicyAlgo, ValueAlgo): + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + + Networks for this algo: critic (potentially ensemble), actor, value function + """ + + # Create nets + self.nets = nn.ModuleDict() + + # Assemble args to pass to actor + actor_args = dict(self.algo_config.actor.net.common) + + # Add network-specific args and define network class + if self.algo_config.actor.net.type == "gaussian": + actor_cls = PolicyNets.GaussianActorNetwork + actor_args.update(dict(self.algo_config.actor.net.gaussian)) + elif self.algo_config.actor.net.type == "gmm": + actor_cls = PolicyNets.GMMActorNetwork + actor_args.update(dict(self.algo_config.actor.net.gmm)) + else: + # Unsupported actor type! + raise ValueError(f"Unsupported actor requested. " + f"Requested: {self.algo_config.actor.net.type}, " + f"valid options are: {['gaussian', 'gmm']}") + + # Actor + self.nets["actor"] = actor_cls( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.actor.layer_dims, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + **actor_args, + ) + + # Critics + self.nets["critic"] = nn.ModuleList() + self.nets["critic_target"] = nn.ModuleList() + for _ in range(self.algo_config.critic.ensemble.n): + for net_list in (self.nets["critic"], self.nets["critic_target"]): + critic = ValueNets.ActionValueNetwork( + obs_shapes=self.obs_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.critic.layer_dims, + goal_shapes=self.goal_shapes, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + net_list.append(critic) + + # Value function network + self.nets["vf"] = ValueNets.ValueNetwork( + obs_shapes=self.obs_shapes, + mlp_layer_dims=self.algo_config.critic.layer_dims, + goal_shapes=self.goal_shapes, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + + # Send networks to appropriate device + self.nets = self.nets.float().to(self.device) + + # sync target networks at beginning of training + with torch.no_grad(): + for critic, critic_target in zip(self.nets["critic"], self.nets["critic_target"]): + TorchUtils.hard_update( + source=critic, + target=critic_target, + ) + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out relevant info and prepare the batch for training. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + + input_batch = dict() + + # remove temporal batches for all + input_batch["obs"] = {k: batch["obs"][k][:, 0, :] for k in batch["obs"]} + input_batch["next_obs"] = {k: batch["next_obs"][k][:, 0, :] for k in batch["next_obs"]} + input_batch["goal_obs"] = batch.get("goal_obs", None) # goals may not be present + input_batch["actions"] = batch["actions"][:, 0, :] + input_batch["dones"] = batch["dones"][:, 0] + input_batch["rewards"] = batch["rewards"][:, 0] + + return TensorUtils.to_device(TensorUtils.to_float(input_batch), self.device) + + def train_on_batch(self, batch, epoch, validate=False): + """ + Training on a single batch of data. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + info = OrderedDict() + + # Set the correct context for this training step + with TorchUtils.maybe_no_grad(no_grad=validate): + # Always run super call first + info = super().train_on_batch(batch, epoch, validate=validate) + + # Compute loss for critic(s) + critic_losses, vf_loss, critic_info = self._compute_critic_loss(batch) + # Compute loss for actor + actor_loss, actor_info = self._compute_actor_loss(batch, critic_info) + + if not validate: + # Critic update + self._update_critic(critic_losses, vf_loss) + + # Actor update + self._update_actor(actor_loss) + + # Update info + info.update(actor_info) + info.update(critic_info) + + # Return stats + return info + + def _compute_critic_loss(self, batch): + """ + Helper function for computing Q and V losses. Called by @train_on_batch + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + Returns: + critic_losses (list): list of critic (Q function) losses + vf_loss (torch.Tensor): value function loss + info (dict): dictionary of Q / V predictions and losses + """ + info = OrderedDict() + + # get batch values + obs = batch["obs"] + actions = batch["actions"] + next_obs = batch["next_obs"] + goal_obs = batch["goal_obs"] + rewards = torch.unsqueeze(batch["rewards"], 1) + dones = torch.unsqueeze(batch["dones"], 1) + + # Q predictions + pred_qs = [critic(obs_dict=obs, acts=actions, goal_dict=goal_obs) + for critic in self.nets["critic"]] + + info["critic/critic1_pred"] = pred_qs[0].mean() + + # Q target values + target_vf_pred = self.nets["vf"](obs_dict=next_obs, goal_dict=goal_obs).detach() + q_target = rewards + (1. - dones) * self.algo_config.discount * target_vf_pred + q_target = q_target.detach() + + # Q losses + critic_losses = [] + td_loss_fcn = nn.SmoothL1Loss() if self.algo_config.critic.use_huber else nn.MSELoss() + for (i, q_pred) in enumerate(pred_qs): + # Calculate td error loss + td_loss = td_loss_fcn(q_pred, q_target) + info[f"critic/critic{i+1}_loss"] = td_loss + critic_losses.append(td_loss) + + # V predictions + pred_qs = [critic(obs_dict=obs, acts=actions, goal_dict=goal_obs) + for critic in self.nets["critic_target"]] + q_pred, _ = torch.cat(pred_qs, dim=1).min(dim=1, keepdim=True) + q_pred = q_pred.detach() + vf_pred = self.nets["vf"](obs) + + # V losses: expectile regression. see section 4.1 in https://arxiv.org/pdf/2110.06169.pdf + vf_err = vf_pred - q_pred + vf_sign = (vf_err > 0).float() + vf_weight = (1 - vf_sign) * self.algo_config.vf_quantile + vf_sign * (1 - self.algo_config.vf_quantile) + vf_loss = (vf_weight * (vf_err ** 2)).mean() + + # update logs for V loss + info["vf/q_pred"] = q_pred + info["vf/v_pred"] = vf_pred + info["vf/v_loss"] = vf_loss + + # Return stats + return critic_losses, vf_loss, info + + def _update_critic(self, critic_losses, vf_loss): + """ + Helper function for updating critic and vf networks. Called by @train_on_batch + + Args: + critic_losses (list): list of critic (Q function) losses + vf_loss (torch.Tensor): value function loss + """ + + # update ensemble of critics + for (critic_loss, critic, critic_target, optimizer) in zip( + critic_losses, self.nets["critic"], self.nets["critic_target"], self.optimizers["critic"] + ): + TorchUtils.backprop_for_loss( + net=critic, + optim=optimizer, + loss=critic_loss, + max_grad_norm=self.algo_config.critic.max_gradient_norm, + retain_graph=False, + ) + + # update target network + with torch.no_grad(): + TorchUtils.soft_update(source=critic, target=critic_target, tau=self.algo_config.target_tau) + + # update V function network + TorchUtils.backprop_for_loss( + net=self.nets["vf"], + optim=self.optimizers["vf"], + loss=vf_loss, + max_grad_norm=self.algo_config.critic.max_gradient_norm, + retain_graph=False, + ) + + def _compute_actor_loss(self, batch, critic_info): + """ + Helper function for computing actor loss. Called by @train_on_batch + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + critic_info (dict): dictionary containing Q and V function predictions, + to be used for computing advantage estimates + + Returns: + actor_loss (torch.Tensor): actor loss + info (dict): dictionary of actor losses, log_probs, advantages, and weights + """ + info = OrderedDict() + + # compute log probability of batch actions + dist = self.nets["actor"].forward_train(obs_dict=batch["obs"], goal_dict=batch["goal_obs"]) + log_prob = dist.log_prob(batch["actions"]) + + info["actor/log_prob"] = log_prob.mean() + + # compute advantage estimate + q_pred = critic_info["vf/q_pred"] + v_pred = critic_info["vf/v_pred"] + adv = q_pred - v_pred + + # compute weights + weights = self._get_adv_weights(adv) + + # compute advantage weighted actor loss. disable gradients through weights + actor_loss = (-log_prob * weights.detach()).mean() + + info["actor/loss"] = actor_loss + + # log adv-related values + info["adv/adv"] = adv + info["adv/adv_weight"] = weights + + # Return stats + return actor_loss, info + + def _update_actor(self, actor_loss): + """ + Helper function for updating actor network. Called by @train_on_batch + + Args: + actor_loss (torch.Tensor): actor loss + """ + + TorchUtils.backprop_for_loss( + net=self.nets["actor"], + optim=self.optimizers["actor"], + loss=actor_loss, + max_grad_norm=self.algo_config.actor.max_gradient_norm, + ) + + def _get_adv_weights(self, adv): + """ + Helper function for computing advantage weights. Called by @_compute_actor_loss + + Args: + adv (torch.Tensor): raw advantage estimates + + Returns: + weights (torch.Tensor): weights computed based on advantage estimates, + in shape (B,) where B is batch size + """ + + # clip raw advantage values + if self.algo_config.adv.clip_adv_value is not None: + adv = adv.clamp(max=self.algo_config.adv.clip_adv_value) + + # compute weights based on advantage values + beta = self.algo_config.adv.beta # temprature factor + weights = torch.exp(adv / beta) + + # clip final weights + if self.algo_config.adv.use_final_clip is True: + weights = weights.clamp(-100.0, 100.0) + + # reshape from (B, 1) to (B,) + return weights[:, 0] + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + log = OrderedDict() + + log["actor/log_prob"] = info["actor/log_prob"].item() + log["actor/loss"] = info["actor/loss"].item() + + log["critic/critic1_pred"] = info["critic/critic1_pred"].item() + log["critic/critic1_loss"] = info["critic/critic1_loss"].item() + + log["vf/v_loss"] = info["vf/v_loss"].item() + + self._log_data_attributes(log, info, "vf/q_pred") + self._log_data_attributes(log, info, "vf/v_pred") + self._log_data_attributes(log, info, "adv/adv") + self._log_data_attributes(log, info, "adv/adv_weight") + + return log + + def _log_data_attributes(self, log, info, key): + """ + Helper function for logging statistics. Moodifies log in-place + + Args: + log (dict): existing log dictionary + log (dict): existing dictionary of tensors containing raw stats + key (str): key to log + """ + log[key + "/max"] = info[key].max().item() + log[key + "/min"] = info[key].min().item() + log[key + "/mean"] = info[key].mean().item() + log[key + "/std"] = info[key].std().item() + + def on_epoch_end(self, epoch): + """ + Called at the end of each epoch. + """ + + # LR scheduling updates + for lr_sc in self.lr_schedulers["critic"]: + if lr_sc is not None: + lr_sc.step() + + if self.lr_schedulers["vf"] is not None: + self.lr_schedulers["vf"].step() + + if self.lr_schedulers["actor"] is not None: + self.lr_schedulers["actor"].step() + + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + action (torch.Tensor): action tensor + """ + assert not self.nets.training + + return self.nets["actor"](obs_dict=obs_dict, goal_dict=goal_dict) \ No newline at end of file diff --git a/aloha-devel/robomimic/algo/iris.py b/aloha-devel/robomimic/algo/iris.py new file mode 100644 index 0000000000000000000000000000000000000000..7b441470c796749f92682ecf2b38a48e0bb3ada5 --- /dev/null +++ b/aloha-devel/robomimic/algo/iris.py @@ -0,0 +1,183 @@ +""" +Implementation of IRIS (https://arxiv.org/abs/1911.05321). +""" +import numpy as np +from collections import OrderedDict +from copy import deepcopy + +import torch + +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.obs_utils as ObsUtils +from robomimic.config.config import Config +from robomimic.algo import register_algo_factory_func, algo_name_to_factory_func, HBC, ValuePlanner, ValueAlgo, GL_VAE + + +@register_algo_factory_func("iris") +def algo_config_to_class(algo_config): + """ + Maps algo config to the IRIS algo class to instantiate, along with additional algo kwargs. + + Args: + algo_config (Config instance): algo config + + Returns: + algo_class: subclass of Algo + algo_kwargs (dict): dictionary of additional kwargs to pass to algorithm + """ + pol_cls, _ = algo_name_to_factory_func("bc")(algo_config.actor) + plan_cls, _ = algo_name_to_factory_func("gl")(algo_config.value_planner.planner) + value_cls, _ = algo_name_to_factory_func("bcq")(algo_config.value_planner.value) + return IRIS, dict(policy_algo_class=pol_cls, planner_algo_class=plan_cls, value_algo_class=value_cls) + + +class IRIS(HBC, ValueAlgo): + """ + Implementation of IRIS (https://arxiv.org/abs/1911.05321). + """ + def __init__( + self, + planner_algo_class, + value_algo_class, + policy_algo_class, + algo_config, + obs_config, + global_config, + obs_key_shapes, + ac_dim, + device, + ): + """ + Args: + planner_algo_class (Algo class): algo class for the planner + + policy_algo_class (Algo class): algo class for the policy + + algo_config (Config object): instance of Config corresponding to the algo section + of the config + + obs_config (Config object): instance of Config corresponding to the observation + section of the config + + global_config (Config object): global training config + + obs_key_shapes (OrderedDict): dictionary that maps input/output observation keys to shapes + + ac_dim (int): action dimension + + device: torch device + """ + self.algo_config = algo_config + self.obs_config = obs_config + self.global_config = global_config + + self.ac_dim = ac_dim + self.device = device + + self._subgoal_step_count = 0 # current step count for deciding when to update subgoal + self._current_subgoal = None # latest subgoal + self._subgoal_update_interval = self.algo_config.subgoal_update_interval # subgoal update frequency + self._subgoal_horizon = self.algo_config.value_planner.planner.subgoal_horizon + self._actor_horizon = self.algo_config.actor.rnn.horizon + + self._algo_mode = self.algo_config.mode + assert self._algo_mode in ["separate", "cascade"] + + self.planner = ValuePlanner( + planner_algo_class=planner_algo_class, + value_algo_class=value_algo_class, + algo_config=algo_config.value_planner, + obs_config=obs_config.value_planner, + global_config=global_config, + obs_key_shapes=obs_key_shapes, + ac_dim=ac_dim, + device=device + ) + + self.actor_goal_shapes = self.planner.subgoal_shapes + assert not algo_config.latent_subgoal.enabled, "IRIS does not support latent subgoals" + + # only for the actor: override goal modalities and shapes to match the subgoal set by the planner + actor_obs_key_shapes = deepcopy(obs_key_shapes) + # make sure we are not modifying existing observation key shapes + for k in self.actor_goal_shapes: + if k in actor_obs_key_shapes: + assert actor_obs_key_shapes[k] == self.actor_goal_shapes[k] + actor_obs_key_shapes.update(self.actor_goal_shapes) + + goal_modalities = {obs_modality: [] for obs_modality in ObsUtils.OBS_MODALITY_CLASSES.keys()} + for k in self.actor_goal_shapes.keys(): + goal_modalities[ObsUtils.OBS_KEYS_TO_MODALITIES[k]].append(k) + + actor_obs_config = deepcopy(obs_config.actor) + with actor_obs_config.unlocked(): + actor_obs_config["goal"] = Config(**goal_modalities) + + self.actor = policy_algo_class( + algo_config=algo_config.actor, + obs_config=actor_obs_config, + global_config=global_config, + obs_key_shapes=actor_obs_key_shapes, + ac_dim=ac_dim, + device=device + ) + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + input_batch = dict() + + input_batch["planner"] = self.planner.process_batch_for_training(batch) + input_batch["actor"] = self.actor.process_batch_for_training(batch) + + if self.algo_config.actor_use_random_subgoals: + # optionally use randomly sampled step between [1, seq_length] as policy goal + policy_subgoal_indices = torch.randint( + low=0, high=self.global_config.train.seq_length, size=(batch["actions"].shape[0],)) + goal_obs = TensorUtils.gather_sequence(batch["next_obs"], policy_subgoal_indices) + goal_obs = TensorUtils.to_float(TensorUtils.to_device(goal_obs, self.device)) + input_batch["actor"]["goal_obs"] = goal_obs + else: + # otherwise, use planner subgoal target as goal for the policy + input_batch["actor"]["goal_obs"] = input_batch["planner"]["planner"]["target_subgoals"] + + # we move to device first before float conversion because image observation modalities will be uint8 - + # this minimizes the amount of data transferred to GPU + return TensorUtils.to_float(TensorUtils.to_device(input_batch, self.device)) + + def get_state_value(self, obs_dict, goal_dict=None): + """ + Get state value outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + value (torch.Tensor): value tensor + """ + return self.planner.get_state_value(obs_dict=obs_dict, goal_dict=goal_dict) + + def get_state_action_value(self, obs_dict, actions, goal_dict=None): + """ + Get state-action value outputs. + + Args: + obs_dict (dict): current observation + actions (torch.Tensor): action + goal_dict (dict): (optional) goal + + Returns: + value (torch.Tensor): value tensor + """ + return self.planner.get_state_action_value(obs_dict=obs_dict, actions=actions, goal_dict=goal_dict) diff --git a/aloha-devel/robomimic/algo/td3_bc.py b/aloha-devel/robomimic/algo/td3_bc.py new file mode 100644 index 0000000000000000000000000000000000000000..e324c54a1614c1c01b3efdb1def9c8f6f11b2c70 --- /dev/null +++ b/aloha-devel/robomimic/algo/td3_bc.py @@ -0,0 +1,567 @@ +""" +Implementation of TD3-BC. +Based on https://github.com/sfujim/TD3_BC +(Paper - https://arxiv.org/abs/1812.02900). + +Note that several parts are exactly the same as the BCQ implementation, +such as @_create_critics, @process_batch_for_training, and +@_train_critic_on_batch. They are replicated here (instead of subclassing +from the BCQ algo class) to be explicit and have implementation details +self-contained in this file. +""" +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import robomimic.models.obs_nets as ObsNets +import robomimic.models.policy_nets as PolicyNets +import robomimic.models.value_nets as ValueNets +import robomimic.models.vae_nets as VAENets +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.torch_utils as TorchUtils +import robomimic.utils.obs_utils as ObsUtils +import robomimic.utils.loss_utils as LossUtils + +from robomimic.algo import register_algo_factory_func, PolicyAlgo, ValueAlgo + + +@register_algo_factory_func("td3_bc") +def algo_config_to_class(algo_config): + """ + Maps algo config to the TD3_BC algo class to instantiate, along with additional algo kwargs. + + Args: + algo_config (Config instance): algo config + + Returns: + algo_class: subclass of Algo + algo_kwargs (dict): dictionary of additional kwargs to pass to algorithm + """ + # only one variant of TD3_BC for now + return TD3_BC, {} + + +class TD3_BC(PolicyAlgo, ValueAlgo): + """ + Default TD3_BC training, based on https://arxiv.org/abs/2106.06860 and + https://github.com/sfujim/TD3_BC. + """ + def __init__(self, **kwargs): + PolicyAlgo.__init__(self, **kwargs) + + # save the discount factor - it may be overriden later + self.set_discount(self.algo_config.discount) + + # initialize actor update counter. This is used to train the actor at a lower freq than critic + self.actor_update_counter = 0 + + def _create_networks(self): + """ + Creates networks and places them into @self.nets. + """ + self.nets = nn.ModuleDict() + + self._create_critics() + self._create_actor() + + # sync target networks at beginning of training + with torch.no_grad(): + for critic_ind in range(len(self.nets["critic"])): + TorchUtils.hard_update( + source=self.nets["critic"][critic_ind], + target=self.nets["critic_target"][critic_ind], + ) + + TorchUtils.hard_update( + source=self.nets["actor"], + target=self.nets["actor_target"], + ) + + self.nets = self.nets.float().to(self.device) + + def _create_critics(self): + """ + Called in @_create_networks to make critic networks. + + Exactly the same as BCQ. + """ + critic_class = ValueNets.ActionValueNetwork + critic_args = dict( + obs_shapes=self.obs_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.critic.layer_dims, + value_bounds=self.algo_config.critic.value_bounds, + goal_shapes=self.goal_shapes, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + + # Q network ensemble and target ensemble + self.nets["critic"] = nn.ModuleList() + self.nets["critic_target"] = nn.ModuleList() + for _ in range(self.algo_config.critic.ensemble.n): + critic = critic_class(**critic_args) + self.nets["critic"].append(critic) + + critic_target = critic_class(**critic_args) + self.nets["critic_target"].append(critic_target) + + def _create_actor(self): + """ + Called in @_create_networks to make actor network. + """ + actor_class = PolicyNets.ActorNetwork + actor_args = dict( + obs_shapes=self.obs_shapes, + goal_shapes=self.goal_shapes, + ac_dim=self.ac_dim, + mlp_layer_dims=self.algo_config.actor.layer_dims, + encoder_kwargs=ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder), + ) + + self.nets["actor"] = actor_class(**actor_args) + self.nets["actor_target"] = actor_class(**actor_args) + + def _check_epoch(self, net_name, epoch): + """ + Helper function to check whether backprop should happen this epoch. + + Args: + net_name (str): name of network in @self.nets and @self.optim_params + epoch (int): epoch number + """ + epoch_start_check = (self.optim_params[net_name]["start_epoch"] == -1) or (epoch >= self.optim_params[net_name]["start_epoch"]) + epoch_end_check = (self.optim_params[net_name]["end_epoch"] == -1) or (epoch < self.optim_params[net_name]["end_epoch"]) + return (epoch_start_check and epoch_end_check) + + def set_discount(self, discount): + """ + Useful function to modify discount factor if necessary (e.g. for n-step returns). + """ + self.discount = discount + + def process_batch_for_training(self, batch): + """ + Processes input batch from a data loader to filter out + relevant information and prepare the batch for training. + + Exactly the same as BCQ. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader + + Returns: + input_batch (dict): processed and filtered batch that + will be used for training + """ + input_batch = dict() + + # n-step returns (default is 1) + n_step = self.algo_config.n_step + assert batch["actions"].shape[1] >= n_step + + # remove temporal batches for all + input_batch["obs"] = {k: batch["obs"][k][:, 0, :] for k in batch["obs"]} + input_batch["next_obs"] = {k: batch["next_obs"][k][:, n_step - 1, :] for k in batch["next_obs"]} + input_batch["goal_obs"] = batch.get("goal_obs", None) # goals may not be present + input_batch["actions"] = batch["actions"][:, 0, :] + + # note: ensure scalar signals (rewards, done) retain last dimension of 1 to be compatible with model outputs + + # single timestep reward is discounted sum of intermediate rewards in sequence + reward_seq = batch["rewards"][:, :n_step] + discounts = torch.pow(self.algo_config.discount, torch.arange(n_step).float()).unsqueeze(0) + input_batch["rewards"] = (reward_seq * discounts).sum(dim=1).unsqueeze(1) + + # discount rate will be gamma^N for computing n-step returns + new_discount = (self.algo_config.discount ** n_step) + self.set_discount(new_discount) + + # consider this n-step seqeunce done if any intermediate dones are present + done_seq = batch["dones"][:, :n_step] + input_batch["dones"] = (done_seq.sum(dim=1) > 0).float().unsqueeze(1) + + if self.algo_config.infinite_horizon: + # scale terminal rewards by 1 / (1 - gamma) for infinite horizon MDPs + done_inds = input_batch["dones"].round().long().nonzero(as_tuple=False)[:, 0] + if done_inds.shape[0] > 0: + input_batch["rewards"][done_inds] = input_batch["rewards"][done_inds] * (1. / (1. - self.discount)) + + # we move to device first before float conversion because image observation modalities will be uint8 - + # this minimizes the amount of data transferred to GPU + return TensorUtils.to_float(TensorUtils.to_device(input_batch, self.device)) + + def _train_critic_on_batch(self, batch, epoch, no_backprop=False): + """ + A modular helper function that can be overridden in case + subclasses would like to modify training behavior for the + critics. + + Exactly the same as BCQ (except for removal of @action_sampler_outputs and @critic_outputs) + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + no_backprop (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + info = OrderedDict() + + # batch variables + s_batch = batch["obs"] + a_batch = batch["actions"] + r_batch = batch["rewards"] + ns_batch = batch["next_obs"] + goal_s_batch = batch["goal_obs"] + + # 1 if not done, 0 otherwise + done_mask_batch = 1. - batch["dones"] + info["done_masks"] = done_mask_batch + + # Bellman backup for Q-targets + q_targets = self._get_target_values( + next_states=ns_batch, + goal_states=goal_s_batch, + rewards=r_batch, + dones=done_mask_batch, + ) + info["critic/q_targets"] = q_targets + + # Train all critics using this set of targets for regression + for critic_ind, critic in enumerate(self.nets["critic"]): + critic_loss = self._compute_critic_loss( + critic=critic, + states=s_batch, + actions=a_batch, + goal_states=goal_s_batch, + q_targets=q_targets, + ) + info["critic/critic{}_loss".format(critic_ind + 1)] = critic_loss + + if not no_backprop: + critic_grad_norms = TorchUtils.backprop_for_loss( + net=self.nets["critic"][critic_ind], + optim=self.optimizers["critic"][critic_ind], + loss=critic_loss, + max_grad_norm=self.algo_config.critic.max_gradient_norm, + ) + info["critic/critic{}_grad_norms".format(critic_ind + 1)] = critic_grad_norms + + return info + + def _train_actor_on_batch(self, batch, epoch, no_backprop=False): + """ + A modular helper function that can be overridden in case + subclasses would like to modify training behavior for the + actor. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + no_backprop (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + info = OrderedDict() + + # Actor loss (update with mixture of DDPG loss and BC loss) + s_batch = batch["obs"] + a_batch = batch["actions"] + goal_s_batch = batch["goal_obs"] + + # lambda mixture weight is combination of hyperparameter (alpha) and Q-value normalization + actor_actions = self.nets["actor"](s_batch, goal_s_batch) + Q_values = self.nets["critic"][0](s_batch, actor_actions, goal_s_batch) + lam = self.algo_config.alpha / Q_values.abs().mean().detach() + actor_loss = -lam * Q_values.mean() + nn.MSELoss()(actor_actions, a_batch) + info["actor/loss"] = actor_loss + + if not no_backprop: + actor_grad_norms = TorchUtils.backprop_for_loss( + net=self.nets["actor"], + optim=self.optimizers["actor"], + loss=actor_loss, + ) + info["actor/grad_norms"] = actor_grad_norms + + return info + + def _get_target_values(self, next_states, goal_states, rewards, dones): + """ + Helper function to get target values for training Q-function with TD-loss. + + Args: + next_states (dict): batch of next observations + goal_states (dict): if not None, batch of goal observations + rewards (torch.Tensor): batch of rewards - should be shape (B, 1) + dones (torch.Tensor): batch of done signals - should be shape (B, 1) + + Returns: + q_targets (torch.Tensor): target Q-values to use for TD loss + """ + + with torch.no_grad(): + # get next actions via target actor and noise + next_target_actions = self.nets["actor_target"](next_states, goal_states) + noise = ( + torch.randn_like(next_target_actions) * self.algo_config.actor.noise_std + ).clamp(-self.algo_config.actor.noise_clip, self.algo_config.actor.noise_clip) + next_actions = (next_target_actions + noise).clamp(-1.0, 1.0) + + # TD3 trick to combine max and min over all Q-ensemble estimates into single target estimates + all_value_targets = self.nets["critic_target"][0](next_states, next_actions, goal_states).reshape(-1, 1) + max_value_targets = all_value_targets + min_value_targets = all_value_targets + for critic_target in self.nets["critic_target"][1:]: + all_value_targets = critic_target(next_states, next_actions, goal_states).reshape(-1, 1) + max_value_targets = torch.max(max_value_targets, all_value_targets) + min_value_targets = torch.min(min_value_targets, all_value_targets) + value_targets = self.algo_config.critic.ensemble.weight * min_value_targets + \ + (1. - self.algo_config.critic.ensemble.weight) * max_value_targets + q_targets = rewards + dones * self.discount * value_targets + + return q_targets + + def _compute_critic_loss(self, critic, states, actions, goal_states, q_targets): + """ + Helper function to compute loss between estimated Q-values and target Q-values. + + Nearly the same as BCQ (return type slightly different). + + Args: + critic (torch.nn.Module): critic network + states (dict): batch of observations + actions (torch.Tensor): batch of actions + goal_states (dict): if not None, batch of goal observations + q_targets (torch.Tensor): batch of target q-values for the TD loss + + Returns: + critic_loss (torch.Tensor): critic loss + """ + q_estimated = critic(states, actions, goal_states) + if self.algo_config.critic.use_huber: + critic_loss = nn.SmoothL1Loss()(q_estimated, q_targets) + else: + critic_loss = nn.MSELoss()(q_estimated, q_targets) + return critic_loss + + def train_on_batch(self, batch, epoch, validate=False): + """ + Training on a single batch of data. + + Args: + batch (dict): dictionary with torch.Tensors sampled + from a data loader and filtered by @process_batch_for_training + + epoch (int): epoch number - required by some Algos that need + to perform staged training and early stopping + + validate (bool): if True, don't perform any learning updates. + + Returns: + info (dict): dictionary of relevant inputs, outputs, and losses + that might be relevant for logging + """ + with TorchUtils.maybe_no_grad(no_grad=validate): + info = PolicyAlgo.train_on_batch(self, batch, epoch, validate=validate) + + # Critic training + no_critic_backprop = validate or (not self._check_epoch(net_name="critic", epoch=epoch)) + with TorchUtils.maybe_no_grad(no_grad=no_critic_backprop): + critic_info = self._train_critic_on_batch( + batch=batch, + epoch=epoch, + no_backprop=no_critic_backprop, + ) + info.update(critic_info) + + # update actor and target networks at lower frequency + if not no_critic_backprop: + # update counter only on critic training gradient steps + self.actor_update_counter += 1 + do_actor_update = (self.actor_update_counter % self.algo_config.actor.update_freq == 0) + + # Actor training + no_actor_backprop = validate or (not self._check_epoch(net_name="actor", epoch=epoch)) + no_actor_backprop = no_actor_backprop or (not do_actor_update) + with TorchUtils.maybe_no_grad(no_grad=no_actor_backprop): + actor_info = self._train_actor_on_batch( + batch=batch, + epoch=epoch, + no_backprop=no_actor_backprop, + ) + info.update(actor_info) + + if not no_actor_backprop: + # to match original implementation, only update target networks on + # actor gradient steps + with torch.no_grad(): + # update the target critic networks + for critic_ind in range(len(self.nets["critic"])): + TorchUtils.soft_update( + source=self.nets["critic"][critic_ind], + target=self.nets["critic_target"][critic_ind], + tau=self.algo_config.target_tau, + ) + + # update target actor network + TorchUtils.soft_update( + source=self.nets["actor"], + target=self.nets["actor_target"], + tau=self.algo_config.target_tau, + ) + + return info + + def log_info(self, info): + """ + Process info dictionary from @train_on_batch to summarize + information to pass to tensorboard for logging. + + Args: + info (dict): dictionary of info + + Returns: + loss_log (dict): name -> summary statistic + """ + loss_log = OrderedDict() + + # record current optimizer learning rates + for k in self.optimizers: + keys = [k] + optims = [self.optimizers[k]] + if k == "critic": + # account for critic having one optimizer per ensemble member + keys = ["{}{}".format(k, critic_ind) for critic_ind in range(len(self.nets["critic"]))] + optims = self.optimizers[k] + for kp, optimizer in zip(keys, optims): + for i, param_group in enumerate(optimizer.param_groups): + loss_log["Optimizer/{}{}_lr".format(kp, i)] = param_group["lr"] + + # extract relevant logs for critic, and actor + loss_log["Loss"] = 0. + for loss_logger in [self._log_critic_info, self._log_actor_info]: + this_log = loss_logger(info) + if "Loss" in this_log: + # manually merge total loss + loss_log["Loss"] += this_log["Loss"] + del this_log["Loss"] + loss_log.update(this_log) + + return loss_log + + def _log_critic_info(self, info): + """ + Helper function to extract critic-relevant information for logging. + """ + loss_log = OrderedDict() + if "done_masks" in info: + loss_log["Critic/Done_Mask_Percentage"] = 100. * torch.mean(info["done_masks"]).item() + if "critic/q_targets" in info: + loss_log["Critic/Q_Targets"] = info["critic/q_targets"].mean().item() + loss_log["Loss"] = 0. + for critic_ind in range(len(self.nets["critic"])): + loss_log["Critic/Critic{}_Loss".format(critic_ind + 1)] = info["critic/critic{}_loss".format(critic_ind + 1)].item() + if "critic/critic{}_grad_norms".format(critic_ind + 1) in info: + loss_log["Critic/Critic{}_Grad_Norms".format(critic_ind + 1)] = info["critic/critic{}_grad_norms".format(critic_ind + 1)] + loss_log["Loss"] += loss_log["Critic/Critic{}_Loss".format(critic_ind + 1)] + return loss_log + + def _log_actor_info(self, info): + """ + Helper function to extract actor-relevant information for logging. + """ + loss_log = OrderedDict() + loss_log["Actor/Loss"] = info["actor/loss"].item() + if "actor/grad_norms" in info: + loss_log["Actor/Grad_Norms"] = info["actor/grad_norms"] + loss_log["Loss"] = loss_log["Actor/Loss"] + return loss_log + + def set_train(self): + """ + Prepare networks for evaluation. Update from super class to make sure + target networks stay in evaluation mode all the time. + """ + self.nets.train() + + # target networks always in eval + for critic_ind in range(len(self.nets["critic_target"])): + self.nets["critic_target"][critic_ind].eval() + + self.nets["actor_target"].eval() + + def on_epoch_end(self, epoch): + """ + Called at the end of each epoch. + """ + + # LR scheduling updates + for lr_sc in self.lr_schedulers["critic"]: + if lr_sc is not None: + lr_sc.step() + + if self.lr_schedulers["actor"] is not None: + self.lr_schedulers["actor"].step() + + def get_action(self, obs_dict, goal_dict=None): + """ + Get policy action outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + action (torch.Tensor): action tensor + """ + assert not self.nets.training + + return self.nets["actor"](obs_dict=obs_dict, goal_dict=goal_dict) + + def get_state_value(self, obs_dict, goal_dict=None): + """ + Get state value outputs. + + Args: + obs_dict (dict): current observation + goal_dict (dict): (optional) goal + + Returns: + value (torch.Tensor): value tensor + """ + assert not self.nets.training + + actions = self.nets["actor"](obs_dict=obs_dict, goal_dict=goal_dict) + return self.nets["critic"][0](obs_dict, actions, goal_dict) + + def get_state_action_value(self, obs_dict, actions, goal_dict=None): + """ + Get state-action value outputs. + + Args: + obs_dict (dict): current observation + actions (torch.Tensor): action + goal_dict (dict): (optional) goal + + Returns: + value (torch.Tensor): value tensor + """ + assert not self.nets.training + + return self.nets["critic"][0](obs_dict, actions, goal_dict) diff --git a/aloha-devel/robomimic/config/__pycache__/__init__.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/__init__.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..53cb05c83927c8e229f4a034262881af492c5547 Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/__init__.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__pycache__/bc_config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/bc_config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2aea8e2d92f935d1e56e7a2158d2d7d072529cab Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/bc_config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__pycache__/bcq_config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/bcq_config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6afaef4becd1f98a0e5e8d4f0097d75085ca0dcf Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/bcq_config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__pycache__/diffusion_policy_config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/diffusion_policy_config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ab68b95378ee508ee098eaa024b9430a66e483e Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/diffusion_policy_config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__pycache__/gl_config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/gl_config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b393e00048184f69690a0766c61b07bd749b8ff2 Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/gl_config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__pycache__/hbc_config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/hbc_config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd2e541af7821ee596f8d9f08bc3890cfd1349d4 Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/hbc_config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__pycache__/iql_config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/iql_config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8aba43a24de51efc573ff052cae9a201043361ed Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/iql_config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__pycache__/iris_config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/iris_config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d5ea741215ff6ce78d06b5004de014bcf077b22c Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/iris_config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/act_config.py b/aloha-devel/robomimic/config/act_config.py new file mode 100644 index 0000000000000000000000000000000000000000..9be3926b46e5d15a517db6c736d0c1950e57cc43 --- /dev/null +++ b/aloha-devel/robomimic/config/act_config.py @@ -0,0 +1,47 @@ +""" +Config for BC algorithm. +""" + +from robomimic.config.base_config import BaseConfig + + +class ACTConfig(BaseConfig): + ALGO_NAME = "act" + + def train_config(self): + """ + BC algorithms don't need "next_obs" from hdf5 - so save on storage and compute by disabling it. + """ + super(ACTConfig, self).train_config() + self.train.hdf5_load_next_obs = False + + def algo_config(self): + """ + This function populates the `config.algo` attribute of the config, and is given to the + `Algo` subclass (see `algo/algo.py`) for each algorithm through the `algo_config` + argument to the constructor. Any parameter that an algorithm needs to determine its + training and test-time behavior should be populated here. + """ + + # optimization parameters + self.algo.optim_params.policy.optimizer_type = "adamw" + self.algo.optim_params.policy.learning_rate.initial = 5e-5 # policy learning rate + self.algo.optim_params.policy.learning_rate.decay_factor = 1 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.policy.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.policy.learning_rate.scheduler_type = "linear" # learning rate scheduler ("multistep", "linear", etc) + self.algo.optim_params.policy.regularization.L2 = 0.0001 # L2 regularization strength + + # loss weights + self.algo.loss.l2_weight = 0.0 # L2 loss weight + self.algo.loss.l1_weight = 1.0 # L1 loss weight + self.algo.loss.cos_weight = 0.0 # cosine loss weight + + # ACT policy settings + self.algo.act.hidden_dim = 512 # length of (s, a) seqeunces to feed to transformer - should usually match train.frame_stack + self.algo.act.dim_feedforward = 3200 # dimension for embeddings used by transformer + self.algo.act.backbone = "resnet18" # number of transformer blocks to stack + self.algo.act.enc_layers = 4 # number of attention heads for each transformer block (should divide embed_dim evenly) + self.algo.act.dec_layers = 7 # dropout probability for embedding inputs in transformer + self.algo.act.nheads = 8 # dropout probability for attention outputs for each transformer block + self.algo.act.latent_dim = 32 # latent dim of VAE + self.algo.act.kl_weight = 20 # KL weight of VAE diff --git a/aloha-devel/robomimic/config/base_config.py b/aloha-devel/robomimic/config/base_config.py new file mode 100644 index 0000000000000000000000000000000000000000..e3a18d6d67adfa903a926a90b1678c310df55fc2 --- /dev/null +++ b/aloha-devel/robomimic/config/base_config.py @@ -0,0 +1,354 @@ +""" +The base config class that is used for all algorithm configs in this repository. +Subclasses get registered into a global dictionary, making it easy to instantiate +the correct config class given the algorithm name. +""" + +import six # preserve metaclass compatibility between python 2 and 3 +from copy import deepcopy + +import robomimic +from robomimic.config.config import Config + +# global dictionary for remembering name - class mappings +REGISTERED_CONFIGS = {} + + +def get_all_registered_configs(): + """ + Give access to dictionary of all registered configs for external use. + """ + return deepcopy(REGISTERED_CONFIGS) + + +def config_factory(algo_name, dic=None): + """ + Creates an instance of a config from the algo name. Optionally pass + a dictionary to instantiate the config from the dictionary. + """ + if algo_name not in REGISTERED_CONFIGS: + raise Exception("Config for algo name {} not found. Make sure it is a registered config among: {}".format( + algo_name, ', '.join(REGISTERED_CONFIGS))) + return REGISTERED_CONFIGS[algo_name](dict_to_load=dic) + + +class ConfigMeta(type): + """ + Define a metaclass for constructing a config class. + It registers configs into the global registry. + """ + def __new__(meta, name, bases, class_dict): + cls = super(ConfigMeta, meta).__new__(meta, name, bases, class_dict) + if cls.__name__ != "BaseConfig": + REGISTERED_CONFIGS[cls.ALGO_NAME] = cls + return cls + + +@six.add_metaclass(ConfigMeta) +class BaseConfig(Config): + def __init__(self, dict_to_load=None): + if dict_to_load is not None: + super(BaseConfig, self).__init__(dict_to_load) + return + + super(BaseConfig, self).__init__() + + # store algo name class property in the config (must be implemented by subclasses) + self.algo_name = type(self).ALGO_NAME + + self.experiment_config() + self.train_config() + self.algo_config() + self.observation_config() + self.meta_config() + + # After Config init, new keys cannot be added to the config, except under nested + # attributes that have called @do_not_lock_keys + self.lock_keys() + + @property + @classmethod + def ALGO_NAME(cls): + # must be specified by subclasses + raise NotImplementedError + + def experiment_config(self): + """ + This function populates the `config.experiment` attribute of the config, + which has several experiment settings such as the name of the training run, + whether to do logging, whether to save models (and how often), whether to render + videos, and whether to do rollouts (and how often). This class has a default + implementation that usually doesn't need to be overriden. + """ + + self.experiment.name = "test" # name of experiment used to make log files + self.experiment.validate = False # whether to do validation or not + self.experiment.logging.terminal_output_to_txt = True # whether to log stdout to txt file + self.experiment.logging.log_tb = True # enable tensorboard logging + self.experiment.logging.log_wandb = False # enable wandb logging + self.experiment.logging.wandb_proj_name = "debug" # project name if using wandb + + # log model prediction MSE + self.experiment.mse.enabled = False # whether to log model prediction MSE + self.experiment.mse.every_n_epochs = 50 # log model prediction MSE every n epochs + self.experiment.mse.on_save_ckpt = True # log model prediction MSE on model checkpoint + self.experiment.mse.num_samples = 20 # number of datapoints to use for MSE prediction + self.experiment.mse.visualize = True # save model prediction visualizations + + ## save config - if and when to save model checkpoints ## + self.experiment.save.enabled = True # whether model saving should be enabled or disabled + self.experiment.save.every_n_seconds = None # save model every n seconds (set to None to disable) + self.experiment.save.every_n_epochs = 50 # save model every n epochs (set to None to disable) + self.experiment.save.epochs = [] # save model on these specific epochs + self.experiment.save.on_best_validation = False # save models that achieve best validation score + self.experiment.save.on_best_rollout_return = False # save models that achieve best rollout return + self.experiment.save.on_best_rollout_success_rate = True # save models that achieve best success rate + + # epoch definitions - if not None, set an epoch to be this many gradient steps, else the full dataset size will be used + self.experiment.epoch_every_n_steps = 100 # number of gradient steps in train epoch (None for full dataset pass) + self.experiment.validation_epoch_every_n_steps = 10 # number of gradient steps in valid epoch (None for full dataset pass) + + # envs to evaluate model on (assuming rollouts are enabled), to override the metadata stored in dataset + self.experiment.env = None # no need to set this (unless you want to override) + self.experiment.additional_envs = None # additional environments that should get evaluated + + + ## rendering config ## + self.experiment.render = False # render on-screen or not + self.experiment.render_video = True # render evaluation rollouts to videos + self.experiment.keep_all_videos = False # save all videos, instead of only saving those for saved model checkpoints + self.experiment.video_skip = 5 # render video frame every n environment steps during rollout + + + ## evaluation rollout config ## + self.experiment.rollout.enabled = True # enable evaluation rollouts + self.experiment.rollout.n = 50 # number of rollouts per evaluation + self.experiment.rollout.horizon = 400 # maximum number of env steps per rollout + self.experiment.rollout.rate = 50 # do rollouts every @rate epochs + self.experiment.rollout.warmstart = 0 # number of epochs to wait before starting rollouts + self.experiment.rollout.terminate_on_success = True # end rollout early after task success + self.experiment.rollout.batched = False # whether to parallelize evaluations over batched environments + self.experiment.rollout.num_batch_envs = 5 # number of batched environments to use (applicable if experiment.rollout.batched is True) + + # for updating the evaluation env meta data + self.experiment.env_meta_update_dict = Config() + self.experiment.env_meta_update_dict.do_not_lock_keys() + + # whether to load in a previously trained model checkpoint + self.experiment.ckpt_path = None + + def train_config(self): + """ + This function populates the `config.train` attribute of the config, which + has several settings related to the training process, such as the dataset + to use for training, and how the data loader should load the data. This + class has a default implementation that usually doesn't need to be overriden. + """ + + # Path to hdf5 dataset to use for training + self.train.data = None + + # Write all results to this directory. A new folder with the timestamp will be created + # in this directory, and it will contain three subfolders - "log", "models", and "videos". + # The "log" directory will contain tensorboard and stdout txt logs. The "models" directory + # will contain saved model checkpoints. The "videos" directory contains evaluation rollout + # videos. + self.train.output_dir = "../{}_trained_models".format(self.algo_name) + + + ## dataset loader config ## + + # num workers for loading data - generally set to 0 for low-dim datasets, and 2 for image datasets + self.train.num_data_workers = 0 + + # One of ["all", "low_dim", or None]. Set to "all" to cache entire hdf5 in memory - this is + # by far the fastest for data loading. Set to "low_dim" to cache all non-image data. Set + # to None to use no caching - in this case, every batch sample is retrieved via file i/o. + # You should almost never set this to None, even for large image datasets. + self.train.hdf5_cache_mode = "all" + + # used for parallel data loading + self.train.hdf5_use_swmr = True + + # whether to load "next_obs" group from hdf5 - only needed for batch / offline RL algorithms + self.train.hdf5_load_next_obs = True + + # if true, normalize observations at train and test time, using the global mean and standard deviation + # of each observation in each dimension, computed across the training set. See SequenceDataset.normalize_obs + # in utils/dataset.py for more information. + self.train.hdf5_normalize_obs = False + + # if provided, use the list of demo keys under the hdf5 group "mask/@hdf5_filter_key" for training, instead + # of the full dataset. This provides a convenient way to train on only a subset of the trajectories in a dataset. + self.train.hdf5_filter_key = None + + # if provided, use the list of demo keys under the hdf5 group "mask/@hdf5_validation_filter_key" for validation. + # Must be provided if @experiment.validate is True. + self.train.hdf5_validation_filter_key = None + + # length of experience sequence to fetch from the dataset + # and whether to pad the beginning / end of the sequence at boundaries of trajectory in dataset + self.train.seq_length = 1 + self.train.pad_seq_length = True + self.train.frame_stack = 1 + self.train.pad_frame_stack = True + + # keys from hdf5 to load into each batch, besides "obs" and "next_obs". If algorithms + # require additional keys from each trajectory in the hdf5, they should be specified here. + self.train.dataset_keys = ( + "actions", + "rewards", + "dones", + ) + + self.train.action_keys = ["actions"] + + # specifing each action keys to load and their corresponding normalization/conversion requirement + # e.g. for dataset keys "action/eef_pos" and "action/eef_rot" + # the desired value of self.train.action_config is: + # { + # "action/eef_pos": { + # "normalization": "min_max", + # "rot_conversion: None + # }, + # "action/eef_rot": { + # "normalization": None, + # "rot_conversion: "axis_angle_to_6d" + # } + # } + # self.train.action_config.actions.normalization = None # "min_max" + # self.train.action_config.actions.rot_conversion = None # "axis_angle_to_6d" + self.train.action_config = {} + # self.train.action_config.do_not_lock_keys() + + # one of [None, "last"] - set to "last" to include goal observations in each batch + self.train.goal_mode = None + + + ## learning config ## + self.train.cuda = True # use GPU or not + self.train.batch_size = 100 # batch size + self.train.num_epochs = 2000 # number of training epochs + self.train.seed = 1 # seed for training (for reproducibility) + + self.train.max_grad_norm = None # clip gradient norms (see `backprop_for_loss` function in torch_utils.py) + + self.train.data_format = "robomimic" # either "robomimic" or "r2d2" + + # list of observation keys to shuffle randomly in the dataset. + # must be list of tuples pairs, with each pair representing + # the corresponding observation key groups to shuffle + self.train.shuffled_obs_key_groups = None + + def algo_config(self): + """ + This function populates the `config.algo` attribute of the config, and is given to the + `Algo` subclass (see `algo/algo.py`) for each algorithm through the `algo_config` + argument to the constructor. Any parameter that an algorithm needs to determine its + training and test-time behavior should be populated here. This function should be + implemented by every subclass. + """ + pass + + def observation_config(self): + """ + This function populates the `config.observation` attribute of the config, and is given + to the `Algo` subclass (see `algo/algo.py`) for each algorithm through the `obs_config` + argument to the constructor. This portion of the config is used to specify what + observation modalities should be used by the networks for training, and how the + observation modalities should be encoded by the networks. While this class has a + default implementation that usually doesn't need to be overriden, certain algorithm + configs may choose to, in order to have seperate configs for different networks + in the algorithm. + """ + + # observation modalities + self.observation.modalities.obs.low_dim = [ # specify low-dim observations for agent + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object", + ] + self.observation.modalities.obs.rgb = [] # specify rgb image observations for agent + self.observation.modalities.obs.depth = [] + self.observation.modalities.obs.scan = [] + self.observation.modalities.goal.low_dim = [] # specify low-dim goal observations to condition agent on + self.observation.modalities.goal.rgb = [] # specify rgb image goal observations to condition agent on + self.observation.modalities.goal.depth = [] + self.observation.modalities.goal.scan = [] + self.observation.modalities.obs.do_not_lock_keys() + self.observation.modalities.goal.do_not_lock_keys() + + # observation encoder architectures (per obs modality) + # This applies to all networks that take observation dicts as input + + # =============== Low Dim default encoder (no encoder) =============== + self.observation.encoder.low_dim.core_class = None + self.observation.encoder.low_dim.core_kwargs = Config() # No kwargs by default + self.observation.encoder.low_dim.core_kwargs.do_not_lock_keys() + + # Low Dim: Obs Randomizer settings + self.observation.encoder.low_dim.obs_randomizer_class = None + self.observation.encoder.low_dim.obs_randomizer_kwargs = Config() # No kwargs by default + self.observation.encoder.low_dim.obs_randomizer_kwargs.do_not_lock_keys() + + # =============== RGB default encoder (ResNet backbone + linear layer output) =============== + self.observation.encoder.rgb.core_class = "VisualCore" # Default VisualCore class combines backbone (like ResNet-18) with pooling operation (like spatial softmax) + self.observation.encoder.rgb.core_kwargs = Config() # See models/obs_core.py for important kwargs to set and defaults used + self.observation.encoder.rgb.core_kwargs.do_not_lock_keys() + + # RGB: Obs Randomizer settings + self.observation.encoder.rgb.obs_randomizer_class = None # Can set to 'CropRandomizer' to use crop randomization + self.observation.encoder.rgb.obs_randomizer_kwargs = Config() # See models/obs_core.py for important kwargs to set and defaults used + self.observation.encoder.rgb.obs_randomizer_kwargs.do_not_lock_keys() + + # Allow for other custom modalities to be specified + self.observation.encoder.do_not_lock_keys() + + # =============== Depth default encoder (same as rgb) =============== + self.observation.encoder.depth = deepcopy(self.observation.encoder.rgb) + + # =============== Scan default encoder (Conv1d backbone + linear layer output) =============== + self.observation.encoder.scan = deepcopy(self.observation.encoder.rgb) + + # Scan: Modify the core class + kwargs, otherwise, is same as rgb encoder + self.observation.encoder.scan.core_class = "ScanCore" # Default ScanCore class uses Conv1D to process this modality + self.observation.encoder.scan.core_kwargs = Config() # See models/obs_core.py for important kwargs to set and defaults used + self.observation.encoder.scan.core_kwargs.do_not_lock_keys() + + def meta_config(self): + """ + This function populates the `config.meta` attribute of the config. This portion of the config + is used to specify job information primarily for hyperparameter sweeps. + It contains hyperparameter keys and values, which are populated automatically + by the hyperparameter config generator (see `utils/hyperparam_utils.py`). + These values are read by the wandb logger (see `utils/log_utils.py`) to set job tags. + """ + + self.meta.hp_base_config_file = None # base config file in hyperparam sweep + self.meta.hp_keys = [] # relevant keys (swept) in hyperparam sweep + self.meta.hp_values = [] # values corresponding to keys in hyperparam sweep + + @property + def use_goals(self): + # whether the agent is goal-conditioned + return len([obs_key for modality in self.observation.modalities.goal.values() for obs_key in modality]) > 0 + + @property + def all_obs_keys(self): + """ + This grabs the union of observation keys over all modalities (e.g.: low_dim, rgb, depth, etc.) and over all + modality groups (e.g: obs, goal, subgoal, etc...) + + Returns: + n-array: all observation keys used for this model + """ + # pool all modalities + return sorted(tuple(set([ + obs_key for group in [ + self.observation.modalities.obs.values(), + self.observation.modalities.goal.values() + ] + for modality in group + for obs_key in modality + ]))) diff --git a/aloha-devel/robomimic/config/bc_config.py b/aloha-devel/robomimic/config/bc_config.py new file mode 100644 index 0000000000000000000000000000000000000000..a010b2fad2bbe04ac93837a5560f74e1f97e9cb2 --- /dev/null +++ b/aloha-devel/robomimic/config/bc_config.py @@ -0,0 +1,110 @@ +""" +Config for BC algorithm. +""" + +from robomimic.config.base_config import BaseConfig + + +class BCConfig(BaseConfig): + ALGO_NAME = "bc" + + def train_config(self): + """ + BC algorithms don't need "next_obs" from hdf5 - so save on storage and compute by disabling it. + """ + super(BCConfig, self).train_config() + self.train.hdf5_load_next_obs = False + + def algo_config(self): + """ + This function populates the `config.algo` attribute of the config, and is given to the + `Algo` subclass (see `algo/algo.py`) for each algorithm through the `algo_config` + argument to the constructor. Any parameter that an algorithm needs to determine its + training and test-time behavior should be populated here. + """ + + # optimization parameters + self.algo.optim_params.policy.optimizer_type = "adam" + self.algo.optim_params.policy.learning_rate.initial = 1e-4 # policy learning rate + self.algo.optim_params.policy.learning_rate.decay_factor = 0.1 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.policy.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.policy.learning_rate.scheduler_type = "multistep" # learning rate scheduler ("multistep", "linear", etc) + self.algo.optim_params.policy.regularization.L2 = 0.00 # L2 regularization strength + + # loss weights + self.algo.loss.l2_weight = 1.0 # L2 loss weight + self.algo.loss.l1_weight = 0.0 # L1 loss weight + self.algo.loss.cos_weight = 0.0 # cosine loss weight + + # MLP network architecture (layers after observation encoder and RNN, if present) + self.algo.actor_layer_dims = (1024, 1024) + + # stochastic Gaussian policy settings + self.algo.gaussian.enabled = False # whether to train a Gaussian policy + self.algo.gaussian.fixed_std = False # whether to train std output or keep it constant + self.algo.gaussian.init_std = 0.1 # initial standard deviation (or constant) + self.algo.gaussian.min_std = 0.01 # minimum std output from network + self.algo.gaussian.std_activation = "softplus" # activation to use for std output from policy net + self.algo.gaussian.low_noise_eval = True # low-std at test-time + + # stochastic GMM policy settings + self.algo.gmm.enabled = False # whether to train a GMM policy + self.algo.gmm.num_modes = 5 # number of GMM modes + self.algo.gmm.min_std = 0.0001 # minimum std output from network + self.algo.gmm.std_activation = "softplus" # activation to use for std output from policy net + self.algo.gmm.low_noise_eval = True # low-std at test-time + + # stochastic VAE policy settings + self.algo.vae.enabled = False # whether to train a VAE policy + self.algo.vae.latent_dim = 14 # VAE latent dimnsion - set to twice the dimensionality of action space + self.algo.vae.latent_clip = None # clip latent space when decoding (set to None to disable) + self.algo.vae.kl_weight = 1. # beta-VAE weight to scale KL loss relative to reconstruction loss in ELBO + + # VAE decoder settings + self.algo.vae.decoder.is_conditioned = True # whether decoder should condition on observation + self.algo.vae.decoder.reconstruction_sum_across_elements = False # sum instead of mean for reconstruction loss + + # VAE prior settings + self.algo.vae.prior.learn = False # learn Gaussian / GMM prior instead of N(0, 1) + self.algo.vae.prior.is_conditioned = False # whether to condition prior on observations + self.algo.vae.prior.use_gmm = False # whether to use GMM prior + self.algo.vae.prior.gmm_num_modes = 10 # number of GMM modes + self.algo.vae.prior.gmm_learn_weights = False # whether to learn GMM weights + self.algo.vae.prior.use_categorical = False # whether to use categorical prior + self.algo.vae.prior.categorical_dim = 10 # the number of categorical classes for each latent dimension + self.algo.vae.prior.categorical_gumbel_softmax_hard = False # use hard selection in forward pass + self.algo.vae.prior.categorical_init_temp = 1.0 # initial gumbel-softmax temp + self.algo.vae.prior.categorical_temp_anneal_step = 0.001 # linear temp annealing rate + self.algo.vae.prior.categorical_min_temp = 0.3 # lowest gumbel-softmax temp + + self.algo.vae.encoder_layer_dims = (300, 400) # encoder MLP layer dimensions + self.algo.vae.decoder_layer_dims = (300, 400) # decoder MLP layer dimensions + self.algo.vae.prior_layer_dims = (300, 400) # prior MLP layer dimensions (if learning conditioned prior) + + # RNN policy settings + self.algo.rnn.enabled = False # whether to train RNN policy + self.algo.rnn.horizon = 10 # unroll length for RNN - should usually match train.seq_length + self.algo.rnn.hidden_dim = 400 # hidden dimension size + self.algo.rnn.rnn_type = "LSTM" # rnn type - one of "LSTM" or "GRU" + self.algo.rnn.num_layers = 2 # number of RNN layers that are stacked + self.algo.rnn.open_loop = False # if True, action predictions are only based on a single observation (not sequence) + self.algo.rnn.kwargs.bidirectional = False # rnn kwargs + self.algo.rnn.kwargs.do_not_lock_keys() + + # Transformer policy settings + self.algo.transformer.enabled = False # whether to train transformer policy + self.algo.transformer.context_length = 10 # length of (s, a) seqeunces to feed to transformer - should usually match train.frame_stack + self.algo.transformer.embed_dim = 512 # dimension for embeddings used by transformer + self.algo.transformer.num_layers = 6 # number of transformer blocks to stack + self.algo.transformer.num_heads = 8 # number of attention heads for each transformer block (should divide embed_dim evenly) + self.algo.transformer.emb_dropout = 0.1 # dropout probability for embedding inputs in transformer + self.algo.transformer.attn_dropout = 0.1 # dropout probability for attention outputs for each transformer block + self.algo.transformer.block_output_dropout = 0.1 # dropout probability for final outputs for each transformer block + self.algo.transformer.sinusoidal_embedding = False # if True, use standard positional encodings (sin/cos) + self.algo.transformer.activation = "gelu" # activation function for MLP in Transformer Block + self.algo.transformer.supervise_all_steps = False # if true, supervise all intermediate actions, otherwise only final one + self.algo.transformer.nn_parameter_for_timesteps = True # if true, use nn.Parameter otherwise use nn.Embedding + self.algo.transformer.pred_future_acs = False # shift action prediction forward to predict future actions instead of past actions + self.algo.transformer.causal = True # whether the transformer is causal + + self.algo.language_conditioned = False # whether policy is language conditioned diff --git a/aloha-devel/robomimic/config/cql_config.py b/aloha-devel/robomimic/config/cql_config.py new file mode 100644 index 0000000000000000000000000000000000000000..26fea048fe49d2d2d03f888eedab6754f37c8dcc --- /dev/null +++ b/aloha-devel/robomimic/config/cql_config.py @@ -0,0 +1,82 @@ +""" +Config for CQL algorithm. +""" + +from robomimic.config.base_config import BaseConfig + + +class CQLConfig(BaseConfig): + ALGO_NAME = "cql" + + def train_config(self): + """ + Update from superclass to change default batch size. + """ + super(CQLConfig, self).train_config() + + # increase batch size to 1024 (found to work better for most manipulation experiments) + self.train.batch_size = 1024 + + def algo_config(self): + """ + This function populates the `config.algo` attribute of the config, and is given to the + `Algo` subclass (see `algo/algo.py`) for each algorithm through the `algo_config` + argument to the constructor. Any parameter that an algorithm needs to determine its + training and test-time behavior should be populated here. + """ + + # optimization parameters + self.algo.optim_params.critic.learning_rate.initial = 1e-3 # critic learning rate + self.algo.optim_params.critic.learning_rate.decay_factor = 0.0 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.critic.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.critic.regularization.L2 = 0.00 # L2 regularization strength + + self.algo.optim_params.actor.learning_rate.initial = 3e-4 # actor learning rate + self.algo.optim_params.actor.learning_rate.decay_factor = 0.0 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.actor.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.actor.regularization.L2 = 0.00 # L2 regularization strength + + # target network related parameters + self.algo.discount = 0.99 # discount factor to use + self.algo.n_step = 1 # for using n-step returns in TD-updates + self.algo.target_tau = 0.005 # update rate for target networks + + # ================== Actor Network Config =================== + self.algo.actor.bc_start_steps = 0 # uses BC policy loss for first n-training steps + self.algo.actor.target_entropy = "default" # None is fixed entropy, otherwise is automatically tuned to match target. Can specify "default" as well for default tuning target + self.algo.actor.max_gradient_norm = None # L2 gradient clipping for actor + + # Actor network settings + self.algo.actor.net.type = "gaussian" # Options are currently only "gaussian" (no support for GMM yet) + + # Actor network settings - shared + self.algo.actor.net.common.std_activation = "exp" # Activation to use for std output from policy net + self.algo.actor.net.common.use_tanh = True # Whether to use tanh at output of actor network + self.algo.actor.net.common.low_noise_eval = True # Whether to use deterministic action sampling at eval stage + + # Actor network settings - gaussian + self.algo.actor.net.gaussian.init_last_fc_weight = 0.001 # If set, will override the initialization of the final fc layer to be uniformly sampled limited by this value + self.algo.actor.net.gaussian.init_std = 0.3 # Relative scaling factor for std from policy net + self.algo.actor.net.gaussian.fixed_std = False # Whether to learn std dev or not + + self.algo.actor.layer_dims = (300, 400) # actor MLP layer dimensions + + # ================== Critic Network Config =================== + self.algo.critic.use_huber = False # Huber Loss instead of L2 for critic + self.algo.critic.max_gradient_norm = None # L2 gradient clipping for critic (None to use no clipping) + + self.algo.critic.value_bounds = None # optional 2-tuple to ensure lower and upper bound on value estimates + + self.algo.critic.num_action_samples = 1 # number of actions to sample per training batch to get target critic value; use maximum Q value from n random sampled actions when doing TD error backup + + # cql settings for critic + self.algo.critic.cql_weight = 1.0 # weighting for cql component of critic loss (only used if target_q_gap is < 0 or None) + self.algo.critic.deterministic_backup = True # if not set, subtract weighted logprob of action when doing backup + self.algo.critic.min_q_weight = 1.0 # min q weight (scaling factor) to apply + self.algo.critic.target_q_gap = 5.0 # if set, sets the diff threshold at which Q-values will be penalized more (note: this overrides cql weight above!) Use None or a negative value if not set + self.algo.critic.num_random_actions = 10 # Number of random actions to sample when calculating CQL loss + + # critic ensemble parameters (TD3 trick) + self.algo.critic.ensemble.n = 2 # number of Q networks in the ensemble + + self.algo.critic.layer_dims = (300, 400) # critic MLP layer dimensions diff --git a/aloha-devel/robomimic/config/diffusion_policy_config.py b/aloha-devel/robomimic/config/diffusion_policy_config.py new file mode 100644 index 0000000000000000000000000000000000000000..daef32dff0313442e6e53e690bc460ff8eedc468 --- /dev/null +++ b/aloha-devel/robomimic/config/diffusion_policy_config.py @@ -0,0 +1,60 @@ +""" +Config for Diffusion Policy algorithm. +""" + +from robomimic.config.base_config import BaseConfig + +class DiffusionPolicyConfig(BaseConfig): + ALGO_NAME = "diffusion_policy" + + def algo_config(self): + """ + This function populates the `config.algo` attribute of the config, and is given to the + `Algo` subclass (see `algo/algo.py`) for each algorithm through the `algo_config` + argument to the constructor. Any parameter that an algorithm needs to determine its + training and test-time behavior should be populated here. + """ + + # optimization parameters + self.algo.optim_params.policy.learning_rate.initial = 1e-4 # policy learning rate + self.algo.optim_params.policy.learning_rate.decay_factor = 0.1 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.policy.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.policy.regularization.L2 = 0.00 # L2 regularization strength + + # horizon parameters + self.algo.horizon.observation_horizon = 2 + self.algo.horizon.action_horizon = 8 + self.algo.horizon.prediction_horizon = 16 + + # UNet parameters + self.algo.unet.enabled = True + self.algo.unet.diffusion_step_embed_dim = 256 + self.algo.unet.down_dims = [256,512,1024] + self.algo.unet.kernel_size = 5 + self.algo.unet.n_groups = 8 + + # EMA parameters + self.algo.ema.enabled = True + self.algo.ema.power = 0.75 + + # Noise Scheduler + ## DDPM + self.algo.ddpm.enabled = True + self.algo.ddpm.num_train_timesteps = 100 + self.algo.ddpm.num_inference_timesteps = 100 + self.algo.ddpm.beta_schedule = 'squaredcos_cap_v2' + self.algo.ddpm.clip_sample = True + self.algo.ddpm.prediction_type = 'epsilon' + + ## DDIM + self.algo.ddim.enabled = False + self.algo.ddim.num_train_timesteps = 100 + self.algo.ddim.num_inference_timesteps = 10 + self.algo.ddim.beta_schedule = 'squaredcos_cap_v2' + self.algo.ddim.clip_sample = True + self.algo.ddim.set_alpha_to_one = True + self.algo.ddim.steps_offset = 0 + self.algo.ddim.prediction_type = 'epsilon' + + self.algo.language_conditioned = False # whether policy is language conditioned + diff --git a/aloha-devel/robomimic/config/gl_config.py b/aloha-devel/robomimic/config/gl_config.py new file mode 100644 index 0000000000000000000000000000000000000000..939103e65dd5f7519fb7be2c9fa1928d5b430bf2 --- /dev/null +++ b/aloha-devel/robomimic/config/gl_config.py @@ -0,0 +1,89 @@ +""" +Config for Goal Learning (sub-algorithm used by hierarchical models like HBC and IRIS). +This class of model predicts (or samples) subgoal observations given a current observation. +""" + +from robomimic.config.base_config import BaseConfig + + +class GLConfig(BaseConfig): + ALGO_NAME = "gl" + + def algo_config(self): + """ + This function populates the `config.algo` attribute of the config, and is given to the + `Algo` subclass (see `algo/algo.py`) for each algorithm through the `algo_config` + argument to the constructor. Any parameter that an algorithm needs to determine its + training and test-time behavior should be populated here. + """ + + # optimization parameters + self.algo.optim_params.goal_network.learning_rate.initial = 1e-4 # goal network learning rate + self.algo.optim_params.goal_network.learning_rate.decay_factor = 0.1 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.goal_network.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.goal_network.regularization.L2 = 0.00 + + # subgoal definition: observation that is @subgoal_horizon number of timesteps in future from current observation + self.algo.subgoal_horizon = 10 + + # MLP size for deterministic goal network (unused if VAE is enabled) + self.algo.ae.planner_layer_dims = (300, 400) + + # ================== VAE config ================== + self.algo.vae.enabled = True # set to true to use VAE network + self.algo.vae.latent_dim = 16 # VAE latent dimension + self.algo.vae.latent_clip = None # clip latent space when decoding (set to None to disable) + self.algo.vae.kl_weight = 1. # beta-VAE weight to scale KL loss relative to reconstruction loss in ELBO + + # VAE decoder settings + self.algo.vae.decoder.is_conditioned = True # whether decoder should condition on observation + self.algo.vae.decoder.reconstruction_sum_across_elements = False # sum instead of mean for reconstruction loss + + # VAE prior settings + self.algo.vae.prior.learn = False # learn Gaussian / GMM prior instead of N(0, 1) + self.algo.vae.prior.is_conditioned = False # whether to condition prior on observations + self.algo.vae.prior.use_gmm = False # whether to use GMM prior + self.algo.vae.prior.gmm_num_modes = 10 # number of GMM modes + self.algo.vae.prior.gmm_learn_weights = False # whether to learn GMM weights + self.algo.vae.prior.use_categorical = False # whether to use categorical prior + self.algo.vae.prior.categorical_dim = 10 # the number of categorical classes for each latent dimension + self.algo.vae.prior.categorical_gumbel_softmax_hard = False # use hard selection in forward pass + self.algo.vae.prior.categorical_init_temp = 1.0 # initial gumbel-softmax temp + self.algo.vae.prior.categorical_temp_anneal_step = 0.001 # linear temp annealing rate + self.algo.vae.prior.categorical_min_temp = 0.3 # lowest gumbel-softmax temp + + self.algo.vae.encoder_layer_dims = (300, 400) # encoder MLP layer dimensions + self.algo.vae.decoder_layer_dims = (300, 400) # decoder MLP layer dimensions + self.algo.vae.prior_layer_dims = (300, 400) # prior MLP layer dimensions (if learning conditioned prior) + + def observation_config(self): + """ + Update from superclass to specify subgoal modalities. + """ + super(GLConfig, self).observation_config() + self.observation.modalities.subgoal.low_dim = [ # specify low-dim subgoal observations for agent to predict + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object", + ] + self.observation.modalities.subgoal.rgb = [] # specify rgb image subgoal observations for agent to predict + self.observation.modalities.subgoal.depth = [] + self.observation.modalities.subgoal.scan = [] + self.observation.modalities.subgoal.do_not_lock_keys() + + @property + def all_obs_keys(self): + """ + Update from superclass to include subgoals. + """ + # pool all modalities + return sorted(tuple(set([ + obs_key for group in [ + self.observation.modalities.obs.values(), + self.observation.modalities.goal.values(), + self.observation.modalities.subgoal.values(), + ] + for modality in group + for obs_key in modality + ]))) diff --git a/aloha-devel/robomimic/config/iris_config.py b/aloha-devel/robomimic/config/iris_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c03328cead61f1a977d76bb4b684613586c2a08c --- /dev/null +++ b/aloha-devel/robomimic/config/iris_config.py @@ -0,0 +1,99 @@ +""" +Config for IRIS algorithm. +""" + +from robomimic.config.bcq_config import BCQConfig +from robomimic.config.gl_config import GLConfig +from robomimic.config.bc_config import BCConfig +from robomimic.config.hbc_config import HBCConfig + + +class IRISConfig(HBCConfig): + ALGO_NAME = "iris" + + def algo_config(self): + """ + This function populates the `config.algo` attribute of the config, and is given to the + `Algo` subclass (see `algo/algo.py`) for each algorithm through the `algo_config` + argument to the constructor. Any parameter that an algorithm needs to determine its + training and test-time behavior should be populated here. + """ + + # One of ["separate", "cascade"]. In "separate" mode (default), + # the planner and actor are trained independently and then the planner subgoal predictions are + # used to condition the actor at test-time. In "cascade" mode, the actor is trained directly + # on planner subgoal predictions. In "actor_only" mode, only the actor is trained, and in + # "planner_only" mode, only the planner is trained. + self.algo.mode = "separate" + + self.algo.actor_use_random_subgoals = False # whether to sample subgoal index from [1, subgoal_horizon] + self.algo.subgoal_update_interval = 10 # how frequently the subgoal should be updated at test-time (usually matches train.seq_length) + + # ================== Latent Subgoal Config ================== + + # NOTE: latent subgoals are not supported by IRIS, but superclass expects this config + self.algo.latent_subgoal.enabled = False + self.algo.latent_subgoal.prior_correction.enabled = False + self.algo.latent_subgoal.prior_correction.num_samples = 100 + + # ================== Planner Config ================== + + # The ValuePlanner planner component is a Goal Learning VAE model + self.algo.value_planner.planner = GLConfig().algo # config for goal learning + # set subgoal horizon explicitly + self.algo.value_planner.planner.subgoal_horizon = 10 + # ensure VAE is used + self.algo.value_planner.planner.vae.enabled = True + + # The ValuePlanner value component is a BCQ model + self.algo.value_planner.value = BCQConfig().algo + self.algo.value_planner.value.actor.enabled = False # ensure no BCQ actor + # number of subgoal samples to use for value planner + self.algo.value_planner.num_samples = 100 + + # ================== Actor Config =================== + self.algo.actor = BCConfig().algo + # use RNN + self.algo.actor.rnn.enabled = True + self.algo.actor.rnn.horizon = 10 + # remove unused parts of BCConfig algo config + del self.algo.actor.gaussian + del self.algo.actor.gmm + del self.algo.actor.vae + + def observation_config(self): + """ + Update from superclass so that value planner and actor each get their own obs config. + """ + self.observation.value_planner.planner = GLConfig().observation + self.observation.value_planner.value = BCQConfig().observation + self.observation.actor = BCConfig().observation + + @property + def use_goals(self): + """ + Update from superclass - value planner goal modalities determine goal-conditioning. + """ + return len( + self.observation.value_planner.planner.modalities.goal.low_dim + + self.observation.value_planner.planner.modalities.goal.rgb) > 0 + + @property + def all_obs_keys(self): + """ + Update from superclass to include modalities from value planner and actor. + """ + # pool all modalities + return sorted(tuple(set([ + obs_key for group in [ + self.observation.value_planner.planner.modalities.obs.values(), + self.observation.value_planner.planner.modalities.goal.values(), + self.observation.value_planner.planner.modalities.subgoal.values(), + self.observation.value_planner.value.modalities.obs.values(), + self.observation.value_planner.value.modalities.goal.values(), + self.observation.actor.modalities.obs.values(), + self.observation.actor.modalities.goal.values(), + ] + for modality in group + for obs_key in modality + ]))) diff --git a/aloha-devel/robomimic/config/td3_bc_config.py b/aloha-devel/robomimic/config/td3_bc_config.py new file mode 100644 index 0000000000000000000000000000000000000000..036a2591a91b4a4f5da4e2415dd035117e587900 --- /dev/null +++ b/aloha-devel/robomimic/config/td3_bc_config.py @@ -0,0 +1,111 @@ +""" +Config for TD3_BC. +""" + +from robomimic.config.base_config import BaseConfig + + +class TD3_BCConfig(BaseConfig): + ALGO_NAME = "td3_bc" + + def experiment_config(self): + """ + Update from subclass to set paper defaults for gym envs. + """ + super(TD3_BCConfig, self).experiment_config() + + # no validation and no video rendering + self.experiment.validate = False + self.experiment.render_video = False + + # save 10 checkpoints throughout training + self.experiment.save.every_n_epochs = 20 + + # save models that achieve best rollout return instead of best success rate + self.experiment.save.on_best_rollout_return = True + self.experiment.save.on_best_rollout_success_rate = False + + # epoch definition - 5000 gradient steps per epoch, with 200 epochs = 1M gradient steps, and eval every 1 epochs + self.experiment.epoch_every_n_steps = 5000 + + # evaluate with normal environment rollouts + self.experiment.rollout.enabled = True + self.experiment.rollout.n = 50 # paper uses 10, but we can afford to do 50 + self.experiment.rollout.horizon = 1000 + self.experiment.rollout.rate = 1 # rollout every epoch to match paper + + def train_config(self): + """ + Update from subclass to set paper defaults for gym envs. + """ + super(TD3_BCConfig, self).train_config() + + # update to normalize observations + self.train.hdf5_normalize_obs = True + + # increase batch size to 256 + self.train.batch_size = 256 + + # 200 epochs, with each epoch lasting 5000 gradient steps, for 1M total steps + self.train.num_epochs = 200 + + def algo_config(self): + """ + This function populates the `config.algo` attribute of the config, and is given to the + `Algo` subclass (see `algo/algo.py`) for each algorithm through the `algo_config` + argument to the constructor. Any parameter that an algorithm needs to determine its + training and test-time behavior should be populated here. + """ + + # optimization parameters + self.algo.optim_params.critic.learning_rate.initial = 3e-4 # critic learning rate + self.algo.optim_params.critic.learning_rate.decay_factor = 0.1 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.critic.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.critic.regularization.L2 = 0.00 # L2 regularization strength + self.algo.optim_params.critic.start_epoch = -1 # number of epochs before starting critic training (-1 means start right away) + self.algo.optim_params.critic.end_epoch = -1 # number of epochs before ending critic training (-1 means start right away) + + self.algo.optim_params.actor.learning_rate.initial = 3e-4 # actor learning rate + self.algo.optim_params.actor.learning_rate.decay_factor = 0.1 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.actor.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.actor.regularization.L2 = 0.00 # L2 regularization strength + self.algo.optim_params.actor.start_epoch = -1 # number of epochs before starting actor training (-1 means start right away) + self.algo.optim_params.actor.end_epoch = -1 # number of epochs before ending actor training (-1 means start right away) + + # alpha value - for weighting critic loss vs. BC loss + self.algo.alpha = 2.5 + + # target network related parameters + self.algo.discount = 0.99 # discount factor to use + self.algo.n_step = 1 # for using n-step returns in TD-updates + self.algo.target_tau = 0.005 # update rate for target networks + self.algo.infinite_horizon = False # if True, scale terminal rewards by 1 / (1 - discount) to treat as infinite horizon + + # ================== Critic Network Config =================== + self.algo.critic.use_huber = False # Huber Loss instead of L2 for critic + self.algo.critic.max_gradient_norm = None # L2 gradient clipping for critic (None to use no clipping) + self.algo.critic.value_bounds = None # optional 2-tuple to ensure lower and upper bound on value estimates + + # critic ensemble parameters (TD3 trick) + self.algo.critic.ensemble.n = 2 # number of Q networks in the ensemble + self.algo.critic.ensemble.weight = 1.0 # weighting for mixing min and max for target Q value + + self.algo.critic.layer_dims = (256, 256) # size of critic MLP + + # ================== Actor Network Config =================== + + # update actor and target networks every n gradients steps for each critic gradient step + self.algo.actor.update_freq = 2 + + # exploration noise used to form target action for Q-update - clipped Gaussian noise + self.algo.actor.noise_std = 0.2 # zero-mean gaussian noise with this std is applied to actions + self.algo.actor.noise_clip = 0.5 # noise is clipped in each dimension to (-noise_clip, noise_clip) + + self.algo.actor.layer_dims = (256, 256) # size of actor MLP + + def observation_config(self): + """ + Update from superclass to use flat observations from gym envs. + """ + super(TD3_BCConfig, self).observation_config() + self.observation.modalities.obs.low_dim = ["flat"] diff --git a/aloha-devel/robomimic/exps/templates/bcq.json b/aloha-devel/robomimic/exps/templates/bcq.json new file mode 100644 index 0000000000000000000000000000000000000000..5ae9d907466f4278b418bcc1fb93aacb7fcb1e2a --- /dev/null +++ b/aloha-devel/robomimic/exps/templates/bcq.json @@ -0,0 +1,235 @@ +{ + "algo_name": "bcq", + "experiment": { + "name": "test", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": true, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../bcq_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "hdf5_validation_filter_key": null, + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "critic": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + }, + "action_sampler": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + }, + "actor": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + } + }, + "discount": 0.99, + "n_step": 1, + "target_tau": 0.005, + "infinite_horizon": false, + "critic": { + "use_huber": false, + "max_gradient_norm": null, + "value_bounds": null, + "num_action_samples": 10, + "num_action_samples_rollout": 100, + "ensemble": { + "n": 2, + "weight": 0.75 + }, + "distributional": { + "enabled": false, + "num_atoms": 51 + }, + "layer_dims": [ + 300, + 400 + ] + }, + "action_sampler": { + "actor_layer_dims": [ + 1024, + 1024 + ], + "gmm": { + "enabled": false, + "num_modes": 5, + "min_std": 0.0001, + "std_activation": "softplus", + "low_noise_eval": true + }, + "vae": { + "enabled": true, + "latent_dim": 14, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + }, + "freeze_encoder_epoch": -1 + }, + "actor": { + "enabled": false, + "perturbation_scale": 0.05, + "layer_dims": [ + 300, + 400 + ] + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} \ No newline at end of file diff --git a/aloha-devel/robomimic/exps/templates/diffusion_policy.json b/aloha-devel/robomimic/exps/templates/diffusion_policy.json new file mode 100644 index 0000000000000000000000000000000000000000..c943fee71025e51a7a8b91865af1fbf59390e1f9 --- /dev/null +++ b/aloha-devel/robomimic/exps/templates/diffusion_policy.json @@ -0,0 +1,175 @@ +{ + "algo_name": "diffusion_policy", + "experiment": { + "name": "test", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "mse":{}, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": true, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir":"../diffusion_policy_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "low_dim", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": false, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "seq_length": 15, + "pad_seq_length": true, + "frame_stack": 2, + "pad_frame_stack": true, + "dataset_keys": [ + "actions" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 256, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "policy": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + } + }, + "horizon": { + "observation_horizon": 2, + "action_horizon": 8, + "prediction_horizon": 16 + }, + "unet": { + "enabled": true, + "diffusion_step_embed_dim": 256, + "down_dims": [256,512,1024], + "kernel_size": 5, + "n_groups": 8 + }, + "ema": { + "enabled": true, + "power": 0.75 + }, + "ddpm": { + "enabled": true, + "num_train_timesteps": 100, + "num_inference_timesteps": 100, + "beta_schedule": "squaredcos_cap_v2", + "clip_sample": true, + "prediction_type": "epsilon" + }, + "ddim": { + "enabled": false, + "num_train_timesteps": 100, + "num_inference_timesteps": 10, + "beta_schedule": "squaredcos_cap_v2", + "clip_sample": true, + "set_alpha_to_one": true, + "steps_offset": 0, + "prediction_type": "epsilon" + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": { + "feature_dimension": 64, + "backbone_class": "ResNet18Conv", + "backbone_kwargs": { + "pretrained": false, + "input_coord_conv": false + }, + "pool_class": "SpatialSoftmax", + "pool_kwargs": { + "num_kp": 32, + "learnable_temperature": false, + "temperature": 1.0, + "noise_std": 0.0 + } + }, + "obs_randomizer_class": "CropRandomizer", + "obs_randomizer_kwargs": { + "crop_height": 76, + "crop_width": 76, + "num_crops": 1, + "pos_enc": false + } + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + } +} \ No newline at end of file diff --git a/aloha-devel/robomimic/models/__init__.py b/aloha-devel/robomimic/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7460f9309af64c4578b547e0944c7e1366b5946c --- /dev/null +++ b/aloha-devel/robomimic/models/__init__.py @@ -0,0 +1 @@ +from .obs_core import EncoderCore, Randomizer diff --git a/aloha-devel/robomimic/models/__pycache__/obs_core.cpython-38.pyc b/aloha-devel/robomimic/models/__pycache__/obs_core.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..00e89878396e7b3076a44a8fccc7d7a0c369e06b Binary files /dev/null and b/aloha-devel/robomimic/models/__pycache__/obs_core.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/models/__pycache__/obs_nets.cpython-38.pyc b/aloha-devel/robomimic/models/__pycache__/obs_nets.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b8da6df6dc53bf37b321c85b6a4a54091c60415 Binary files /dev/null and b/aloha-devel/robomimic/models/__pycache__/obs_nets.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/models/__pycache__/vae_nets.cpython-38.pyc b/aloha-devel/robomimic/models/__pycache__/vae_nets.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d140848ef73ec7864bcbe5e5c69ba6fcd40ca525 Binary files /dev/null and b/aloha-devel/robomimic/models/__pycache__/vae_nets.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/models/distributions.py b/aloha-devel/robomimic/models/distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..411efb1a8bbc6b0da7ac6f628357dc9c178b8780 --- /dev/null +++ b/aloha-devel/robomimic/models/distributions.py @@ -0,0 +1,123 @@ +""" +Contains distribution models used as parts of other networks. These +classes usually inherit or emulate torch distributions. +""" +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.distributions as D + + +class TanhWrappedDistribution(D.Distribution): + """ + Class that wraps another valid torch distribution, such that sampled values from the base distribution are + passed through a tanh layer. The corresponding (log) probabilities are also modified accordingly. + Tanh Normal distribution - adapted from rlkit and CQL codebase + (https://github.com/aviralkumar2907/CQL/blob/d67dbe9cf5d2b96e3b462b6146f249b3d6569796/d4rl/rlkit/torch/distributions.py#L6). + """ + def __init__(self, base_dist, scale=1.0, epsilon=1e-6): + """ + Args: + base_dist (Distribution): Distribution to wrap with tanh output + scale (float): Scale of output + epsilon (float): Numerical stability epsilon when computing log-prob. + """ + self.base_dist = base_dist + self.scale = scale + self.tanh_epsilon = epsilon + super(TanhWrappedDistribution, self).__init__() + + def log_prob(self, value, pre_tanh_value=None): + """ + Args: + value (torch.Tensor): some tensor to compute log probabilities for + pre_tanh_value: If specified, will not calculate atanh manually from @value. More numerically stable + """ + value = value / self.scale + if pre_tanh_value is None: + one_plus_x = (1. + value).clamp(min=self.tanh_epsilon) + one_minus_x = (1. - value).clamp(min=self.tanh_epsilon) + pre_tanh_value = 0.5 * torch.log(one_plus_x / one_minus_x) + lp = self.base_dist.log_prob(pre_tanh_value) + tanh_lp = torch.log(1 - value * value + self.tanh_epsilon) + # In case the base dist already sums up the log probs, make sure we do the same + return lp - tanh_lp if len(lp.shape) == len(tanh_lp.shape) else lp - tanh_lp.sum(-1) + + def sample(self, sample_shape=torch.Size(), return_pretanh_value=False): + """ + Gradients will and should *not* pass through this operation. + See https://github.com/pytorch/pytorch/issues/4620 for discussion. + """ + z = self.base_dist.sample(sample_shape=sample_shape).detach() + + if return_pretanh_value: + return torch.tanh(z) * self.scale, z + else: + return torch.tanh(z) * self.scale + + def rsample(self, sample_shape=torch.Size(), return_pretanh_value=False): + """ + Sampling in the reparameterization case - for differentiable samples. + """ + z = self.base_dist.rsample(sample_shape=sample_shape) + + if return_pretanh_value: + return torch.tanh(z) * self.scale, z + else: + return torch.tanh(z) * self.scale + + @property + def mean(self): + return self.base_dist.mean + + @property + def stddev(self): + return self.base_dist.stddev + + +class DiscreteValueDistribution(object): + """ + Extension to torch categorical probability distribution in order to keep track + of the support (categorical values, or in this case, value atoms). This is + used for distributional value networks. + """ + def __init__(self, values, probs=None, logits=None): + """ + Creates a categorical distribution parameterized by either @probs or + @logits (but not both). Expects inputs to be consistent in shape + for broadcasting operations (e.g. multiplication). + """ + self._values = values + self._categorical_dist = D.Categorical(probs=probs, logits=logits) + + @property + def values(self): + return self._values + + @property + def probs(self): + return self._categorical_dist.probs + + @property + def logits(self): + return self._categorical_dist.logits + + def mean(self): + """ + Categorical distribution mean, taking the value support into account. + """ + return (self._categorical_dist.probs * self._values).sum(dim=-1) + + def variance(self): + """ + Categorical distribution variance, taking the value support into account. + """ + dist_squared = (self.mean().unsqueeze(-1) - self.values).pow(2) + return (self._categorical_dist.probs * dist_squared).sum(dim=-1) + + def sample(self, sample_shape=torch.Size()): + """ + Sample from the distribution. Make sure to return value atoms, not categorical class indices. + """ + inds = self._categorical_dist.sample(sample_shape=sample_shape) + return torch.gather(self.values, inds, dim=-1) diff --git a/aloha-devel/robomimic/models/obs_core.py b/aloha-devel/robomimic/models/obs_core.py new file mode 100644 index 0000000000000000000000000000000000000000..3566f82951a89ff6bd2da628dceeddc8b5cd85ad --- /dev/null +++ b/aloha-devel/robomimic/models/obs_core.py @@ -0,0 +1,829 @@ +""" +Contains torch Modules for core observation processing blocks +such as encoders (e.g. EncoderCore, VisualCore, ScanCore, ...) +and randomizers (e.g. Randomizer, CropRandomizer). +""" + +import abc +import numpy as np +import textwrap +import random + +import torch +import torch.nn as nn + +import robomimic.models.base_nets as BaseNets +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.obs_utils as ObsUtils +from robomimic.utils.python_utils import extract_class_init_kwargs_from_dict + +# NOTE: this is required for the backbone classes to be found by the `eval` call in the core networks +from robomimic.models.base_nets import * +from robomimic.utils.vis_utils import visualize_image_randomizer +from robomimic.macros import VISUALIZE_RANDOMIZER + +import torchvision.transforms.functional as TVF +from torchvision.transforms import Lambda, Compose + +""" +================================================ +Encoder Core Networks (Abstract class) +================================================ +""" +class EncoderCore(BaseNets.Module): + """ + Abstract class used to categorize all cores used to encode observations + """ + def __init__(self, input_shape): + self.input_shape = input_shape + super(EncoderCore, self).__init__() + + def __init_subclass__(cls, **kwargs): + """ + Hook method to automatically register all valid subclasses so we can keep track of valid observation encoders + in a global dict. + + This global dict stores mapping from observation encoder network name to class. + We keep track of these registries to enable automated class inference at runtime, allowing + users to simply extend our base encoder class and refer to that class in string form + in their config, without having to manually register their class internally. + This also future-proofs us for any additional encoder classes we would + like to add ourselves. + """ + ObsUtils.register_encoder_core(cls) + + +""" +================================================ +Visual Core Networks (Backbone + Pool) +================================================ +""" +class VisualCore(EncoderCore, BaseNets.ConvBase): + """ + A network block that combines a visual backbone network with optional pooling + and linear layers. + """ + def __init__( + self, + input_shape, + backbone_class="ResNet18Conv", + pool_class="SpatialSoftmax", + backbone_kwargs=None, + pool_kwargs=None, + flatten=True, + feature_dimension=64, + ): + """ + Args: + input_shape (tuple): shape of input (not including batch dimension) + backbone_class (str): class name for the visual backbone network. Defaults + to "ResNet18Conv". + pool_class (str): class name for the visual feature pooler (optional) + Common options are "SpatialSoftmax" and "SpatialMeanPool". Defaults to + "SpatialSoftmax". + backbone_kwargs (dict): kwargs for the visual backbone network (optional) + pool_kwargs (dict): kwargs for the visual feature pooler (optional) + flatten (bool): whether to flatten the visual features + feature_dimension (int): if not None, add a Linear layer to + project output into a desired feature dimension + """ + super(VisualCore, self).__init__(input_shape=input_shape) + self.flatten = flatten + + if backbone_kwargs is None: + backbone_kwargs = dict() + + # add input channel dimension to visual core inputs + backbone_kwargs["input_channel"] = input_shape[0] + + # extract only relevant kwargs for this specific backbone + backbone_kwargs = extract_class_init_kwargs_from_dict(cls=eval(backbone_class), dic=backbone_kwargs, copy=True) + + # visual backbone + assert isinstance(backbone_class, str) + self.backbone = eval(backbone_class)(**backbone_kwargs) + + assert isinstance(self.backbone, BaseNets.ConvBase) + + feat_shape = self.backbone.output_shape(input_shape) + net_list = [self.backbone] + + # maybe make pool net + if pool_class is not None: + assert isinstance(pool_class, str) + # feed output shape of backbone to pool net + if pool_kwargs is None: + pool_kwargs = dict() + # extract only relevant kwargs for this specific backbone + pool_kwargs["input_shape"] = feat_shape + pool_kwargs = extract_class_init_kwargs_from_dict(cls=eval(pool_class), dic=pool_kwargs, copy=True) + self.pool = eval(pool_class)(**pool_kwargs) + assert isinstance(self.pool, BaseNets.Module) + + feat_shape = self.pool.output_shape(feat_shape) + net_list.append(self.pool) + else: + self.pool = None + + # flatten layer + if self.flatten: + net_list.append(torch.nn.Flatten(start_dim=1, end_dim=-1)) + + # maybe linear layer + self.feature_dimension = feature_dimension + if feature_dimension is not None: + assert self.flatten + linear = torch.nn.Linear(int(np.prod(feat_shape)), feature_dimension) + net_list.append(linear) + + self.nets = nn.Sequential(*net_list) + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + if self.feature_dimension is not None: + # linear output + return [self.feature_dimension] + feat_shape = self.backbone.output_shape(input_shape) + if self.pool is not None: + # pool output + feat_shape = self.pool.output_shape(feat_shape) + # backbone + flat output + if self.flatten: + return [np.prod(feat_shape)] + else: + return feat_shape + + def forward(self, inputs): + """ + Forward pass through visual core. + """ + ndim = len(self.input_shape) + assert tuple(inputs.shape)[-ndim:] == tuple(self.input_shape) + return super(VisualCore, self).forward(inputs) + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = '' + indent = ' ' * 2 + msg += textwrap.indent( + "\ninput_shape={}\noutput_shape={}".format(self.input_shape, self.output_shape(self.input_shape)), indent) + msg += textwrap.indent("\nbackbone_net={}".format(self.backbone), indent) + msg += textwrap.indent("\npool_net={}".format(self.pool), indent) + msg = header + '(' + msg + '\n)' + return msg + + +""" +================================================ +Scan Core Networks (Conv1D Sequential + Pool) +================================================ +""" +class ScanCore(EncoderCore, BaseNets.ConvBase): + """ + A network block that combines a Conv1D backbone network with optional pooling + and linear layers. + """ + def __init__( + self, + input_shape, + conv_kwargs=None, + conv_activation="relu", + pool_class=None, + pool_kwargs=None, + flatten=True, + feature_dimension=None, + ): + """ + Args: + input_shape (tuple): shape of input (not including batch dimension) + conv_kwargs (dict): kwargs for the conv1d backbone network. Should contain lists for the following values: + out_channels (int) + kernel_size (int) + stride (int) + ... + + If not specified, or an empty dictionary is specified, some default settings will be used. + conv_activation (str or None): Activation to use between conv layers. Default is relu. + Currently, valid options are {relu} + pool_class (str): class name for the visual feature pooler (optional) + Common options are "SpatialSoftmax" and "SpatialMeanPool" + pool_kwargs (dict): kwargs for the visual feature pooler (optional) + flatten (bool): whether to flatten the network output + feature_dimension (int): if not None, add a Linear layer to + project output into a desired feature dimension (note: flatten must be set to True!) + """ + super(ScanCore, self).__init__(input_shape=input_shape) + self.flatten = flatten + self.feature_dimension = feature_dimension + + if conv_kwargs is None: + conv_kwargs = dict() + + # Generate backbone network + self.backbone = BaseNets.Conv1dBase( + input_channel=1, + activation=conv_activation, + **conv_kwargs, + ) + feat_shape = self.backbone.output_shape(input_shape=input_shape) + + # Create netlist of all generated networks + net_list = [self.backbone] + + # Possibly add pooling network + if pool_class is not None: + # Add an unsqueeze network so that the shape is correct to pass to pooling network + self.unsqueeze = Unsqueeze(dim=-1) + net_list.append(self.unsqueeze) + # Get output shape + feat_shape = self.unsqueeze.output_shape(feat_shape) + # Create pooling network + self.pool = eval(pool_class)(input_shape=feat_shape, **pool_kwargs) + net_list.append(self.pool) + feat_shape = self.pool.output_shape(feat_shape) + else: + self.unsqueeze, self.pool = None, None + + # flatten layer + if self.flatten: + net_list.append(torch.nn.Flatten(start_dim=1, end_dim=-1)) + + # maybe linear layer + if self.feature_dimension is not None: + assert self.flatten + linear = torch.nn.Linear(int(np.prod(feat_shape)), self.feature_dimension) + net_list.append(linear) + + # Generate final network + self.nets = nn.Sequential(*net_list) + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + if self.feature_dimension is not None: + # linear output + return [self.feature_dimension] + feat_shape = self.backbone.output_shape(input_shape) + if self.pool is not None: + # pool output + feat_shape = self.pool.output_shape(self.unsqueeze.output_shape(feat_shape)) + # backbone + flat output + return [np.prod(feat_shape)] if self.flatten else feat_shape + + def forward(self, inputs): + """ + Forward pass through visual core. + """ + ndim = len(self.input_shape) + assert tuple(inputs.shape)[-ndim:] == tuple(self.input_shape) + return super(ScanCore, self).forward(inputs) + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = '' + indent = ' ' * 2 + msg += textwrap.indent( + "\ninput_shape={}\noutput_shape={}".format(self.input_shape, self.output_shape(self.input_shape)), indent) + msg += textwrap.indent("\nbackbone_net={}".format(self.backbone), indent) + msg += textwrap.indent("\npool_net={}".format(self.pool), indent) + msg = header + '(' + msg + '\n)' + return msg + + +""" +================================================ +Observation Randomizer Networks +================================================ +""" +class Randomizer(BaseNets.Module): + """ + Base class for randomizer networks. Each randomizer should implement the @output_shape_in, + @output_shape_out, @forward_in, and @forward_out methods. The randomizer's @forward_in + method is invoked on raw inputs, and @forward_out is invoked on processed inputs + (usually processed by a @VisualCore instance). Note that the self.training property + can be used to change the randomizer's behavior at train vs. test time. + """ + def __init__(self): + super(Randomizer, self).__init__() + + def __init_subclass__(cls, **kwargs): + """ + Hook method to automatically register all valid subclasses so we can keep track of valid observation randomizers + in a global dict. + + This global dict stores mapping from observation randomizer network name to class. + We keep track of these registries to enable automated class inference at runtime, allowing + users to simply extend our base randomizer class and refer to that class in string form + in their config, without having to manually register their class internally. + This also future-proofs us for any additional randomizer classes we would + like to add ourselves. + """ + ObsUtils.register_randomizer(cls) + + def output_shape(self, input_shape=None): + """ + This function is unused. See @output_shape_in and @output_shape_out. + """ + raise NotImplementedError + + @abc.abstractmethod + def output_shape_in(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. Corresponds to + the @forward_in operation, where raw inputs (usually observation modalities) + are passed in. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + raise NotImplementedError + + @abc.abstractmethod + def output_shape_out(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. Corresponds to + the @forward_out operation, where processed inputs (usually encoded observation + modalities) are passed in. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + raise NotImplementedError + + def forward_in(self, inputs): + """ + Randomize raw inputs if training. + """ + if self.training: + randomized_inputs = self._forward_in(inputs=inputs) + if VISUALIZE_RANDOMIZER: + num_samples_to_visualize = min(4, inputs.shape[0]) + self._visualize(inputs, randomized_inputs, num_samples_to_visualize=num_samples_to_visualize) + return randomized_inputs + else: + return self._forward_in_eval(inputs) + + def forward_out(self, inputs): + """ + Processing for network outputs. + """ + if self.training: + return self._forward_out(inputs) + else: + return self._forward_out_eval(inputs) + + @abc.abstractmethod + def _forward_in(self, inputs): + """ + Randomize raw inputs. + """ + raise NotImplementedError + + def _forward_in_eval(self, inputs): + """ + Test-time behavior for the randomizer + """ + return inputs + + @abc.abstractmethod + def _forward_out(self, inputs): + """ + Processing for network outputs. + """ + return inputs + + def _forward_out_eval(self, inputs): + """ + Test-time behavior for the randomizer + """ + return inputs + + @abc.abstractmethod + def _visualize(self, pre_random_input, randomized_input, num_samples_to_visualize=2): + """ + Visualize the original input and the randomized input for _forward_in for debugging purposes. + """ + pass + + +class CropRandomizer(Randomizer): + """ + Randomly sample crops at input, and then average across crop features at output. + """ + def __init__( + self, + input_shape, + crop_height=76, + crop_width=76, + num_crops=1, + pos_enc=False, + ): + """ + Args: + input_shape (tuple, list): shape of input (not including batch dimension) + crop_height (int): crop height + crop_width (int): crop width + num_crops (int): number of random crops to take + pos_enc (bool): if True, add 2 channels to the output to encode the spatial + location of the cropped pixels in the source image + """ + super(CropRandomizer, self).__init__() + + assert len(input_shape) == 3 # (C, H, W) + assert crop_height < input_shape[1] + assert crop_width < input_shape[2] + + self.input_shape = input_shape + self.crop_height = crop_height + self.crop_width = crop_width + self.num_crops = num_crops + self.pos_enc = pos_enc + + def output_shape_in(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. Corresponds to + the @forward_in operation, where raw inputs (usually observation modalities) + are passed in. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + + # outputs are shape (C, CH, CW), or maybe C + 2 if using position encoding, because + # the number of crops are reshaped into the batch dimension, increasing the batch + # size from B to B * N + out_c = self.input_shape[0] + 2 if self.pos_enc else self.input_shape[0] + return [out_c, self.crop_height, self.crop_width] + + def output_shape_out(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. Corresponds to + the @forward_out operation, where processed inputs (usually encoded observation + modalities) are passed in. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + + # since the forward_out operation splits [B * N, ...] -> [B, N, ...] + # and then pools to result in [B, ...], only the batch dimension changes, + # and so the other dimensions retain their shape. + return list(input_shape) + + def _forward_in(self, inputs): + """ + Samples N random crops for each input in the batch, and then reshapes + inputs to [B * N, ...]. + """ + assert len(inputs.shape) >= 3 # must have at least (C, H, W) dimensions + out, _ = ObsUtils.sample_random_image_crops( + images=inputs, + crop_height=self.crop_height, + crop_width=self.crop_width, + num_crops=self.num_crops, + pos_enc=self.pos_enc, + ) + # [B, N, ...] -> [B * N, ...] + return TensorUtils.join_dimensions(out, 0, 1) + + def _forward_in_eval(self, inputs): + """ + Do center crops during eval + """ + assert len(inputs.shape) >= 3 # must have at least (C, H, W) dimensions + inputs = inputs.permute(*range(inputs.dim()-3), inputs.dim()-2, inputs.dim()-1, inputs.dim()-3) + out = ObsUtils.center_crop(inputs, self.crop_height, self.crop_width) + out = out.permute(*range(out.dim()-3), out.dim()-1, out.dim()-3, out.dim()-2) + return out + + def _forward_out(self, inputs): + """ + Splits the outputs from shape [B * N, ...] -> [B, N, ...] and then average across N + to result in shape [B, ...] to make sure the network output is consistent with + what would have happened if there were no randomization. + """ + batch_size = (inputs.shape[0] // self.num_crops) + out = TensorUtils.reshape_dimensions(inputs, begin_axis=0, end_axis=0, + target_dims=(batch_size, self.num_crops)) + return out.mean(dim=1) + + def _visualize(self, pre_random_input, randomized_input, num_samples_to_visualize=2): + batch_size = pre_random_input.shape[0] + random_sample_inds = torch.randint(0, batch_size, size=(num_samples_to_visualize,)) + pre_random_input_np = TensorUtils.to_numpy(pre_random_input)[random_sample_inds] + randomized_input = TensorUtils.reshape_dimensions( + randomized_input, + begin_axis=0, + end_axis=0, + target_dims=(batch_size, self.num_crops) + ) # [B * N, ...] -> [B, N, ...] + randomized_input_np = TensorUtils.to_numpy(randomized_input[random_sample_inds]) + + pre_random_input_np = pre_random_input_np.transpose((0, 2, 3, 1)) # [B, C, H, W] -> [B, H, W, C] + randomized_input_np = randomized_input_np.transpose((0, 1, 3, 4, 2)) # [B, N, C, H, W] -> [B, N, H, W, C] + + visualize_image_randomizer( + pre_random_input_np, + randomized_input_np, + randomizer_name='{}'.format(str(self.__class__.__name__)) + ) + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = header + "(input_shape={}, crop_size=[{}, {}], num_crops={})".format( + self.input_shape, self.crop_height, self.crop_width, self.num_crops) + return msg + + +class ColorRandomizer(Randomizer): + """ + Randomly sample color jitter at input, and then average across color jtters at output. + """ + def __init__( + self, + input_shape, + brightness=0.3, + contrast=0.3, + saturation=0.3, + hue=0.3, + num_samples=1, + ): + """ + Args: + input_shape (tuple, list): shape of input (not including batch dimension) + brightness (None or float or 2-tuple): How much to jitter brightness. brightness_factor is chosen uniformly + from [max(0, 1 - brightness), 1 + brightness] or the given [min, max]. Should be non negative numbers. + contrast (None or float or 2-tuple): How much to jitter contrast. contrast_factor is chosen uniformly + from [max(0, 1 - contrast), 1 + contrast] or the given [min, max]. Should be non negative numbers. + saturation (None or float or 2-tuple): How much to jitter saturation. saturation_factor is chosen uniformly + from [max(0, 1 - saturation), 1 + saturation] or the given [min, max]. Should be non negative numbers. + hue (None or float or 2-tuple): How much to jitter hue. hue_factor is chosen uniformly from [-hue, hue] or + the given [min, max]. Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5. To jitter hue, the pixel + values of the input image has to be non-negative for conversion to HSV space; thus it does not work + if you normalize your image to an interval with negative values, or use an interpolation that + generates negative values before using this function. + num_samples (int): number of random color jitters to take + """ + super(ColorRandomizer, self).__init__() + + assert len(input_shape) == 3 # (C, H, W) + + self.input_shape = input_shape + self.brightness = [max(0, 1 - brightness), 1 + brightness] if type(brightness) in {float, int} else brightness + self.contrast = [max(0, 1 - contrast), 1 + contrast] if type(contrast) in {float, int} else contrast + self.saturation = [max(0, 1 - saturation), 1 + saturation] if type(saturation) in {float, int} else saturation + self.hue = [-hue, hue] if type(hue) in {float, int} else hue + self.num_samples = num_samples + + @torch.jit.unused + def get_transform(self): + """ + Get a randomized transform to be applied on image. + + Implementation taken directly from: + + https://github.com/pytorch/vision/blob/2f40a483d73018ae6e1488a484c5927f2b309969/torchvision/transforms/transforms.py#L1053-L1085 + + Returns: + Transform: Transform which randomly adjusts brightness, contrast and + saturation in a random order. + """ + transforms = [] + + if self.brightness is not None: + brightness_factor = random.uniform(self.brightness[0], self.brightness[1]) + transforms.append(Lambda(lambda img: TVF.adjust_brightness(img, brightness_factor))) + + if self.contrast is not None: + contrast_factor = random.uniform(self.contrast[0], self.contrast[1]) + transforms.append(Lambda(lambda img: TVF.adjust_contrast(img, contrast_factor))) + + if self.saturation is not None: + saturation_factor = random.uniform(self.saturation[0], self.saturation[1]) + transforms.append(Lambda(lambda img: TVF.adjust_saturation(img, saturation_factor))) + + if self.hue is not None: + hue_factor = random.uniform(self.hue[0], self.hue[1]) + transforms.append(Lambda(lambda img: TVF.adjust_hue(img, hue_factor))) + + random.shuffle(transforms) + transform = Compose(transforms) + + return transform + + def get_batch_transform(self, N): + """ + Generates a batch transform, where each set of sample(s) along the batch (first) dimension will have the same + @N unique ColorJitter transforms applied. + + Args: + N (int): Number of ColorJitter transforms to apply per set of sample(s) along the batch (first) dimension + + Returns: + Lambda: Aggregated transform which will autoamtically apply a different ColorJitter transforms to + each sub-set of samples along batch dimension, assumed to be the FIRST dimension in the inputted tensor + Note: This function will MULTIPLY the first dimension by N + """ + return Lambda(lambda x: torch.stack([self.get_transform()(x_) for x_ in x for _ in range(N)])) + + def output_shape_in(self, input_shape=None): + # outputs are same shape as inputs + return list(input_shape) + + def output_shape_out(self, input_shape=None): + # since the forward_out operation splits [B * N, ...] -> [B, N, ...] + # and then pools to result in [B, ...], only the batch dimension changes, + # and so the other dimensions retain their shape. + return list(input_shape) + + def _forward_in(self, inputs): + """ + Samples N random color jitters for each input in the batch, and then reshapes + inputs to [B * N, ...]. + """ + assert len(inputs.shape) >= 3 # must have at least (C, H, W) dimensions + + # Make sure shape is exactly 4 + if len(inputs.shape) == 3: + inputs = torch.unsqueeze(inputs, dim=0) + + # TODO: Make more efficient other than implicit for-loop? + # Create lambda to aggregate all color randomizings at once + transform = self.get_batch_transform(N=self.num_samples) + + return transform(inputs) + + def _forward_out(self, inputs): + """ + Splits the outputs from shape [B * N, ...] -> [B, N, ...] and then average across N + to result in shape [B, ...] to make sure the network output is consistent with + what would have happened if there were no randomization. + """ + batch_size = (inputs.shape[0] // self.num_samples) + out = TensorUtils.reshape_dimensions(inputs, begin_axis=0, end_axis=0, + target_dims=(batch_size, self.num_samples)) + return out.mean(dim=1) + + def _visualize(self, pre_random_input, randomized_input, num_samples_to_visualize=2): + batch_size = pre_random_input.shape[0] + random_sample_inds = torch.randint(0, batch_size, size=(num_samples_to_visualize,)) + pre_random_input_np = TensorUtils.to_numpy(pre_random_input)[random_sample_inds] + randomized_input = TensorUtils.reshape_dimensions( + randomized_input, + begin_axis=0, + end_axis=0, + target_dims=(batch_size, self.num_samples) + ) # [B * N, ...] -> [B, N, ...] + randomized_input_np = TensorUtils.to_numpy(randomized_input[random_sample_inds]) + + pre_random_input_np = pre_random_input_np.transpose((0, 2, 3, 1)) # [B, C, H, W] -> [B, H, W, C] + randomized_input_np = randomized_input_np.transpose((0, 1, 3, 4, 2)) # [B, N, C, H, W] -> [B, N, H, W, C] + + visualize_image_randomizer( + pre_random_input_np, + randomized_input_np, + randomizer_name='{}'.format(str(self.__class__.__name__)) + ) + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = header + f"(input_shape={self.input_shape}, brightness={self.brightness}, contrast={self.contrast}, " \ + f"saturation={self.saturation}, hue={self.hue}, num_samples={self.num_samples})" + return msg + + +class GaussianNoiseRandomizer(Randomizer): + """ + Randomly sample gaussian noise at input, and then average across noises at output. + """ + def __init__( + self, + input_shape, + noise_mean=0.0, + noise_std=0.3, + limits=None, + num_samples=1, + ): + """ + Args: + input_shape (tuple, list): shape of input (not including batch dimension) + noise_mean (float): Mean of noise to apply + noise_std (float): Standard deviation of noise to apply + limits (None or 2-tuple): If specified, should be the (min, max) values to clamp all noisied samples to + num_samples (int): number of random color jitters to take + """ + super(GaussianNoiseRandomizer, self).__init__() + + self.input_shape = input_shape + self.noise_mean = noise_mean + self.noise_std = noise_std + self.limits = limits + self.num_samples = num_samples + + def output_shape_in(self, input_shape=None): + # outputs are same shape as inputs + return list(input_shape) + + def output_shape_out(self, input_shape=None): + # since the forward_out operation splits [B * N, ...] -> [B, N, ...] + # and then pools to result in [B, ...], only the batch dimension changes, + # and so the other dimensions retain their shape. + return list(input_shape) + + def _forward_in(self, inputs): + """ + Samples N random gaussian noises for each input in the batch, and then reshapes + inputs to [B * N, ...]. + """ + out = TensorUtils.repeat_by_expand_at(inputs, repeats=self.num_samples, dim=0) + + # Sample noise across all samples + out = torch.rand(size=out.shape) * self.noise_std + self.noise_mean + out + + # Possibly clamp + if self.limits is not None: + out = torch.clip(out, min=self.limits[0], max=self.limits[1]) + + return out + + def _forward_out(self, inputs): + """ + Splits the outputs from shape [B * N, ...] -> [B, N, ...] and then average across N + to result in shape [B, ...] to make sure the network output is consistent with + what would have happened if there were no randomization. + """ + batch_size = (inputs.shape[0] // self.num_samples) + out = TensorUtils.reshape_dimensions(inputs, begin_axis=0, end_axis=0, + target_dims=(batch_size, self.num_samples)) + return out.mean(dim=1) + + def _visualize(self, pre_random_input, randomized_input, num_samples_to_visualize=2): + batch_size = pre_random_input.shape[0] + random_sample_inds = torch.randint(0, batch_size, size=(num_samples_to_visualize,)) + pre_random_input_np = TensorUtils.to_numpy(pre_random_input)[random_sample_inds] + randomized_input = TensorUtils.reshape_dimensions( + randomized_input, + begin_axis=0, + end_axis=0, + target_dims=(batch_size, self.num_samples) + ) # [B * N, ...] -> [B, N, ...] + randomized_input_np = TensorUtils.to_numpy(randomized_input[random_sample_inds]) + + pre_random_input_np = pre_random_input_np.transpose((0, 2, 3, 1)) # [B, C, H, W] -> [B, H, W, C] + randomized_input_np = randomized_input_np.transpose((0, 1, 3, 4, 2)) # [B, N, C, H, W] -> [B, N, H, W, C] + + visualize_image_randomizer( + pre_random_input_np, + randomized_input_np, + randomizer_name='{}'.format(str(self.__class__.__name__)) + ) + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = header + f"(input_shape={self.input_shape}, noise_mean={self.noise_mean}, noise_std={self.noise_std}, " \ + f"limits={self.limits}, num_samples={self.num_samples})" + return msg diff --git a/aloha-devel/robomimic/models/policy_nets.py b/aloha-devel/robomimic/models/policy_nets.py new file mode 100644 index 0000000000000000000000000000000000000000..8dba1d934cbb6b6a6f2d5c6475d699c48eb2a302 --- /dev/null +++ b/aloha-devel/robomimic/models/policy_nets.py @@ -0,0 +1,1570 @@ +""" +Contains torch Modules for policy networks. These networks take an +observation dictionary as input (and possibly additional conditioning, +such as subgoal or goal dictionaries) and produce action predictions, +samples, or distributions as outputs. Note that actions +are assumed to lie in [-1, 1], and most networks will have a final +tanh activation to help ensure this range. +""" +import textwrap +import numpy as np +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.distributions as D + +import robomimic.utils.tensor_utils as TensorUtils +from robomimic.models.base_nets import Module +from robomimic.models.transformers import GPT_Backbone +from robomimic.models.obs_nets import MIMO_MLP, RNN_MIMO_MLP, MIMO_Transformer, ObservationDecoder +from robomimic.models.vae_nets import VAE +from robomimic.models.distributions import TanhWrappedDistribution + + +class ActorNetwork(MIMO_MLP): + """ + A basic policy network that predicts actions from observations. + Can optionally be goal conditioned on future observations. + """ + def __init__( + self, + obs_shapes, + ac_dim, + mlp_layer_dims, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + obs_shapes (OrderedDict): a dictionary that maps observation keys to + expected shapes for observations. + + ac_dim (int): dimension of action space. + + mlp_layer_dims ([int]): sequence of integers for the MLP hidden layers sizes. + + goal_shapes (OrderedDict): a dictionary that maps observation keys to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-observation key information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + assert isinstance(obs_shapes, OrderedDict) + self.obs_shapes = obs_shapes + self.ac_dim = ac_dim + + # set up different observation groups for @MIMO_MLP + observation_group_shapes = OrderedDict() + observation_group_shapes["obs"] = OrderedDict(self.obs_shapes) + + self._is_goal_conditioned = False + if goal_shapes is not None and len(goal_shapes) > 0: + assert isinstance(goal_shapes, OrderedDict) + self._is_goal_conditioned = True + self.goal_shapes = OrderedDict(goal_shapes) + observation_group_shapes["goal"] = OrderedDict(self.goal_shapes) + else: + self.goal_shapes = OrderedDict() + + output_shapes = self._get_output_shapes() + super(ActorNetwork, self).__init__( + input_obs_group_shapes=observation_group_shapes, + output_shapes=output_shapes, + layer_dims=mlp_layer_dims, + encoder_kwargs=encoder_kwargs, + ) + + def _get_output_shapes(self): + """ + Allow subclasses to re-define outputs from @MIMO_MLP, since we won't + always directly predict actions, but may instead predict the parameters + of a action distribution. + """ + return OrderedDict(action=(self.ac_dim,)) + + def output_shape(self, input_shape=None): + return [self.ac_dim] + + def forward(self, obs_dict, goal_dict=None): + actions = super(ActorNetwork, self).forward(obs=obs_dict, goal=goal_dict)["action"] + # apply tanh squashing to ensure actions are in [-1, 1] + return torch.tanh(actions) + + def _to_string(self): + """Info to pretty print.""" + return "action_dim={}".format(self.ac_dim) + + +class PerturbationActorNetwork(ActorNetwork): + """ + An action perturbation network - primarily used in BCQ. + It takes states and actions and returns action perturbations. + """ + def __init__( + self, + obs_shapes, + ac_dim, + mlp_layer_dims, + perturbation_scale=0.05, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + obs_shapes (OrderedDict): a dictionary that maps observation keys to + expected shapes for observations. + + ac_dim (int): dimension of action space. + + mlp_layer_dims ([int]): sequence of integers for the MLP hidden layers sizes. + + perturbation_scale (float): the perturbation network output is always squashed to + lie in +/- @perturbation_scale. The final action output is equal to the original + input action added to the output perturbation (and clipped to lie in [-1, 1]). + + goal_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + self.perturbation_scale = perturbation_scale + + # add in action as a modality + new_obs_shapes = OrderedDict(obs_shapes) + new_obs_shapes["action"] = (ac_dim,) + + # pass to super class to instantiate network + super(PerturbationActorNetwork, self).__init__( + obs_shapes=new_obs_shapes, + ac_dim=ac_dim, + mlp_layer_dims=mlp_layer_dims, + goal_shapes=goal_shapes, + encoder_kwargs=encoder_kwargs, + ) + + def forward(self, obs_dict, acts, goal_dict=None): + """Forward pass through perturbation actor.""" + # add in actions + inputs = dict(obs_dict) + inputs["action"] = acts + perturbations = super(PerturbationActorNetwork, self).forward(inputs, goal_dict) + + # add perturbations from network to original actions, and ensure the new actions lie in [-1, 1] + output_actions = acts + self.perturbation_scale * perturbations + output_actions = output_actions.clamp(-1.0, 1.0) + return output_actions + + def _to_string(self): + """Info to pretty print.""" + return "action_dim={}, perturbation_scale={}".format(self.ac_dim, self.perturbation_scale) + + +class GaussianActorNetwork(ActorNetwork): + """ + Variant of actor network that learns a diagonal unimodal Gaussian distribution + over actions. + """ + def __init__( + self, + obs_shapes, + ac_dim, + mlp_layer_dims, + fixed_std=False, + std_activation="softplus", + init_last_fc_weight=None, + init_std=0.3, + mean_limits=(-9.0, 9.0), + std_limits=(0.007, 7.5), + low_noise_eval=True, + use_tanh=False, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + obs_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for observations. + + ac_dim (int): dimension of action space. + + mlp_layer_dims ([int]): sequence of integers for the MLP hidden layers sizes. + + fixed_std (bool): if True, std is not learned, but kept constant at @init_std + + std_activation (None or str): type of activation to use for std deviation. Options are: + + None: no activation applied (not recommended unless using fixed std) + + `'softplus'`: Only applicable if not using fixed std. Softplus activation applied, after which the + output is scaled by init_std / softplus(0) + + `'exp'`: Only applicable if not using fixed std. Exp applied; this corresponds to network output + as being interpreted as log_std instead of std + + NOTE: In all cases, the final result is clipped to be within @std_limits + + init_last_fc_weight (None or float): if specified, will intialize the final layer network weights to be + uniformly sampled from [-init_weight, init_weight] + + init_std (None or float): approximate initial scaling for standard deviation outputs + from network. If None + + mean_limits (2-array): (min, max) to clamp final mean output by + + std_limits (2-array): (min, max) to clamp final std output by + + low_noise_eval (float): if True, model will output means of Gaussian distribution + at eval time. + + use_tanh (bool): if True, use a tanh-Gaussian distribution + + goal_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + + # parameters specific to Gaussian actor + self.fixed_std = fixed_std + self.init_std = init_std + self.mean_limits = np.array(mean_limits) + self.std_limits = np.array(std_limits) + + # Define activations to use + def softplus_scaled(x): + out = F.softplus(x) + out = out * (self.init_std / F.softplus(torch.zeros(1).to(x.device))) + return out + + self.activations = { + None: lambda x: x, + "softplus": softplus_scaled, + "exp": torch.exp, + } + assert std_activation in self.activations, \ + "std_activation must be one of: {}; instead got: {}".format(self.activations.keys(), std_activation) + self.std_activation = std_activation if not self.fixed_std else None + + self.low_noise_eval = low_noise_eval + self.use_tanh = use_tanh + + super(GaussianActorNetwork, self).__init__( + obs_shapes=obs_shapes, + ac_dim=ac_dim, + mlp_layer_dims=mlp_layer_dims, + goal_shapes=goal_shapes, + encoder_kwargs=encoder_kwargs, + ) + + # If initialization weight was specified, make sure all final layer network weights are specified correctly + if init_last_fc_weight is not None: + with torch.no_grad(): + for name, layer in self.nets["decoder"].nets.items(): + torch.nn.init.uniform_(layer.weight, -init_last_fc_weight, init_last_fc_weight) + torch.nn.init.uniform_(layer.bias, -init_last_fc_weight, init_last_fc_weight) + + def _get_output_shapes(self): + """ + Tells @MIMO_MLP superclass about the output dictionary that should be generated + at the last layer. Network outputs parameters of Gaussian distribution. + """ + return OrderedDict( + mean=(self.ac_dim,), + scale=(self.ac_dim,), + ) + + def forward_train(self, obs_dict, goal_dict=None): + """ + Return full Gaussian distribution, which is useful for computing + quantities necessary at train-time, like log-likelihood, KL + divergence, etc. + + Args: + obs_dict (dict): batch of observations + goal_dict (dict): if not None, batch of goal observations + + Returns: + dist (Distribution): Gaussian distribution + """ + out = MIMO_MLP.forward(self, obs=obs_dict, goal=goal_dict) + mean = out["mean"] + # Use either constant std or learned std depending on setting + scale = out["scale"] if not self.fixed_std else torch.ones_like(mean) * self.init_std + + # Clamp the mean + mean = torch.clamp(mean, min=self.mean_limits[0], max=self.mean_limits[1]) + + # apply tanh squashing to mean if not using tanh-Gaussian to ensure mean is in [-1, 1] + if not self.use_tanh: + mean = torch.tanh(mean) + + # Calculate scale + if self.low_noise_eval and (not self.training): + # override std value so that you always approximately sample the mean + scale = torch.ones_like(mean) * 1e-4 + else: + # Post-process the scale accordingly + scale = self.activations[self.std_activation](scale) + # Clamp the scale + scale = torch.clamp(scale, min=self.std_limits[0], max=self.std_limits[1]) + + + # the Independent call will make it so that `batch_shape` for dist will be equal to batch size + # while `event_shape` will be equal to action dimension - ensuring that log-probability + # computations are summed across the action dimension + dist = D.Normal(loc=mean, scale=scale) + dist = D.Independent(dist, 1) + + if self.use_tanh: + # Wrap distribution with Tanh + dist = TanhWrappedDistribution(base_dist=dist, scale=1.) + + return dist + + def forward(self, obs_dict, goal_dict=None): + """ + Samples actions from the policy distribution. + + Args: + obs_dict (dict): batch of observations + goal_dict (dict): if not None, batch of goal observations + + Returns: + action (torch.Tensor): batch of actions from policy distribution + """ + dist = self.forward_train(obs_dict, goal_dict) + if self.low_noise_eval and (not self.training): + if self.use_tanh: + # # scaling factor lets us output actions like [-1. 1.] and is consistent with the distribution transform + # return (1. + 1e-6) * torch.tanh(dist.base_dist.mean) + return torch.tanh(dist.mean) + return dist.mean + return dist.sample() + + def _to_string(self): + """Info to pretty print.""" + msg = "action_dim={}\nfixed_std={}\nstd_activation={}\ninit_std={}\nmean_limits={}\nstd_limits={}\nlow_noise_eval={}".format( + self.ac_dim, self.fixed_std, self.std_activation, self.init_std, self.mean_limits, self.std_limits, self.low_noise_eval) + return msg + + +class GMMActorNetwork(ActorNetwork): + """ + Variant of actor network that learns a multimodal Gaussian mixture distribution + over actions. + """ + def __init__( + self, + obs_shapes, + ac_dim, + mlp_layer_dims, + num_modes=5, + min_std=0.01, + std_activation="softplus", + low_noise_eval=True, + use_tanh=False, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + obs_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for observations. + + ac_dim (int): dimension of action space. + + mlp_layer_dims ([int]): sequence of integers for the MLP hidden layers sizes. + + num_modes (int): number of GMM modes + + min_std (float): minimum std output from network + + std_activation (None or str): type of activation to use for std deviation. Options are: + + `'softplus'`: Softplus activation applied + + `'exp'`: Exp applied; this corresponds to network output being interpreted as log_std instead of std + + low_noise_eval (float): if True, model will sample from GMM with low std, so that + one of the GMM modes will be sampled (approximately) + + use_tanh (bool): if True, use a tanh-Gaussian distribution + + goal_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + + # parameters specific to GMM actor + self.num_modes = num_modes + self.min_std = min_std + self.low_noise_eval = low_noise_eval + self.use_tanh = use_tanh + + # Define activations to use + self.activations = { + "softplus": F.softplus, + "exp": torch.exp, + } + assert std_activation in self.activations, \ + "std_activation must be one of: {}; instead got: {}".format(self.activations.keys(), std_activation) + self.std_activation = std_activation + + super(GMMActorNetwork, self).__init__( + obs_shapes=obs_shapes, + ac_dim=ac_dim, + mlp_layer_dims=mlp_layer_dims, + goal_shapes=goal_shapes, + encoder_kwargs=encoder_kwargs, + ) + + def _get_output_shapes(self): + """ + Tells @MIMO_MLP superclass about the output dictionary that should be generated + at the last layer. Network outputs parameters of GMM distribution. + """ + return OrderedDict( + mean=(self.num_modes, self.ac_dim), + scale=(self.num_modes, self.ac_dim), + logits=(self.num_modes,), + ) + + def forward_train(self, obs_dict, goal_dict=None): + """ + Return full GMM distribution, which is useful for computing + quantities necessary at train-time, like log-likelihood, KL + divergence, etc. + + Args: + obs_dict (dict): batch of observations + goal_dict (dict): if not None, batch of goal observations + + Returns: + dist (Distribution): GMM distribution + """ + out = MIMO_MLP.forward(self, obs=obs_dict, goal=goal_dict) + means = out["mean"] + scales = out["scale"] + logits = out["logits"] + + # apply tanh squashing to means if not using tanh-GMM to ensure means are in [-1, 1] + if not self.use_tanh: + means = torch.tanh(means) + + # Calculate scale + if self.low_noise_eval and (not self.training): + # low-noise for all Gaussian dists + scales = torch.ones_like(means) * 1e-4 + else: + # post-process the scale accordingly + scales = self.activations[self.std_activation](scales) + self.min_std + + # mixture components - make sure that `batch_shape` for the distribution is equal + # to (batch_size, num_modes) since MixtureSameFamily expects this shape + component_distribution = D.Normal(loc=means, scale=scales) + component_distribution = D.Independent(component_distribution, 1) + + # unnormalized logits to categorical distribution for mixing the modes + mixture_distribution = D.Categorical(logits=logits) + + dist = D.MixtureSameFamily( + mixture_distribution=mixture_distribution, + component_distribution=component_distribution, + ) + + if self.use_tanh: + # Wrap distribution with Tanh + dist = TanhWrappedDistribution(base_dist=dist, scale=1.) + + return dist + + def forward(self, obs_dict, goal_dict=None): + """ + Samples actions from the policy distribution. + + Args: + obs_dict (dict): batch of observations + goal_dict (dict): if not None, batch of goal observations + + Returns: + action (torch.Tensor): batch of actions from policy distribution + """ + dist = self.forward_train(obs_dict, goal_dict) + return dist.sample() + + def _to_string(self): + """Info to pretty print.""" + return "action_dim={}\nnum_modes={}\nmin_std={}\nstd_activation={}\nlow_noise_eval={}".format( + self.ac_dim, self.num_modes, self.min_std, self.std_activation, self.low_noise_eval) + + +class RNNActorNetwork(RNN_MIMO_MLP): + """ + An RNN policy network that predicts actions from observations. + """ + def __init__( + self, + obs_shapes, + ac_dim, + mlp_layer_dims, + rnn_hidden_dim, + rnn_num_layers, + rnn_type="LSTM", # [LSTM, GRU] + rnn_kwargs=None, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + obs_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for observations. + + ac_dim (int): dimension of action space. + + mlp_layer_dims ([int]): sequence of integers for the MLP hidden layers sizes. + + rnn_hidden_dim (int): RNN hidden dimension + + rnn_num_layers (int): number of RNN layers + + rnn_type (str): [LSTM, GRU] + + rnn_kwargs (dict): kwargs for the torch.nn.LSTM / GRU + + goal_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + self.ac_dim = ac_dim + + assert isinstance(obs_shapes, OrderedDict) + self.obs_shapes = obs_shapes + + # set up different observation groups for @RNN_MIMO_MLP + observation_group_shapes = OrderedDict() + observation_group_shapes["obs"] = OrderedDict(self.obs_shapes) + + self._is_goal_conditioned = False + if goal_shapes is not None and len(goal_shapes) > 0: + assert isinstance(goal_shapes, OrderedDict) + self._is_goal_conditioned = True + self.goal_shapes = OrderedDict(goal_shapes) + observation_group_shapes["goal"] = OrderedDict(self.goal_shapes) + else: + self.goal_shapes = OrderedDict() + + output_shapes = self._get_output_shapes() + super(RNNActorNetwork, self).__init__( + input_obs_group_shapes=observation_group_shapes, + output_shapes=output_shapes, + mlp_layer_dims=mlp_layer_dims, + mlp_activation=nn.ReLU, + mlp_layer_func=nn.Linear, + rnn_hidden_dim=rnn_hidden_dim, + rnn_num_layers=rnn_num_layers, + rnn_type=rnn_type, + rnn_kwargs=rnn_kwargs, + per_step=True, + encoder_kwargs=encoder_kwargs, + ) + + def _get_output_shapes(self): + """ + Allow subclasses to re-define outputs from @RNN_MIMO_MLP, since we won't + always directly predict actions, but may instead predict the parameters + of a action distribution. + """ + return OrderedDict(action=(self.ac_dim,)) + + def output_shape(self, input_shape): + # note: @input_shape should be dictionary (key: mod) + # infers temporal dimension from input shape + mod = list(self.obs_shapes.keys())[0] + T = input_shape[mod][0] + TensorUtils.assert_size_at_dim(input_shape, size=T, dim=0, + msg="RNNActorNetwork: input_shape inconsistent in temporal dimension") + return [T, self.ac_dim] + + def forward(self, obs_dict, goal_dict=None, rnn_init_state=None, return_state=False): + """ + Forward a sequence of inputs through the RNN and the per-step network. + + Args: + obs_dict (dict): batch of observations - each tensor in the dictionary + should have leading dimensions batch and time [B, T, ...] + goal_dict (dict): if not None, batch of goal observations + rnn_init_state: rnn hidden state, initialize to zero state if set to None + return_state (bool): whether to return hidden state + + Returns: + actions (torch.Tensor): predicted action sequence + rnn_state: return rnn state at the end if return_state is set to True + """ + if self._is_goal_conditioned: + assert goal_dict is not None + # repeat the goal observation in time to match dimension with obs_dict + mod = list(obs_dict.keys())[0] + goal_dict = TensorUtils.unsqueeze_expand_at(goal_dict, size=obs_dict[mod].shape[1], dim=1) + + outputs = super(RNNActorNetwork, self).forward( + obs=obs_dict, goal=goal_dict, rnn_init_state=rnn_init_state, return_state=return_state) + + if return_state: + actions, state = outputs + else: + actions = outputs + state = None + + # apply tanh squashing to ensure actions are in [-1, 1] + actions = torch.tanh(actions["action"]) + + if return_state: + return actions, state + else: + return actions + + def forward_step(self, obs_dict, goal_dict=None, rnn_state=None): + """ + Unroll RNN over single timestep to get actions. + + Args: + obs_dict (dict): batch of observations. Should not contain + time dimension. + goal_dict (dict): if not None, batch of goal observations + rnn_state: rnn hidden state, initialize to zero state if set to None + + Returns: + actions (torch.Tensor): batch of actions - does not contain time dimension + state: updated rnn state + """ + obs_dict = TensorUtils.to_sequence(obs_dict) + action, state = self.forward( + obs_dict, goal_dict, rnn_init_state=rnn_state, return_state=True) + return action[:, 0], state + + def _to_string(self): + """Info to pretty print.""" + return "action_dim={}".format(self.ac_dim) + + +class RNNGMMActorNetwork(RNNActorNetwork): + """ + An RNN GMM policy network that predicts sequences of action distributions from observation sequences. + """ + def __init__( + self, + obs_shapes, + ac_dim, + mlp_layer_dims, + rnn_hidden_dim, + rnn_num_layers, + rnn_type="LSTM", # [LSTM, GRU] + rnn_kwargs=None, + num_modes=5, + min_std=0.01, + std_activation="softplus", + low_noise_eval=True, + use_tanh=False, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + + rnn_hidden_dim (int): RNN hidden dimension + + rnn_num_layers (int): number of RNN layers + + rnn_type (str): [LSTM, GRU] + + rnn_kwargs (dict): kwargs for the torch.nn.LSTM / GRU + + num_modes (int): number of GMM modes + + min_std (float): minimum std output from network + + std_activation (None or str): type of activation to use for std deviation. Options are: + + `'softplus'`: Softplus activation applied + + `'exp'`: Exp applied; this corresponds to network output being interpreted as log_std instead of std + + low_noise_eval (float): if True, model will sample from GMM with low std, so that + one of the GMM modes will be sampled (approximately) + + use_tanh (bool): if True, use a tanh-Gaussian distribution + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + + # parameters specific to GMM actor + self.num_modes = num_modes + self.min_std = min_std + self.low_noise_eval = low_noise_eval + self.use_tanh = use_tanh + + # Define activations to use + self.activations = { + "softplus": F.softplus, + "exp": torch.exp, + } + assert std_activation in self.activations, \ + "std_activation must be one of: {}; instead got: {}".format(self.activations.keys(), std_activation) + self.std_activation = std_activation + + super(RNNGMMActorNetwork, self).__init__( + obs_shapes=obs_shapes, + ac_dim=ac_dim, + mlp_layer_dims=mlp_layer_dims, + rnn_hidden_dim=rnn_hidden_dim, + rnn_num_layers=rnn_num_layers, + rnn_type=rnn_type, + rnn_kwargs=rnn_kwargs, + goal_shapes=goal_shapes, + encoder_kwargs=encoder_kwargs, + ) + + def _get_output_shapes(self): + """ + Tells @MIMO_MLP superclass about the output dictionary that should be generated + at the last layer. Network outputs parameters of GMM distribution. + """ + return OrderedDict( + mean=(self.num_modes, self.ac_dim), + scale=(self.num_modes, self.ac_dim), + logits=(self.num_modes,), + ) + + def forward_train(self, obs_dict, goal_dict=None, rnn_init_state=None, return_state=False): + """ + Return full GMM distribution, which is useful for computing + quantities necessary at train-time, like log-likelihood, KL + divergence, etc. + + Args: + obs_dict (dict): batch of observations + goal_dict (dict): if not None, batch of goal observations + rnn_init_state: rnn hidden state, initialize to zero state if set to None + return_state (bool): whether to return hidden state + + Returns: + dists (Distribution): sequence of GMM distributions over the timesteps + rnn_state: return rnn state at the end if return_state is set to True + """ + if self._is_goal_conditioned: + assert goal_dict is not None + # repeat the goal observation in time to match dimension with obs_dict + mod = list(obs_dict.keys())[0] + goal_dict = TensorUtils.unsqueeze_expand_at(goal_dict, size=obs_dict[mod].shape[1], dim=1) + + outputs = RNN_MIMO_MLP.forward( + self, obs=obs_dict, goal=goal_dict, rnn_init_state=rnn_init_state, return_state=return_state) + + if return_state: + outputs, state = outputs + else: + state = None + + means = outputs["mean"] + scales = outputs["scale"] + logits = outputs["logits"] + + # apply tanh squashing to mean if not using tanh-GMM to ensure means are in [-1, 1] + if not self.use_tanh: + means = torch.tanh(means) + + if self.low_noise_eval and (not self.training): + # low-noise for all Gaussian dists + scales = torch.ones_like(means) * 1e-4 + else: + # post-process the scale accordingly + scales = self.activations[self.std_activation](scales) + self.min_std + + # mixture components - make sure that `batch_shape` for the distribution is equal + # to (batch_size, timesteps, num_modes) since MixtureSameFamily expects this shape + component_distribution = D.Normal(loc=means, scale=scales) + component_distribution = D.Independent(component_distribution, 1) # shift action dim to event shape + + # unnormalized logits to categorical distribution for mixing the modes + mixture_distribution = D.Categorical(logits=logits) + + dists = D.MixtureSameFamily( + mixture_distribution=mixture_distribution, + component_distribution=component_distribution, + ) + + if self.use_tanh: + # Wrap distribution with Tanh + dists = TanhWrappedDistribution(base_dist=dists, scale=1.) + + if return_state: + return dists, state + else: + return dists + + def forward(self, obs_dict, goal_dict=None, rnn_init_state=None, return_state=False): + """ + Samples actions from the policy distribution. + + Args: + obs_dict (dict): batch of observations + goal_dict (dict): if not None, batch of goal observations + + Returns: + action (torch.Tensor): batch of actions from policy distribution + """ + out = self.forward_train(obs_dict=obs_dict, goal_dict=goal_dict, rnn_init_state=rnn_init_state, return_state=return_state) + if return_state: + ad, state = out + return ad.sample(), state + return out.sample() + + def forward_train_step(self, obs_dict, goal_dict=None, rnn_state=None): + """ + Unroll RNN over single timestep to get action GMM distribution, which + is useful for computing quantities necessary at train-time, like + log-likelihood, KL divergence, etc. + + Args: + obs_dict (dict): batch of observations. Should not contain + time dimension. + goal_dict (dict): if not None, batch of goal observations + rnn_state: rnn hidden state, initialize to zero state if set to None + + Returns: + ad (Distribution): GMM action distributions + state: updated rnn state + """ + obs_dict = TensorUtils.to_sequence(obs_dict) + ad, state = self.forward_train( + obs_dict, goal_dict, rnn_init_state=rnn_state, return_state=True) + + # to squeeze time dimension, make another action distribution + assert ad.component_distribution.base_dist.loc.shape[1] == 1 + assert ad.component_distribution.base_dist.scale.shape[1] == 1 + assert ad.mixture_distribution.logits.shape[1] == 1 + component_distribution = D.Normal( + loc=ad.component_distribution.base_dist.loc.squeeze(1), + scale=ad.component_distribution.base_dist.scale.squeeze(1), + ) + component_distribution = D.Independent(component_distribution, 1) + mixture_distribution = D.Categorical(logits=ad.mixture_distribution.logits.squeeze(1)) + ad = D.MixtureSameFamily( + mixture_distribution=mixture_distribution, + component_distribution=component_distribution, + ) + return ad, state + + def forward_step(self, obs_dict, goal_dict=None, rnn_state=None): + """ + Unroll RNN over single timestep to get sampled actions. + + Args: + obs_dict (dict): batch of observations. Should not contain + time dimension. + goal_dict (dict): if not None, batch of goal observations + rnn_state: rnn hidden state, initialize to zero state if set to None + + Returns: + acts (torch.Tensor): batch of actions - does not contain time dimension + state: updated rnn state + """ + obs_dict = TensorUtils.to_sequence(obs_dict) + acts, state = self.forward( + obs_dict, goal_dict, rnn_init_state=rnn_state, return_state=True) + assert acts.shape[1] == 1 + return acts[:, 0], state + + def _to_string(self): + """Info to pretty print.""" + msg = "action_dim={}, std_activation={}, low_noise_eval={}, num_nodes={}, min_std={}".format( + self.ac_dim, self.std_activation, self.low_noise_eval, self.num_modes, self.min_std) + return msg + + +class TransformerActorNetwork(MIMO_Transformer): + """ + An Transformer policy network that predicts actions from observation sequences (assumed to be frame stacked + from previous observations) and possible from previous actions as well (in an autoregressive manner). + """ + def __init__( + self, + obs_shapes, + ac_dim, + transformer_embed_dim, + transformer_num_layers, + transformer_num_heads, + transformer_context_length, + transformer_emb_dropout=0.1, + transformer_attn_dropout=0.1, + transformer_block_output_dropout=0.1, + transformer_sinusoidal_embedding=False, + transformer_activation="gelu", + transformer_nn_parameter_for_timesteps=False, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + + obs_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for observations. + + ac_dim (int): dimension of action space. + + transformer_embed_dim (int): dimension for embeddings used by transformer + + transformer_num_layers (int): number of transformer blocks to stack + + transformer_num_heads (int): number of attention heads for each + transformer block - must divide @transformer_embed_dim evenly. Self-attention is + computed over this many partitions of the embedding dimension separately. + + transformer_context_length (int): expected length of input sequences + + transformer_embedding_dropout (float): dropout probability for embedding inputs in transformer + + transformer_attn_dropout (float): dropout probability for attention outputs for each transformer block + + transformer_block_output_dropout (float): dropout probability for final outputs for each transformer block + + goal_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + self.ac_dim = ac_dim + + assert isinstance(obs_shapes, OrderedDict) + self.obs_shapes = obs_shapes + + self.transformer_nn_parameter_for_timesteps = transformer_nn_parameter_for_timesteps + + # set up different observation groups for @RNN_MIMO_MLP + observation_group_shapes = OrderedDict() + observation_group_shapes["obs"] = OrderedDict(self.obs_shapes) + + self._is_goal_conditioned = False + if goal_shapes is not None and len(goal_shapes) > 0: + assert isinstance(goal_shapes, OrderedDict) + self._is_goal_conditioned = True + self.goal_shapes = OrderedDict(goal_shapes) + observation_group_shapes["goal"] = OrderedDict(self.goal_shapes) + else: + self.goal_shapes = OrderedDict() + + output_shapes = self._get_output_shapes() + super(TransformerActorNetwork, self).__init__( + input_obs_group_shapes=observation_group_shapes, + output_shapes=output_shapes, + transformer_embed_dim=transformer_embed_dim, + transformer_num_layers=transformer_num_layers, + transformer_num_heads=transformer_num_heads, + transformer_context_length=transformer_context_length, + transformer_emb_dropout=transformer_emb_dropout, + transformer_attn_dropout=transformer_attn_dropout, + transformer_block_output_dropout=transformer_block_output_dropout, + transformer_sinusoidal_embedding=transformer_sinusoidal_embedding, + transformer_activation=transformer_activation, + transformer_nn_parameter_for_timesteps=transformer_nn_parameter_for_timesteps, + + encoder_kwargs=encoder_kwargs, + ) + + def _get_output_shapes(self): + """ + Allow subclasses to re-define outputs from @MIMO_Transformer, since we won't + always directly predict actions, but may instead predict the parameters + of a action distribution. + """ + output_shapes = OrderedDict(action=(self.ac_dim,)) + return output_shapes + + def output_shape(self, input_shape): + # note: @input_shape should be dictionary (key: mod) + # infers temporal dimension from input shape + mod = list(self.obs_shapes.keys())[0] + T = input_shape[mod][0] + TensorUtils.assert_size_at_dim(input_shape, size=T, dim=0, + msg="TransformerActorNetwork: input_shape inconsistent in temporal dimension") + return [T, self.ac_dim] + + def forward(self, obs_dict, actions=None, goal_dict=None): + """ + Forward a sequence of inputs through the Transformer. + Args: + obs_dict (dict): batch of observations - each tensor in the dictionary + should have leading dimensions batch and time [B, T, ...] + actions (torch.Tensor): batch of actions of shape [B, T, D] + goal_dict (dict): if not None, batch of goal observations + Returns: + outputs (torch.Tensor or dict): contains predicted action sequence, or dictionary + with predicted action sequence and predicted observation sequences + """ + if self._is_goal_conditioned: + assert goal_dict is not None + # repeat the goal observation in time to match dimension with obs_dict + mod = list(obs_dict.keys())[0] + goal_dict = TensorUtils.unsqueeze_expand_at(goal_dict, size=obs_dict[mod].shape[1], dim=1) + + forward_kwargs = dict(obs=obs_dict, goal=goal_dict) + outputs = super(TransformerActorNetwork, self).forward(**forward_kwargs) + + # apply tanh squashing to ensure actions are in [-1, 1] + outputs["action"] = torch.tanh(outputs["action"]) + + return outputs["action"] # only action sequences + + def _to_string(self): + """Info to pretty print.""" + return "action_dim={}".format(self.ac_dim) + + +class TransformerGMMActorNetwork(TransformerActorNetwork): + """ + A Transformer GMM policy network that predicts sequences of action distributions from observation + sequences (assumed to be frame stacked from previous observations). + """ + def __init__( + self, + obs_shapes, + ac_dim, + transformer_embed_dim, + transformer_num_layers, + transformer_num_heads, + transformer_context_length, + transformer_emb_dropout=0.1, + transformer_attn_dropout=0.1, + transformer_block_output_dropout=0.1, + transformer_sinusoidal_embedding=False, + transformer_activation="gelu", + transformer_nn_parameter_for_timesteps=False, + num_modes=5, + min_std=0.01, + std_activation="softplus", + low_noise_eval=True, + use_tanh=False, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + + obs_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for observations. + + ac_dim (int): dimension of action space. + + transformer_embed_dim (int): dimension for embeddings used by transformer + + transformer_num_layers (int): number of transformer blocks to stack + + transformer_num_heads (int): number of attention heads for each + transformer block - must divide @transformer_embed_dim evenly. Self-attention is + computed over this many partitions of the embedding dimension separately. + + transformer_context_length (int): expected length of input sequences + + transformer_embedding_dropout (float): dropout probability for embedding inputs in transformer + + transformer_attn_dropout (float): dropout probability for attention outputs for each transformer block + + transformer_block_output_dropout (float): dropout probability for final outputs for each transformer block + + num_modes (int): number of GMM modes + + min_std (float): minimum std output from network + + std_activation (None or str): type of activation to use for std deviation. Options are: + + `'softplus'`: Softplus activation applied + + `'exp'`: Exp applied; this corresponds to network output being interpreted as log_std instead of std + + low_noise_eval (float): if True, model will sample from GMM with low std, so that + one of the GMM modes will be sampled (approximately) + + use_tanh (bool): if True, use a tanh-Gaussian distribution + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + + # parameters specific to GMM actor + self.num_modes = num_modes + self.min_std = min_std + self.low_noise_eval = low_noise_eval + self.use_tanh = use_tanh + + # Define activations to use + self.activations = { + "softplus": F.softplus, + "exp": torch.exp, + } + assert std_activation in self.activations, \ + "std_activation must be one of: {}; instead got: {}".format(self.activations.keys(), std_activation) + self.std_activation = std_activation + + super(TransformerGMMActorNetwork, self).__init__( + obs_shapes=obs_shapes, + ac_dim=ac_dim, + transformer_embed_dim=transformer_embed_dim, + transformer_num_layers=transformer_num_layers, + transformer_num_heads=transformer_num_heads, + transformer_context_length=transformer_context_length, + transformer_emb_dropout=transformer_emb_dropout, + transformer_attn_dropout=transformer_attn_dropout, + transformer_block_output_dropout=transformer_block_output_dropout, + transformer_sinusoidal_embedding=transformer_sinusoidal_embedding, + transformer_activation=transformer_activation, + transformer_nn_parameter_for_timesteps=transformer_nn_parameter_for_timesteps, + encoder_kwargs=encoder_kwargs, + goal_shapes=goal_shapes, + ) + + def _get_output_shapes(self): + """ + Tells @MIMO_Transformer superclass about the output dictionary that should be generated + at the last layer. Network outputs parameters of GMM distribution. + """ + return OrderedDict( + mean=(self.num_modes, self.ac_dim), + scale=(self.num_modes, self.ac_dim), + logits=(self.num_modes,), + ) + + def forward_train(self, obs_dict, actions=None, goal_dict=None, low_noise_eval=None): + """ + Return full GMM distribution, which is useful for computing + quantities necessary at train-time, like log-likelihood, KL + divergence, etc. + Args: + obs_dict (dict): batch of observations + actions (torch.Tensor): batch of actions + goal_dict (dict): if not None, batch of goal observations + Returns: + dists (Distribution): sequence of GMM distributions over the timesteps + """ + if self._is_goal_conditioned: + assert goal_dict is not None + # repeat the goal observation in time to match dimension with obs_dict + mod = list(obs_dict.keys())[0] + goal_dict = TensorUtils.unsqueeze_expand_at(goal_dict, size=obs_dict[mod].shape[1], dim=1) + + forward_kwargs = dict(obs=obs_dict, goal=goal_dict) + + outputs = MIMO_Transformer.forward(self, **forward_kwargs) + + means = outputs["mean"] + scales = outputs["scale"] + logits = outputs["logits"] + + # apply tanh squashing to mean if not using tanh-GMM to ensure means are in [-1, 1] + if not self.use_tanh: + means = torch.tanh(means) + + if low_noise_eval is None: + low_noise_eval = self.low_noise_eval + if low_noise_eval and (not self.training): + # low-noise for all Gaussian dists + scales = torch.ones_like(means) * 1e-4 + else: + # post-process the scale accordingly + scales = self.activations[self.std_activation](scales) + self.min_std + + # mixture components - make sure that `batch_shape` for the distribution is equal + # to (batch_size, timesteps, num_modes) since MixtureSameFamily expects this shape + component_distribution = D.Normal(loc=means, scale=scales) + component_distribution = D.Independent(component_distribution, 1) # shift action dim to event shape + + # unnormalized logits to categorical distribution for mixing the modes + mixture_distribution = D.Categorical(logits=logits) + + dists = D.MixtureSameFamily( + mixture_distribution=mixture_distribution, + component_distribution=component_distribution, + ) + + if self.use_tanh: + # Wrap distribution with Tanh + dists = TanhWrappedDistribution(base_dist=dists, scale=1.) + + return dists + + def forward(self, obs_dict, actions=None, goal_dict=None): + """ + Samples actions from the policy distribution. + Args: + obs_dict (dict): batch of observations + actions (torch.Tensor): batch of actions + goal_dict (dict): if not None, batch of goal observations + Returns: + action (torch.Tensor): batch of actions from policy distribution + """ + out = self.forward_train(obs_dict=obs_dict, actions=actions, goal_dict=goal_dict) + return out.sample() + + def _to_string(self): + """Info to pretty print.""" + msg = "action_dim={}, std_activation={}, low_noise_eval={}, num_nodes={}, min_std={}".format( + self.ac_dim, self.std_activation, self.low_noise_eval, self.num_modes, self.min_std) + return msg + + +class VAEActor(Module): + """ + A VAE that models a distribution of actions conditioned on observations. + The VAE prior and decoder are used at test-time as the policy. + """ + def __init__( + self, + obs_shapes, + ac_dim, + encoder_layer_dims, + decoder_layer_dims, + latent_dim, + device, + decoder_is_conditioned=True, + decoder_reconstruction_sum_across_elements=False, + latent_clip=None, + prior_learn=False, + prior_is_conditioned=False, + prior_layer_dims=(), + prior_use_gmm=False, + prior_gmm_num_modes=10, + prior_gmm_learn_weights=False, + prior_use_categorical=False, + prior_categorical_dim=10, + prior_categorical_gumbel_softmax_hard=False, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + obs_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for observations. + + ac_dim (int): dimension of action space. + + goal_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + super(VAEActor, self).__init__() + + self.obs_shapes = obs_shapes + self.ac_dim = ac_dim + action_shapes = OrderedDict(action=(self.ac_dim,)) + + # ensure VAE decoder will squash actions into [-1, 1] + output_squash = ['action'] + output_scales = OrderedDict(action=1.) + + self._vae = VAE( + input_shapes=action_shapes, + output_shapes=action_shapes, + encoder_layer_dims=encoder_layer_dims, + decoder_layer_dims=decoder_layer_dims, + latent_dim=latent_dim, + device=device, + condition_shapes=self.obs_shapes, + decoder_is_conditioned=decoder_is_conditioned, + decoder_reconstruction_sum_across_elements=decoder_reconstruction_sum_across_elements, + latent_clip=latent_clip, + output_squash=output_squash, + output_scales=output_scales, + prior_learn=prior_learn, + prior_is_conditioned=prior_is_conditioned, + prior_layer_dims=prior_layer_dims, + prior_use_gmm=prior_use_gmm, + prior_gmm_num_modes=prior_gmm_num_modes, + prior_gmm_learn_weights=prior_gmm_learn_weights, + prior_use_categorical=prior_use_categorical, + prior_categorical_dim=prior_categorical_dim, + prior_categorical_gumbel_softmax_hard=prior_categorical_gumbel_softmax_hard, + goal_shapes=goal_shapes, + encoder_kwargs=encoder_kwargs, + ) + + def encode(self, actions, obs_dict, goal_dict=None): + """ + Args: + actions (torch.Tensor): a batch of actions + + obs_dict (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to the observation modalities + used for conditioning in either the decoder or the prior (or both). + + goal_dict (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to goal modalities. + + Returns: + posterior params (dict): dictionary with the following keys: + + mean (torch.Tensor): posterior encoder means + + logvar (torch.Tensor): posterior encoder logvars + """ + inputs = OrderedDict(action=actions) + return self._vae.encode(inputs=inputs, conditions=obs_dict, goals=goal_dict) + + def decode(self, obs_dict=None, goal_dict=None, z=None, n=None): + """ + Thin wrapper around @VaeNets.VAE implementation. + + Args: + obs_dict (dict): a dictionary that maps modalities to torch.Tensor + batches. Only needs to be provided if @decoder_is_conditioned + or @z is None (since the prior will require it to generate z). + + goal_dict (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to goal modalities. + + z (torch.Tensor): if provided, these latents are used to generate + reconstructions from the VAE, and the prior is not sampled. + + n (int): this argument is used to specify the number of samples to + generate from the prior. Only required if @z is None - i.e. + sampling takes place + + Returns: + recons (dict): dictionary of reconstructed inputs (this will be a dictionary + with a single "action" key) + """ + return self._vae.decode(conditions=obs_dict, goals=goal_dict, z=z, n=n) + + def sample_prior(self, obs_dict=None, goal_dict=None, n=None): + """ + Thin wrapper around @VaeNets.VAE implementation. + + Args: + n (int): this argument is used to specify the number + of samples to generate from the prior. + + obs_dict (dict): a dictionary that maps modalities to torch.Tensor + batches. Only needs to be provided if @prior_is_conditioned. + + goal_dict (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to goal modalities. + + Returns: + z (torch.Tensor): latents sampled from the prior + """ + return self._vae.sample_prior(n=n, conditions=obs_dict, goals=goal_dict) + + def set_gumbel_temperature(self, temperature): + """ + Used by external algorithms to schedule Gumbel-Softmax temperature, + which is used during reparametrization at train-time. Should only be + used if @prior_use_categorical is True. + """ + self._vae.set_gumbel_temperature(temperature) + + def get_gumbel_temperature(self): + """ + Return current Gumbel-Softmax temperature. Should only be used if + @prior_use_categorical is True. + """ + return self._vae.get_gumbel_temperature() + + def output_shape(self, input_shape=None): + """ + This implementation is required by the Module superclass, but is unused since we + never chain this module to other ones. + """ + return [self.ac_dim] + + def forward_train(self, actions, obs_dict, goal_dict=None, freeze_encoder=False): + """ + A full pass through the VAE network used during training to construct KL + and reconstruction losses. See @VAE class for more info. + + Args: + actions (torch.Tensor): a batch of actions + + obs_dict (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to the observation modalities + used for conditioning in either the decoder or the prior (or both). + + goal_dict (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to goal modalities. + + Returns: + vae_outputs (dict): a dictionary that contains the following outputs. + + encoder_params (dict): parameters for the posterior distribution + from the encoder forward pass + + encoder_z (torch.Tensor): latents sampled from the encoder posterior + + decoder_outputs (dict): action reconstructions from the decoder + + kl_loss (torch.Tensor): KL loss over the batch of data + + reconstruction_loss (torch.Tensor): reconstruction loss over the batch of data + """ + action_inputs = OrderedDict(action=actions) + return self._vae.forward( + inputs=action_inputs, + outputs=action_inputs, + conditions=obs_dict, + goals=goal_dict, + freeze_encoder=freeze_encoder) + + def forward(self, obs_dict, goal_dict=None, z=None): + """ + Samples actions from the policy distribution. + + Args: + obs_dict (dict): batch of observations + goal_dict (dict): if not None, batch of goal observations + z (torch.Tensor): if not None, use the provided batch of latents instead + of sampling from the prior + + Returns: + action (torch.Tensor): batch of actions from policy distribution + """ + n = None + if z is None: + # prior will be sampled - so we must provide number of samples explicitly + mod = list(obs_dict.keys())[0] + n = obs_dict[mod].shape[0] + return self.decode(obs_dict=obs_dict, goal_dict=goal_dict, z=z, n=n)["action"] diff --git a/aloha-devel/robomimic/models/transformers.py b/aloha-devel/robomimic/models/transformers.py new file mode 100644 index 0000000000000000000000000000000000000000..309bff301d02ad561a34021dbea5d370249cef0f --- /dev/null +++ b/aloha-devel/robomimic/models/transformers.py @@ -0,0 +1,426 @@ +""" +Implementation of transformers, mostly based on Andrej's minGPT model. +See https://github.com/karpathy/minGPT/blob/master/mingpt/model.py +for more details. +""" + +import math +import numpy as np + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from robomimic.models.base_nets import Module +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.torch_utils as TorchUtils + +class GEGLU(nn.Module): + """ + References: + Shazeer et al., "GLU Variants Improve Transformer," 2020. + https://arxiv.org/abs/2002.05202 + Implementation: https://github.com/pfnet-research/deep-table/blob/237c8be8a405349ce6ab78075234c60d9bfe60b7/deep_table/nn/layers/activation.py + """ + + def geglu(self, x): + assert x.shape[-1] % 2 == 0 + a, b = x.chunk(2, dim=-1) + return a * F.gelu(b) + + def forward(self, x): + return self.geglu(x) + + +class PositionalEncoding(nn.Module): + """ + Taken from https://pytorch.org/tutorials/beginner/transformer_tutorial.html. + """ + + def __init__(self, embed_dim): + """ + Standard sinusoidal positional encoding scheme in transformers. + + Positional encoding of the k'th position in the sequence is given by: + p(k, 2i) = sin(k/n^(i/d)) + p(k, 2i+1) = sin(k/n^(i/d)) + + n: set to 10K in original Transformer paper + d: the embedding dimension + i: positions along the projected embedding space (ranges from 0 to d/2) + + Args: + embed_dim: The number of dimensions to project the timesteps into. + """ + super().__init__() + self.embed_dim = embed_dim + + def forward(self, x): + """ + Input timestep of shape BxT + """ + position = x + + # computing 1/n^(i/d) in log space and then exponentiating and fixing the shape + div_term = ( + torch.exp( + torch.arange(0, self.embed_dim, 2, device=x.device) + * (-math.log(10000.0) / self.embed_dim) + ) + .unsqueeze(0) + .unsqueeze(0) + .repeat(x.shape[0], x.shape[1], 1) + ) + pe = torch.zeros((x.shape[0], x.shape[1], self.embed_dim), device=x.device) + pe[:, :, 0::2] = torch.sin(position.unsqueeze(-1) * div_term) + pe[:, :, 1::2] = torch.cos(position.unsqueeze(-1) * div_term) + return pe.detach() + + +class CausalSelfAttention(Module): + def __init__( + self, + embed_dim, + num_heads, + context_length, + attn_dropout=0.1, + output_dropout=0.1, + ): + """ + Multi-head masked self-attention layer + projection (MLP layer). + + For normal self-attention (@num_heads = 1), every single input in the sequence is + mapped to a key, query, and value embedding of size @embed_dim. For each input, + its query vector is compared (using dot-product) with all other key vectors in the + sequence, and softmax normalized to compute an attention over all members of the + sequence. This is used to take a linear combination of corresponding value embeddings. + + The @num_heads argument is for multi-head attention, where the self-attention operation above + is performed in parallel over equal size partitions of the @embed_dim, allowing for different + portions of the embedding dimension to model different kinds of attention. The attention + output for each head is concatenated together. + + Finally, we use a causal mask here to ensure that each output only depends on inputs that come + before it. + + Args: + embed_dim (int): dimension of embeddings to use for keys, queries, and values + used in self-attention + + num_heads (int): number of attention heads - must divide @embed_dim evenly. Self-attention is + computed over this many partitions of the embedding dimension separately. + + context_length (int): expected length of input sequences + + attn_dropout (float): dropout probability for attention outputs + + output_dropout (float): dropout probability for final outputs + """ + super(CausalSelfAttention, self).__init__() + + assert ( + embed_dim % num_heads == 0 + ), "num_heads: {} does not divide embed_dim: {} exactly".format(num_heads, embed_dim) + + self.embed_dim = embed_dim + self.num_heads = num_heads + self.context_length = context_length + self.attn_dropout = attn_dropout + self.output_dropout = output_dropout + self.nets = nn.ModuleDict() + + # projection layers for key, query, value, across all attention heads + self.nets["qkv"] = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=False) + + # dropout layers + self.nets["attn_dropout"] = nn.Dropout(self.attn_dropout) + self.nets["output_dropout"] = nn.Dropout(self.output_dropout) + + # output layer + self.nets["output"] = nn.Linear(self.embed_dim, self.embed_dim) + + # causal mask (ensures attention is only over previous inputs) - just a lower triangular matrix of 1s + mask = torch.tril(torch.ones(context_length, context_length)).view( + 1, 1, context_length, context_length + ) + self.register_buffer("mask", mask) + + def forward(self, x): + """ + Forward pass through Self-Attention block. + Input should be shape (B, T, D) where B is batch size, T is seq length (@self.context_length), and + D is input dimension (@self.embed_dim). + """ + + # enforce shape consistency + assert len(x.shape) == 3 + B, T, D = x.shape + assert ( + T <= self.context_length + ), "self-attention module can only handle sequences up to {} in length but got length {}".format( + self.context_length, T + ) + assert D == self.embed_dim + NH = self.num_heads # number of attention heads + DH = D // NH # embed dimension for each attention head + + # compute key, query, and value vectors for each member of sequence, and split across attention heads + qkv = self.nets["qkv"](x) + q, k, v = torch.chunk(qkv, 3, dim=-1) + k = k.view(B, T, NH, DH).transpose(1, 2) # [B, NH, T, DH] + q = q.view(B, T, NH, DH).transpose(1, 2) # [B, NH, T, DH] + v = v.view(B, T, NH, DH).transpose(1, 2) # [B, NH, T, DH] + + # causal self-attention mechanism + + # batched matrix multiplication between queries and keys to get all pair-wise dot-products. + # We broadcast across batch and attention heads and get pair-wise dot-products between all pairs of timesteps + # [B, NH, T, DH] x [B, NH, DH, T] -> [B, NH, T, T] + att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) + + # use mask to replace entries in dot products with negative inf to ensure they don't contribute to softmax, + # then take softmax over last dimension to end up with attention score for each member of sequence. + # Note the use of [:T, :T] - this makes it so we can handle sequences less than @self.context_length in length. + att = att.masked_fill(self.mask[..., :T, :T] == 0, float("-inf")) + att = F.softmax( + att, dim=-1 + ) # shape [B, NH, T, T], last dimension has score over all T for each sequence member + + # dropout on attention + att = self.nets["attn_dropout"](att) + + # take weighted sum of value vectors over whole sequence according to attention, with batched matrix multiplication + # [B, NH, T, T] x [B, NH, T, DH] -> [B, NH, T, DH] + y = att @ v + # reshape [B, NH, T, DH] -> [B, T, NH, DH] -> [B, T, NH * DH] = [B, T, D] + y = y.transpose(1, 2).contiguous().view(B, T, D) + + # pass through output layer + dropout + y = self.nets["output"](y) + y = self.nets["output_dropout"](y) + return y + + def output_shape(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + + # this module doesn't modify the size of the input, it goes from (B, T, D) -> (B, T, D) + return list(input_shape) + + +class SelfAttentionBlock(Module): + """ + A single Transformer Block, that can be chained together repeatedly. + It consists of a @CausalSelfAttention module and a small MLP, along with + layer normalization and residual connections on each input. + """ + + def __init__( + self, + embed_dim, + num_heads, + context_length, + attn_dropout=0.1, + output_dropout=0.1, + activation=nn.GELU(), + ): + """ + Args: + embed_dim (int): dimension of embeddings to use for keys, queries, and values + used in self-attention + + num_heads (int): number of attention heads - must divide @embed_dim evenly. Self-attention is + computed over this many partitions of the embedding dimension separately. + + context_length (int): expected length of input sequences + + attn_dropout (float): dropout probability for attention outputs + + output_dropout (float): dropout probability for final outputs + + activation (str): string denoting the activation function to use in each transformer block + """ + super(SelfAttentionBlock, self).__init__() + + self.embed_dim = embed_dim + self.num_heads = num_heads + self.context_length = context_length + self.attn_dropout = attn_dropout + self.output_dropout = output_dropout + self.nets = nn.ModuleDict() + + # self-attention block + self.nets["attention"] = CausalSelfAttention( + embed_dim=embed_dim, + num_heads=num_heads, + context_length=context_length, + attn_dropout=attn_dropout, + output_dropout=output_dropout, + ) + + if type(activation) == GEGLU: + mult = 2 + else: + mult = 1 + + # small 2-layer MLP + self.nets["mlp"] = nn.Sequential( + nn.Linear(embed_dim, 4 * embed_dim * mult), + activation, + nn.Linear(4 * embed_dim, embed_dim), + nn.Dropout(output_dropout) + ) + + # layer normalization for inputs to self-attention module and MLP + self.nets["ln1"] = nn.LayerNorm(embed_dim) + self.nets["ln2"] = nn.LayerNorm(embed_dim) + + def forward(self, x): + """ + Forward pass - chain self-attention + MLP blocks, with residual connections and layer norms. + """ + x = x + self.nets["attention"](self.nets["ln1"](x)) + x = x + self.nets["mlp"](self.nets["ln2"](x)) + return x + + def output_shape(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + + # this module doesn't modify the size of the input, it goes from (B, T, D) -> (B, T, D) + return list(input_shape) + + +class GPT_Backbone(Module): + """the GPT model, with a context size of block_size""" + + def __init__( + self, + embed_dim, + context_length, + attn_dropout=0.1, + block_output_dropout=0.1, + num_layers=6, + num_heads=8, + activation="gelu", + ): + """ + Args: + embed_dim (int): dimension of embeddings to use for keys, queries, and values + used in self-attention + + context_length (int): expected length of input sequences + + attn_dropout (float): dropout probability for attention outputs for each transformer block + + block_output_dropout (float): dropout probability for final outputs for each transformer block + + num_layers (int): number of transformer blocks to stack + + num_heads (int): number of attention heads - must divide @embed_dim evenly. Self-attention is + computed over this many partitions of the embedding dimension separately. + + activation (str): string denoting the activation function to use in each transformer block + + """ + super(GPT_Backbone, self).__init__() + + self.embed_dim = embed_dim + self.num_layers = num_layers + self.num_heads = num_heads + self.context_length = context_length + self.attn_dropout = attn_dropout + self.block_output_dropout = block_output_dropout + + if activation == "gelu": + self.activation = nn.GELU() + elif activation == "geglu": + self.activation = GEGLU() + + # create networks + self._create_networks() + + # initialize weights + self.apply(self._init_weights) + + print( + "Created {} model with number of parameters: {}".format( + self.__class__.__name__, sum(p.numel() for p in self.parameters()) + ) + ) + + def _create_networks(self): + """ + Helper function to create networks. + """ + self.nets = nn.ModuleDict() + + # transformer - cascaded transformer blocks + self.nets["transformer"] = nn.Sequential( + *[ + SelfAttentionBlock( + embed_dim=self.embed_dim, + num_heads=self.num_heads, + context_length=self.context_length, + attn_dropout=self.attn_dropout, + output_dropout=self.block_output_dropout, + activation=self.activation, + ) + for _ in range(self.num_layers) + ] + ) + + # decoder head + self.nets["output_ln"] = nn.LayerNorm(self.embed_dim) + + def _init_weights(self, module): + """ + Weight initializer. + """ + if isinstance(module, (nn.Linear, nn.Embedding)): + module.weight.data.normal_(mean=0.0, std=0.02) + if isinstance(module, nn.Linear) and module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + + def output_shape(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + + # this module takes inputs (B, T, @self.input_dim) and produces outputs (B, T, @self.output_dim) + return input_shape[:-1] + [self.output_dim] + + def forward(self, inputs): + assert inputs.shape[1:] == (self.context_length, self.embed_dim), inputs.shape + x = self.nets["transformer"](inputs) + transformer_output = self.nets["output_ln"](x) + return transformer_output \ No newline at end of file diff --git a/aloha-devel/robomimic/models/vae_nets.py b/aloha-devel/robomimic/models/vae_nets.py new file mode 100644 index 0000000000000000000000000000000000000000..91b4e7f02352126f17fcc92a0e651080a4e0bee6 --- /dev/null +++ b/aloha-devel/robomimic/models/vae_nets.py @@ -0,0 +1,1386 @@ +""" +Contains an implementation of Variational Autoencoder (VAE) and other +variants, including other priors, and RNN-VAEs. +""" +import textwrap +import numpy as np +from copy import deepcopy +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.distributions as D + +import robomimic.utils.loss_utils as LossUtils +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.torch_utils as TorchUtils +from robomimic.models.base_nets import Module +from robomimic.models.obs_nets import MIMO_MLP + + +def vae_args_from_config(vae_config): + """ + Generate a set of VAE args that are read from the VAE-specific part + of a config (for example see `config.algo.vae` in BCConfig). + """ + vae_args = dict( + encoder_layer_dims=vae_config.encoder_layer_dims, + decoder_layer_dims=vae_config.decoder_layer_dims, + latent_dim=vae_config.latent_dim, + decoder_is_conditioned=vae_config.decoder.is_conditioned, + decoder_reconstruction_sum_across_elements=vae_config.decoder.reconstruction_sum_across_elements, + latent_clip=vae_config.latent_clip, + prior_learn=vae_config.prior.learn, + prior_is_conditioned=vae_config.prior.is_conditioned, + prior_layer_dims=vae_config.prior_layer_dims, + prior_use_gmm=vae_config.prior.use_gmm, + prior_gmm_num_modes=vae_config.prior.gmm_num_modes, + prior_gmm_learn_weights=vae_config.prior.gmm_learn_weights, + prior_use_categorical=vae_config.prior.use_categorical, + prior_categorical_dim=vae_config.prior.categorical_dim, + prior_categorical_gumbel_softmax_hard=vae_config.prior.categorical_gumbel_softmax_hard, + ) + return vae_args + + +class Prior(Module): + """ + Base class for VAE priors. It's basically the same as a @MIMO_MLP network (it + instantiates one) but it supports additional methods such as KL loss computation + and sampling, and also may learn prior parameters as observation-independent + torch Parameters instead of observation-dependent mappings. + """ + def __init__( + self, + param_shapes, + param_obs_dependent, + obs_shapes=None, + mlp_layer_dims=(), + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + param_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for parameters that determine the prior + distribution. + + param_obs_dependent (OrderedDict): a dictionary with boolean + values consistent with @param_shapes which determines whether + to learn parameters as part of the (obs-dependent) network or + directly as learnable parameters. + + obs_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for observations. + + mlp_layer_dims ([int]): sequence of integers for the MLP hidden layer sizes + + goal_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + super(Prior, self).__init__() + + assert isinstance(param_shapes, OrderedDict) and isinstance(param_obs_dependent, OrderedDict) + assert set(param_shapes.keys()) == set(param_obs_dependent.keys()) + self.param_shapes = param_shapes + self.param_obs_dependent = param_obs_dependent + + net_kwargs = dict( + obs_shapes=obs_shapes, + mlp_layer_dims=mlp_layer_dims, + goal_shapes=goal_shapes, + encoder_kwargs=encoder_kwargs, + ) + self._create_layers(net_kwargs) + + def _create_layers(self, net_kwargs): + """ + Create networks and parameters needed by the prior. + """ + self.prior_params = nn.ParameterDict() + + self._is_obs_dependent = False + mlp_output_shapes = OrderedDict() + for pp in self.param_shapes: + if self.param_obs_dependent[pp]: + # prior parameters will be a function of observations using a network + mlp_output_shapes[pp] = self.param_shapes[pp] + else: + # learnable prior parameters independent of observation + param_init = torch.randn(*self.param_shapes[pp]) / np.sqrt(np.prod(self.param_shapes[pp])) + self.prior_params[pp] = torch.nn.Parameter(param_init) + + # only make networks if we have obs-dependent prior parameters + self.prior_module = None + if len(mlp_output_shapes) > 0: + # create @MIMO_MLP that takes obs and goal dicts and returns prior params + self._is_obs_dependent = True + obs_shapes = net_kwargs["obs_shapes"] + goal_shapes = net_kwargs["goal_shapes"] + obs_group_shapes = OrderedDict() + assert isinstance(obs_shapes, OrderedDict) + obs_group_shapes["obs"] = OrderedDict(obs_shapes) + if goal_shapes is not None and len(goal_shapes) > 0: + assert isinstance(goal_shapes, OrderedDict) + obs_group_shapes["goal"] = OrderedDict(goal_shapes) + self.prior_module = MIMO_MLP( + input_obs_group_shapes=obs_group_shapes, + output_shapes=mlp_output_shapes, + layer_dims=net_kwargs["mlp_layer_dims"], + encoder_kwargs=net_kwargs["encoder_kwargs"], + ) + + def sample(self, n, obs_dict=None, goal_dict=None): + """ + Returns a batch of samples from the prior distribution. + + Args: + n (int): this argument is used to specify the number + of samples to generate from the prior. + + obs_dict (dict): inputs according to @obs_shapes. Only needs to be provided + if any prior parameters are obs-dependent. Leading dimension should + be consistent with @n, the number of samples to generate. + + goal_dict (dict): inputs according to @goal_shapes (only if using goal observations) + + Returns: + z (torch.Tensor): batch of sampled latent vectors. + """ + raise NotImplementedError + + def kl_loss(self, posterior_params, z=None, obs_dict=None, goal_dict=None): + """ + Computes sample-based KL divergence loss between the Gaussian distribution + given by @mu, @logvar and the prior distribution. + + Args: + posterior_params (dict): dictionary with keys "mu" and "logvar" corresponding + to torch.Tensor batch of means and log-variances of posterior Gaussian + distribution. + + z (torch.Tensor): samples from the Gaussian distribution parametrized by + @mu and @logvar. May not be needed depending on the prior. + + obs_dict (dict): inputs according to @obs_shapes. Only needs to be provided + if any prior parameters are obs-dependent. + + goal_dict (dict): inputs according to @goal_shapes (only if using goal observations) + + Returns: + kl_loss (torch.Tensor): KL divergence loss + """ + raise NotImplementedError + + def output_shape(self, input_shape=None): + """ + Returns output shape for this module, which is a dictionary instead + of a list since outputs are dictionaries. + """ + if self.prior_module is not None: + return self.prior_module.output_shape(input_shape) + return { k : list(self.param_shapes[k]) for k in self.param_shapes } + + def forward(self, batch_size, obs_dict=None, goal_dict=None): + """ + Computes prior parameters. + + Args: + batch_size (int): batch size - this is needed for parameters that are + not obs-dependent, to make sure the leading dimension is correct + for downstream sampling and loss computation purposes + + obs_dict (dict): inputs according to @obs_shapes. Only needs to be provided + if any prior parameters are obs-dependent. + + goal_dict (dict): inputs according to @goal_shapes (only if using goal observations) + + Returns: + prior_params (dict): dictionary containing prior parameters + """ + prior_params = dict() + if self._is_obs_dependent: + # forward through network for obs-dependent params + prior_params = self.prior_module.forward(obs=obs_dict, goal=goal_dict) + + # return params that do not depend on obs as well + for pp in self.param_shapes: + if not self.param_obs_dependent[pp]: + # ensure leading dimension will be consistent with other params + prior_params[pp] = TensorUtils.expand_at(self.prior_params[pp], size=batch_size, dim=0) + + # ensure leading dimensions are all consistent + TensorUtils.assert_size_at_dim(prior_params, size=batch_size, dim=0, + msg="prior params dim 0 mismatch in forward") + + return prior_params + + +class GaussianPrior(Prior): + """ + A class that holds functionality for learning both unimodal Gaussian priors and + multimodal Gaussian Mixture Model priors for use in VAEs. + """ + def __init__( + self, + latent_dim, + device, + latent_clip=None, + learnable=False, + use_gmm=False, + gmm_num_modes=10, + gmm_learn_weights=False, + obs_shapes=None, + mlp_layer_dims=(), + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + latent_dim (int): size of latent dimension for the prior + + device (torch.Device): where the module should live (i.e. cpu, gpu) + + latent_clip (float): if provided, clip all latents sampled at + test-time in each dimension to (-@latent_clip, @latent_clip) + + learnable (bool): if True, learn the parameters of the prior (as opposed + to a default N(0, 1) prior) + + use_gmm (bool): if True, learn a Gaussian Mixture Model (GMM) + prior instead of a unimodal Gaussian prior. To use this option, + @learnable must be set to True. + + gmm_num_modes (int): number of GMM modes to learn. Only + used if @use_gmm is True. + + gmm_learn_weights (bool): if True, learn the weights of the GMM + model instead of setting them to be uniform across all the modes. + Only used if @use_gmm is True. + + obs_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for observations. If provided, assumes that + the prior should depend on observation inputs, and networks + will be created to output prior parameters. + + mlp_layer_dims ([int]): sequence of integers for the MLP hidden layer sizes + + goal_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + self.device = device + self.latent_dim = latent_dim + self.latent_clip = latent_clip + self.learnable = learnable + + self.use_gmm = use_gmm + if self.use_gmm: + self.num_modes = gmm_num_modes + else: + # unimodal Gaussian prior + self.num_modes = 1 + self.gmm_learn_weights = gmm_learn_weights + + self._input_dependent = (obs_shapes is not None) and (len(obs_shapes) > 0) + + if self._input_dependent: + assert learnable + assert isinstance(obs_shapes, OrderedDict) + + # network will generate mean and logvar + param_shapes = OrderedDict( + mean=(self.num_modes, self.latent_dim,), + logvar=(self.num_modes, self.latent_dim,), + ) + param_obs_dependent = OrderedDict(mean=True, logvar=True) + + if self.use_gmm and self.gmm_learn_weights: + # network generates GMM weights + param_shapes["weight"] = (self.num_modes,) + param_obs_dependent["weight"] = True + else: + # learn obs-indep mean / logvar + param_shapes = OrderedDict( + mean=(1, self.num_modes, self.latent_dim), + logvar=(1, self.num_modes, self.latent_dim), + ) + param_obs_dependent = OrderedDict(mean=False, logvar=False) + + if self.use_gmm and self.gmm_learn_weights: + # learn obs-indep GMM weights + param_shapes["weight"] = (1, self.num_modes) + param_obs_dependent["weight"] = False + + super(GaussianPrior, self).__init__( + param_shapes=param_shapes, + param_obs_dependent=param_obs_dependent, + obs_shapes=obs_shapes, + mlp_layer_dims=mlp_layer_dims, + goal_shapes=goal_shapes, + encoder_kwargs=encoder_kwargs, + ) + + def _create_layers(self, net_kwargs): + """ + Update from superclass to only create parameters / networks if not using + N(0, 1) Gaussian prior. + """ + if self.learnable: + super(GaussianPrior, self)._create_layers(net_kwargs) + + def sample(self, n, obs_dict=None, goal_dict=None): + """ + Returns a batch of samples from the prior distribution. + + Args: + n (int): this argument is used to specify the number + of samples to generate from the prior. + + obs_dict (dict): inputs according to @obs_shapes. Only needs to be provided + if any prior parameters are obs-dependent. Leading dimension should + be consistent with @n, the number of samples to generate. + + goal_dict (dict): inputs according to @goal_shapes (only if using goal observations) + + Returns: + z (torch.Tensor): batch of sampled latent vectors. + """ + + # check consistency between n and obs_dict + if self._input_dependent: + TensorUtils.assert_size_at_dim(obs_dict, size=n, dim=0, + msg="obs dict and n mismatch in @sample") + + if self.learnable: + + # forward to get parameters + out = self.forward(batch_size=n, obs_dict=obs_dict, goal_dict=goal_dict) + prior_means, prior_logvars, prior_logweights = out["means"], out["logvars"], out["logweights"] + + if prior_logweights is not None: + prior_weights = torch.exp(prior_logweights) + + if self.use_gmm: + # learned GMM + + # make uniform weights (in the case that weights were not learned) + if not self.gmm_learn_weights: + prior_weights = torch.ones(n, self.num_modes).to(prior_means.device) / self.num_modes + + # sample modes + gmm_mode_indices = D.Categorical(prior_weights).sample() + + # get GMM centers and sample using reparametrization trick + selected_means = TensorUtils.gather_sequence(prior_means, indices=gmm_mode_indices) + selected_logvars = TensorUtils.gather_sequence(prior_logvars, indices=gmm_mode_indices) + z = TorchUtils.reparameterize(selected_means, selected_logvars) + + else: + # learned unimodal Gaussian - remove mode dim and sample from Gaussian using reparametrization trick + z = TorchUtils.reparameterize(prior_means[:, 0, :], prior_logvars[:, 0, :]) + + else: + # sample from N(0, 1) + z = torch.randn(n, self.latent_dim).float().to(self.device) + + if self.latent_clip is not None: + z = z.clamp(-self.latent_clip, self.latent_clip) + + return z + + def kl_loss(self, posterior_params, z=None, obs_dict=None, goal_dict=None): + """ + Computes sample-based KL divergence loss between the Gaussian distribution + given by @mu, @logvar and the prior distribution. + + Args: + posterior_params (dict): dictionary with keys "mu" and "logvar" corresponding + to torch.Tensor batch of means and log-variances of posterior Gaussian + distribution. + + z (torch.Tensor): samples from the Gaussian distribution parametrized by + @mu and @logvar. Only needed if @self.use_gmm is True. + + obs_dict (dict): inputs according to @obs_shapes. Only needs to be provided + if any prior parameters are obs-dependent. + + goal_dict (dict): inputs according to @goal_shapes (only if using goal observations) + + Returns: + kl_loss (torch.Tensor): KL divergence loss + """ + mu = posterior_params["mean"] + logvar = posterior_params["logvar"] + + if not self.learnable: + # closed-form Gaussian KL from N(0, 1) prior + return LossUtils.KLD_0_1_loss(mu=mu, logvar=logvar) + + # forward to get parameters + out = self.forward(batch_size=mu.shape[0], obs_dict=obs_dict, goal_dict=goal_dict) + prior_means, prior_logvars, prior_logweights = out["means"], out["logvars"], out["logweights"] + + if not self.use_gmm: + # collapse mode dimension and compute Gaussian KL in closed-form + prior_means = prior_means[:, 0, :] + prior_logvars = prior_logvars[:, 0, :] + return LossUtils.KLD_gaussian_loss( + mu_1=mu, + logvar_1=logvar, + mu_2=prior_means, + logvar_2=prior_logvars, + ) + + # GMM KL loss computation + var = torch.exp(logvar.clamp(-8, 30)) # clamp for numerical stability + prior_vars = torch.exp(prior_logvars.clamp(-8, 30)) + kl_loss = LossUtils.log_normal(x=z, m=mu, v=var) \ + - LossUtils.log_normal_mixture(x=z, m=prior_means, v=prior_vars, log_w=prior_logweights) + return kl_loss.mean() + + def forward(self, batch_size, obs_dict=None, goal_dict=None): + """ + Computes means, logvars, and GMM weights (if using GMM and learning weights). + + Args: + batch_size (int): batch size - this is needed for parameters that are + not obs-dependent, to make sure the leading dimension is correct + for downstream sampling and loss computation purposes + + obs_dict (dict): inputs according to @obs_shapes. Only needs to be provided + if any prior parameters are obs-dependent. + + goal_dict (dict): inputs according to @goal_shapes (only if using goal observations) + + Returns: + prior_params (dict): dictionary containing prior parameters + """ + assert self.learnable + prior_params = super(GaussianPrior, self).forward( + batch_size=batch_size, obs_dict=obs_dict, goal_dict=goal_dict) + + if self.use_gmm and self.gmm_learn_weights: + # normalize learned weight outputs to sum to 1 + logweights = F.log_softmax(prior_params["weight"], dim=-1) + else: + logweights = None + assert "weight" not in prior_params + + out = dict(means=prior_params["mean"], logvars=prior_params["logvar"], logweights=logweights) + return out + + def __repr__(self): + """Pretty print network""" + header = '{}'.format(str(self.__class__.__name__)) + msg = '' + indent = ' ' * 4 + msg += textwrap.indent("latent_dim={}\n".format(self.latent_dim), indent) + msg += textwrap.indent("latent_clip={}\n".format(self.latent_clip), indent) + msg += textwrap.indent("learnable={}\n".format(self.learnable), indent) + msg += textwrap.indent("input_dependent={}\n".format(self._input_dependent), indent) + msg += textwrap.indent("use_gmm={}\n".format(self.use_gmm), indent) + if self.use_gmm: + msg += textwrap.indent("gmm_num_nodes={}\n".format(self.num_modes), indent) + msg += textwrap.indent("gmm_learn_weights={}\n".format(self.gmm_learn_weights), indent) + if self.learnable: + if self.prior_module is not None: + msg += textwrap.indent("\nprior_module={}\n".format(self.prior_module), indent) + msg += textwrap.indent("prior_params={}\n".format(self.prior_params), indent) + msg = header + '(\n' + msg + ')' + return msg + + +class CategoricalPrior(Prior): + """ + A class that holds functionality for learning categorical priors for use + in VAEs. + """ + def __init__( + self, + latent_dim, + categorical_dim, + device, + learnable=False, + obs_shapes=None, + mlp_layer_dims=(), + goal_shapes=None, + encoder_kwargs=None, + + ): + """ + Args: + latent_dim (int): size of latent dimension for the prior + + categorical_dim (int): size of categorical dimension (number of classes + for each dimension of latent space) + + device (torch.Device): where the module should live (i.e. cpu, gpu) + + learnable (bool): if True, learn the parameters of the prior (as opposed + to a default N(0, 1) prior) + + obs_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for observations. If provided, assumes that + the prior should depend on observation inputs, and networks + will be created to output prior parameters. + + mlp_layer_dims ([int]): sequence of integers for the MLP hidden layer sizes + + goal_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + self.device = device + self.latent_dim = latent_dim + self.categorical_dim = categorical_dim + self.learnable = learnable + + self._input_dependent = (obs_shapes is not None) and (len(obs_shapes) > 0) + + if self._input_dependent: + assert learnable + assert isinstance(obs_shapes, OrderedDict) + + # network will generate logits for categorical distributions + param_shapes = OrderedDict( + logit=(self.latent_dim, self.categorical_dim,) + ) + param_obs_dependent = OrderedDict(logit=True) + else: + # learn obs-indep mean / logvar + param_shapes = OrderedDict( + logit=(1, self.latent_dim, self.categorical_dim), + ) + param_obs_dependent = OrderedDict(logit=False) + + super(CategoricalPrior, self).__init__( + param_shapes=param_shapes, + param_obs_dependent=param_obs_dependent, + obs_shapes=obs_shapes, + mlp_layer_dims=mlp_layer_dims, + goal_shapes=goal_shapes, + encoder_kwargs=encoder_kwargs, + ) + + def _create_layers(self, net_kwargs): + """ + Update from superclass to only create parameters / networks if not using + uniform categorical prior. + """ + if self.learnable: + super(CategoricalPrior, self)._create_layers(net_kwargs) + + def sample(self, n, obs_dict=None, goal_dict=None): + """ + Returns a batch of samples from the prior distribution. + + Args: + n (int): this argument is used to specify the number + of samples to generate from the prior. + + obs_dict (dict): inputs according to @obs_shapes. Only needs to be provided + if any prior parameters are obs-dependent. Leading dimension should + be consistent with @n, the number of samples to generate. + + goal_dict (dict): inputs according to @goal_shapes (only if using goal observations) + + Returns: + z (torch.Tensor): batch of sampled latent vectors. + """ + + # check consistency between n and obs_dict + if self._input_dependent: + TensorUtils.assert_size_at_dim(obs_dict, size=n, dim=0, + msg="obs dict and n mismatch in @sample") + + if self.learnable: + + # forward to get parameters + out = self.forward(batch_size=n, obs_dict=obs_dict, goal_dict=goal_dict) + prior_logits = out["logit"] + + # sample one-hot latents from categorical distribution + dist = D.Categorical(logits=prior_logits) + z = TensorUtils.to_one_hot(dist.sample(), num_class=self.categorical_dim) + + else: + # try to include a categorical sample for each class if possible (ensuring rough uniformity) + if (self.latent_dim == 1) and (self.categorical_dim <= n): + # include samples [0, 1, ..., C - 1] and then repeat until batch is filled + dist_samples = torch.arange(n).remainder(self.categorical_dim).unsqueeze(-1).to(self.device) + else: + # sample one-hot latents from uniform categorical distribution for each latent dimension + probs = torch.ones(n, self.latent_dim, self.categorical_dim).float().to(self.device) + dist_samples = D.Categorical(probs=probs).sample() + z = TensorUtils.to_one_hot(dist_samples, num_class=self.categorical_dim) + + # reshape [B, D, C] to [B, D * C] to be consistent with other priors that return flat latents + z = z.reshape(*z.shape[:-2], -1) + return z + + def kl_loss(self, posterior_params, z=None, obs_dict=None, goal_dict=None): + """ + Computes KL divergence loss between the Categorical distribution + given by the unnormalized logits @logits and the prior distribution. + + Args: + posterior_params (dict): dictionary with key "logits" corresponding + to torch.Tensor batch of unnormalized logits of shape [B, D * C] + that corresponds to the posterior categorical distribution + + z (torch.Tensor): samples from encoder - unused for this prior + + obs_dict (dict): inputs according to @obs_shapes. Only needs to be provided + if any prior parameters are obs-dependent. + + goal_dict (dict): inputs according to @goal_shapes (only if using goal observations) + + Returns: + kl_loss (torch.Tensor): KL divergence loss + """ + logits = posterior_params["logit"].reshape(-1, self.latent_dim, self.categorical_dim) + if not self.learnable: + # prior logits correspond to uniform categorical distribution + prior_logits = torch.zeros_like(logits) + else: + # forward to get parameters + out = self.forward(batch_size=posterior_params["logit"].shape[0], obs_dict=obs_dict, goal_dict=goal_dict) + prior_logits = out["logit"] + + prior_dist = D.Categorical(logits=prior_logits) + posterior_dist = D.Categorical(logits=logits) + + # sum over latent dimensions, but average over batch dimension + kl_loss = D.kl_divergence(posterior_dist, prior_dist) + assert len(kl_loss.shape) == 2 + return kl_loss.sum(-1).mean() + + def forward(self, batch_size, obs_dict=None, goal_dict=None): + """ + Computes prior logits (unnormalized log-probs). + + Args: + batch_size (int): batch size - this is needed for parameters that are + not obs-dependent, to make sure the leading dimension is correct + for downstream sampling and loss computation purposes + + obs_dict (dict): inputs according to @obs_shapes. Only needs to be provided + if any prior parameters are obs-dependent. + + goal_dict (dict): inputs according to @goal_shapes (only if using goal observations) + + Returns: + prior_params (dict): dictionary containing prior parameters + """ + assert self.learnable + return super(CategoricalPrior, self).forward( + batch_size=batch_size, obs_dict=obs_dict, goal_dict=goal_dict) + + def __repr__(self): + """Pretty print network""" + header = '{}'.format(str(self.__class__.__name__)) + msg = '' + indent = ' ' * 4 + msg += textwrap.indent("latent_dim={}\n".format(self.latent_dim), indent) + msg += textwrap.indent("categorical_dim={}\n".format(self.categorical_dim), indent) + msg += textwrap.indent("learnable={}\n".format(self.learnable), indent) + msg += textwrap.indent("input_dependent={}\n".format(self._input_dependent), indent) + if self.learnable: + if self.prior_module is not None: + msg += textwrap.indent("\nprior_module={}\n".format(self.prior_module), indent) + msg += textwrap.indent("prior_params={}\n".format(self.prior_params), indent) + msg = header + '(\n' + msg + ')' + return msg + + +class VAE(torch.nn.Module): + """ + A Variational Autoencoder (VAE), as described in https://arxiv.org/abs/1312.6114. + + Models a distribution p(X) or a conditional distribution p(X | Y), where each + variable can consist of multiple modalities. The target variable X, whose + distribution is modeled, is specified through the @input_shapes argument, + which is a map between modalities (strings) and expected shapes. In this way, + a variable that consists of multiple kinds of data (e.g. image and flat-dimensional) + can be modeled as well. A separate @output_shapes argument is used to specify the + expected reconstructions - this allows for asymmetric reconstruction (for example, + reconstructing low-resolution images). + + This implementation supports learning conditional distributions as well (cVAE). + The conditioning variable Y is specified through the @condition_shapes argument, + which is also a map between modalities (strings) and expected shapes. In this way, + variables with multiple kinds of data (e.g. image and flat-dimensional) can + jointly be conditioned on. By default, the decoder takes the conditioning + variable Y as input. To force the decoder to reconstruct from just the latent, + set @decoder_is_conditioned to False (in this case, the prior must be conditioned). + + The implementation also supports learning expressive priors instead of using + the usual N(0, 1) prior. There are three kinds of priors supported - Gaussian, + Gaussian Mixture Model (GMM), and Categorical. For each prior, the parameters can + be learned as independent parameters, or be learned as functions of the conditioning + variable Y (by setting @prior_is_conditioned). + """ + def __init__( + self, + input_shapes, + output_shapes, + encoder_layer_dims, + decoder_layer_dims, + latent_dim, + device, + condition_shapes=None, + decoder_is_conditioned=True, + decoder_reconstruction_sum_across_elements=False, + latent_clip=None, + output_squash=(), + output_scales=None, + output_ranges=None, + prior_learn=False, + prior_is_conditioned=False, + prior_layer_dims=(), + prior_use_gmm=False, + prior_gmm_num_modes=10, + prior_gmm_learn_weights=False, + prior_use_categorical=False, + prior_categorical_dim=10, + prior_categorical_gumbel_softmax_hard=False, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + input_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for all encoder-specific inputs. This corresponds + to the variable X whose distribution we are learning. + + output_shapes (OrderedDict): a dictionary that maps modality to + expected shape for outputs to reconstruct. Usually, this is + the same as @input_shapes but this argument allows + for asymmetries, such as reconstructing low-resolution + images. + + encoder_layer_dims ([int]): sequence of integers for the encoder hidden + layer sizes. + + decoder_layer_dims ([int]): sequence of integers for the decoder hidden + layer sizes. + + latent_dim (int): dimension of latent space for the VAE + + device (torch.Device): where the module should live (i.e. cpu, gpu) + + condition_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for all conditioning inputs. If this is provided, + a conditional distribution is modeled (cVAE). Conditioning takes + place in the decoder by default, and optionally, the prior. + + decoder_is_conditioned (bool): whether to condition the decoder + on the conditioning variables. True by default. Only used if + @condition_shapes is not empty. + + decoder_reconstruction_sum_across_elements (bool): by default, VAEs + average across modality elements and modalities when computing + reconstruction loss. If this is True, sum across all dimensions + and modalities instead. + + latent_clip (float): if provided, clip all latents sampled at + test-time in each dimension to (-@latent_clip, @latent_clip) + + output_squash ([str]): an iterable of modalities that should be + a subset of @output_shapes. The decoder outputs for these + modalities will be squashed into a symmetric range [-a, a] + by using a tanh layer and then scaling the output with the + corresponding value in the @output_scales dictionary. + + output_scales (dict): a dictionary that maps modality to a + scaling value. Used in conjunction with @output_squash. + + output_ranges (dict): a dictionary of [a, b] specifying the output range. + when output_ranges is specified (not None), output_scales should be None + + prior_learn (bool): if True, the prior distribution parameters + are also learned through the KL-divergence loss (instead + of being constrained to a N(0, 1) Gaussian distribution). + If @prior_is_conditioned is True, a global set of parameters + are learned, otherwise, a prior network that maps between + modalities in @condition_shapes and prior parameters is + learned. By default, a Gaussian prior is learned, unless + @prior_use_gmm is True, in which case a Gaussian Mixture + Model (GMM) prior is learned. + + prior_is_conditioned (bool): whether to condition the prior + on the conditioning variables. False by default. Only used if + @condition_shapes is not empty. If this is set to True, + @prior_learn must be True. + + prior_layer_dims ([int]): sequence of integers for the prior hidden layer + sizes. Only used for learned priors that take condition variables as + input (i.e. when @prior_learn and @prior_is_conditioned are set to True, + and @condition_shapes is not empty). + + prior_use_gmm (bool): if True, learn a Gaussian Mixture Model (GMM) + prior instead of a unimodal Gaussian prior. To use this option, + @prior_learn must be set to True. + + prior_gmm_num_modes (int): number of GMM modes to learn. Only + used if @prior_use_gmm is True. + + prior_gmm_learn_weights (bool): if True, learn the weights of the GMM + model instead of setting them to be uniform across all the modes. + Only used if @prior_use_gmm is True. + + prior_use_categorical (bool): if True, use a categorical prior instead of + a unimodal Gaussian prior. This will also cause the encoder to output + a categorical distribution, and will use the Gumbel-Softmax trick + for reparametrization. + + prior_categorical_dim (int): categorical dimension - each latent sampled + from the prior will be of shape (@latent_dim, @prior_categorical_dim) + and will be "one-hot" in the latter dimension. Only used if + @prior_use_categorical is True. + + prior_categorical_gumbel_softmax_hard (bool): if True, use the "hard" version of + Gumbel Softmax for reparametrization. Only used if @prior_use_categorical is True. + + goal_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for goal observations. Goals are treates as additional + conditioning inputs. They are usually specified separately because + they have duplicate modalities as the conditioning inputs (otherwise + they could just be added to the set of conditioning inputs). + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + super(VAE, self).__init__() + + self.latent_dim = latent_dim + self.latent_clip = latent_clip + self.device = device + + # encoder and decoder input dicts and output shapes dict for reconstruction + assert isinstance(input_shapes, OrderedDict) + assert isinstance(output_shapes, OrderedDict) + self.input_shapes = deepcopy(input_shapes) + self.output_shapes = deepcopy(output_shapes) + + # check for conditioning (cVAE) + self._is_cvae = False + self.condition_shapes = deepcopy(condition_shapes) if condition_shapes is not None else OrderedDict() + if len(self.condition_shapes) > 0: + # this is a cVAE - we learn a conditional distribution p(X | Y) + assert isinstance(self.condition_shapes, OrderedDict) + self._is_cvae = True + self.decoder_is_conditioned = decoder_is_conditioned + self.prior_is_conditioned = prior_is_conditioned + assert self.decoder_is_conditioned or self.prior_is_conditioned, \ + "cVAE must be conditioned in decoder and/or prior" + if self.prior_is_conditioned: + assert prior_learn, "to pass conditioning inputs to prior, prior must be learned" + + # check for goal conditioning + self._is_goal_conditioned = False + self.goal_shapes = deepcopy(goal_shapes) if goal_shapes is not None else OrderedDict() + if len(self.goal_shapes) > 0: + assert self._is_cvae, "to condition VAE on goals, it must be a cVAE" + assert isinstance(self.goal_shapes, OrderedDict) + self._is_goal_conditioned = True + + self.encoder_layer_dims = encoder_layer_dims + self.decoder_layer_dims = decoder_layer_dims + + # determines whether outputs are squashed with tanh and if so, to what scaling + assert not (output_scales is not None and output_ranges is not None) + self.output_squash = output_squash + self.output_scales = output_scales if output_scales is not None else OrderedDict() + self.output_ranges = output_ranges if output_ranges is not None else OrderedDict() + + assert set(self.output_squash) == set(self.output_scales.keys()) + assert set(self.output_squash).issubset(set(self.output_shapes)) + + # decoder settings + self.decoder_reconstruction_sum_across_elements = decoder_reconstruction_sum_across_elements + + # prior parameters + self.prior_learn = prior_learn + self.prior_layer_dims = prior_layer_dims + self.prior_use_gmm = prior_use_gmm + self.prior_gmm_num_modes = prior_gmm_num_modes + self.prior_gmm_learn_weights = prior_gmm_learn_weights + self.prior_use_categorical = prior_use_categorical + self.prior_categorical_dim = prior_categorical_dim + self.prior_categorical_gumbel_softmax_hard = prior_categorical_gumbel_softmax_hard + assert np.sum([self.prior_use_gmm, self.prior_use_categorical]) <= 1 + + # for obs core + self._encoder_kwargs = encoder_kwargs + + if self.prior_use_gmm: + assert self.prior_learn, "GMM must be learned" + + if self.prior_use_categorical: + # initialize temperature for Gumbel-Softmax + self.set_gumbel_temperature(1.0) + + # create encoder, decoder, prior + self._create_layers() + + def _create_layers(self): + """ + Creates the encoder, decoder, and prior networks. + """ + self.nets = nn.ModuleDict() + + # VAE Encoder + self._create_encoder() + + # VAE Decoder + self._create_decoder() + + # VAE Prior. + self._create_prior() + + def _create_encoder(self): + """ + Helper function to create encoder. + """ + + # encoder takes "input" dictionary and possibly "condition" (if cVAE) and "goal" (if goal-conditioned) + encoder_obs_group_shapes = OrderedDict() + encoder_obs_group_shapes["input"] = OrderedDict(self.input_shapes) + if self._is_cvae: + encoder_obs_group_shapes["condition"] = OrderedDict(self.condition_shapes) + if self._is_goal_conditioned: + encoder_obs_group_shapes["goal"] = OrderedDict(self.goal_shapes) + + # encoder outputs posterior distribution parameters + if self.prior_use_categorical: + encoder_output_shapes = OrderedDict( + logit=(self.latent_dim * self.prior_categorical_dim,), + ) + else: + encoder_output_shapes = OrderedDict( + mean=(self.latent_dim,), + logvar=(self.latent_dim,), + ) + + self.nets["encoder"] = MIMO_MLP( + input_obs_group_shapes=encoder_obs_group_shapes, + output_shapes=encoder_output_shapes, + layer_dims=self.encoder_layer_dims, + encoder_kwargs=self._encoder_kwargs, + ) + + def _create_decoder(self): + """ + Helper function to create decoder. + """ + + # decoder takes latent (included as "input" observation group) and possibly "condition" (if cVAE) and "goal" (if goal-conditioned) + decoder_obs_group_shapes = OrderedDict() + latent_shape = (self.latent_dim,) + if self.prior_use_categorical: + latent_shape = (self.latent_dim * self.prior_categorical_dim,) + decoder_obs_group_shapes["input"] = OrderedDict(latent=latent_shape) + if self._is_cvae: + decoder_obs_group_shapes["condition"] = OrderedDict(self.condition_shapes) + if self._is_goal_conditioned: + decoder_obs_group_shapes["goal"] = OrderedDict(self.goal_shapes) + + self.nets["decoder"] = MIMO_MLP( + input_obs_group_shapes=decoder_obs_group_shapes, + output_shapes=self.output_shapes, + layer_dims=self.decoder_layer_dims, + encoder_kwargs=self._encoder_kwargs, + ) + + def _create_prior(self): + """ + Helper function to create prior. + """ + + # prior possibly takes "condition" (if cVAE) and "goal" (if goal-conditioned) + prior_obs_group_shapes = OrderedDict(condition=None, goal=None) + if self._is_cvae and self.prior_is_conditioned: + prior_obs_group_shapes["condition"] = OrderedDict(self.condition_shapes) + if self._is_goal_conditioned: + prior_obs_group_shapes["goal"] = OrderedDict(self.goal_shapes) + + if self.prior_use_categorical: + self.nets["prior"] = CategoricalPrior( + latent_dim=self.latent_dim, + categorical_dim=self.prior_categorical_dim, + device=self.device, + learnable=self.prior_learn, + obs_shapes=prior_obs_group_shapes["condition"], + mlp_layer_dims=self.prior_layer_dims, + goal_shapes=prior_obs_group_shapes["goal"], + encoder_kwargs=self._encoder_kwargs, + ) + else: + self.nets["prior"] = GaussianPrior( + latent_dim=self.latent_dim, + device=self.device, + latent_clip=self.latent_clip, + learnable=self.prior_learn, + use_gmm=self.prior_use_gmm, + gmm_num_modes=self.prior_gmm_num_modes, + gmm_learn_weights=self.prior_gmm_learn_weights, + obs_shapes=prior_obs_group_shapes["condition"], + mlp_layer_dims=self.prior_layer_dims, + goal_shapes=prior_obs_group_shapes["goal"], + encoder_kwargs=self._encoder_kwargs, + ) + + def encode(self, inputs, conditions=None, goals=None): + """ + Args: + inputs (dict): a dictionary that maps input modalities to torch.Tensor + batches. These should correspond to the encoder-only modalities + (i.e. @self.encoder_only_shapes). + + conditions (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to the modalities used for conditioning + in either the decoder or the prior (or both). Only for cVAEs. + + goals (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to goal modalities. Only for cVAEs. + + Returns: + posterior params (dict): dictionary with posterior parameters + """ + return self.nets["encoder"]( + input=inputs, + condition=conditions, + goal=goals, + ) + + def reparameterize(self, posterior_params): + """ + Args: + posterior params (dict): dictionary from encoder forward pass that + parametrizes the encoder distribution + + Returns: + z (torch.Tensor): sampled latents that are also differentiable + """ + if self.prior_use_categorical: + # reshape to [B, D, C] to take softmax across categorical classes + logits = posterior_params["logit"].reshape(-1, self.latent_dim, self.prior_categorical_dim) + z = F.gumbel_softmax( + logits=logits, + tau=self._gumbel_temperature, + hard=self.prior_categorical_gumbel_softmax_hard, + dim=-1, + ) + # reshape to [B, D * C], since downstream networks expect flat latents + return TensorUtils.flatten(z) + + return TorchUtils.reparameterize( + mu=posterior_params["mean"], + logvar=posterior_params["logvar"], + ) + + def decode(self, conditions=None, goals=None, z=None, n=None): + """ + Pass latents through decoder. Latents should be passed in to + this function at train-time for backpropagation, but they + can be left out at test-time. In this case, latents will + be sampled using the VAE prior. + + Args: + conditions (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to the modalities used for conditioning + in either the decoder or the prior (or both). Only for cVAEs. + + goals (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to goal modalities. Only for cVAEs. + + z (torch.Tensor): if provided, these latents are used to generate + reconstructions from the VAE, and the prior is not sampled. + + n (int): this argument is used to specify the number of samples to + generate from the prior. Only required if @z is None - i.e. + sampling takes place + + Returns: + recons (dict): dictionary of reconstructed inputs + """ + + if z is None: + # sample latents from prior distribution + assert n is not None + z = self.sample_prior(n=n, conditions=conditions, goals=goals) + + # decoder takes latents as input, and maybe condition variables + # and goal variables + inputs = dict( + input=dict(latent=z), + condition=conditions, + goal=goals, + ) + + # pass through decoder to reconstruct variables in @self.output_shapes + recons = self.nets["decoder"](**inputs) + + # apply tanh squashing to output modalities + for k in self.output_squash: + recons[k] = self.output_scales[k] * torch.tanh(recons[k]) + + for k, v_range in self.output_ranges.items(): + assert v_range[1] > v_range[0] + recons[k] = torch.sigmoid(recons[k]) * (v_range[1] - v_range[0]) + v_range[0] + return recons + + def sample_prior(self, n, conditions=None, goals=None): + """ + Samples from the prior using the prior parameters. + + Args: + n (int): this argument is used to specify the number + of samples to generate from the prior. + + conditions (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to the modalities used for conditioning + in either the decoder or the prior (or both). Only for cVAEs. + + goals (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to goal modalities. Only for cVAEs. + + Returns: + z (torch.Tensor): sampled latents from the prior + """ + return self.nets["prior"].sample(n=n, obs_dict=conditions, goal_dict=goals) + + def kl_loss(self, posterior_params, encoder_z=None, conditions=None, goals=None): + """ + Computes KL divergence loss given the results of the VAE encoder forward + pass and the conditioning and goal modalities (if the prior is input-dependent). + + Args: + posterior_params (dict): dictionary with keys "mu" and "logvar" corresponding + to torch.Tensor batch of means and log-variances of posterior Gaussian + distribution. This is the output of @self.encode. + + encoder_z (torch.Tensor): samples from the Gaussian distribution parametrized by + @mu and @logvar. Only required if using a GMM prior. + + conditions (dict): inputs according to @self.condition_shapes. Only needs to be provided + if any prior parameters are input-dependent. + + goal_dict (dict): inputs according to @self.goal_shapes (only if using goal observations) + + Returns: + kl_loss (torch.Tensor): VAE KL divergence loss + """ + return self.nets["prior"].kl_loss( + posterior_params=posterior_params, + z=encoder_z, + obs_dict=conditions, + goal_dict=goals, + ) + + def reconstruction_loss(self, reconstructions, targets): + """ + Reconstruction loss. Note that we compute the average per-dimension error + in each modality and then average across all the modalities. + + The beta term for weighting between reconstruction and kl losses will + need to be tuned in practice for each situation (see + https://twitter.com/memotv/status/973323454350090240 for more + discussion). + + Args: + reconstructions (dict): reconstructed inputs, consistent with + @self.output_shapes + targets (dict): reconstruction targets, consistent with + @self.output_shapes + + Returns: + reconstruction_loss (torch.Tensor): VAE reconstruction loss + """ + random_key = list(reconstructions.keys())[0] + batch_size = reconstructions[random_key].shape[0] + num_mods = len(reconstructions.keys()) + + # collect errors per modality, while preserving shapes in @reconstructions + recons_errors = [] + for k in reconstructions: + L2_loss = (reconstructions[k] - targets[k]).pow(2) + recons_errors.append(L2_loss) + + # reduce errors across modalities and dimensions + if self.decoder_reconstruction_sum_across_elements: + # average across batch but sum across modalities and dimensions + loss = sum([x.sum() for x in recons_errors]) + loss /= batch_size + else: + # compute mse loss in each modality and average across modalities + loss = sum([x.mean() for x in recons_errors]) + loss /= num_mods + return loss + + def forward(self, inputs, outputs, conditions=None, goals=None, freeze_encoder=False): + """ + A full pass through the VAE network to construct KL and reconstruction + losses. + + Args: + inputs (dict): a dictionary that maps input modalities to torch.Tensor + batches. These should correspond to the encoder-only modalities + (i.e. @self.encoder_only_shapes). + + outputs (dict): a dictionary that maps output modalities to torch.Tensor + batches. These should correspond to the modalities used for + reconstruction (i.e. @self.output_shapes). + + conditions (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to the modalities used for conditioning + in either the decoder or the prior (or both). Only for cVAEs. + + goals (dict): a dictionary that maps modalities to torch.Tensor + batches. These should correspond to goal modalities. Only for cVAEs. + + freeze_encoder (bool): if True, don't backprop into encoder by detaching + encoder outputs. Useful for doing staged VAE training. + + Returns: + vae_outputs (dict): a dictionary that contains the following outputs. + + encoder_params (dict): parameters for the posterior distribution + from the encoder forward pass + + encoder_z (torch.Tensor): latents sampled from the encoder posterior + + decoder_outputs (dict): reconstructions from the decoder + + kl_loss (torch.Tensor): KL loss over the batch of data + + reconstruction_loss (torch.Tensor): reconstruction loss over the batch of data + """ + + # In the comments below, X = inputs, Y = conditions, and we seek to learn P(X | Y). + # The decoder and prior only have knowledge about Y and try to reconstruct X. + # Notice that when Y is the empty set, this reduces to a normal VAE. + + # mu, logvar <- Enc(X, Y) + posterior_params = self.encode( + inputs=inputs, + conditions=conditions, + goals=goals, + ) + + if freeze_encoder: + posterior_params = TensorUtils.detach(posterior_params) + + # z ~ Enc(z | X, Y) + encoder_z = self.reparameterize(posterior_params) + + # hat(X) = Dec(z, Y) + reconstructions = self.decode( + conditions=conditions, + goals=goals, + z=encoder_z, + ) + + # this will also train prior network z ~ Prior(z | Y) + kl_loss = self.kl_loss( + posterior_params=posterior_params, + encoder_z=encoder_z, + conditions=conditions, + goals=goals, + ) + + reconstruction_loss = self.reconstruction_loss( + reconstructions=reconstructions, + targets=outputs, + ) + + return { + "encoder_params" : posterior_params, + "encoder_z" : encoder_z, + "decoder_outputs" : reconstructions, + "kl_loss" : kl_loss, + "reconstruction_loss" : reconstruction_loss, + } + + def set_gumbel_temperature(self, temperature): + """ + Used by external algorithms to schedule Gumbel-Softmax temperature, + which is used during reparametrization at train-time. Should only + be used if @self.prior_use_categorical is True. + """ + assert self.prior_use_categorical + self._gumbel_temperature = temperature + + def get_gumbel_temperature(self): + """ + Return current Gumbel-Softmax temperature. Should only be used if + @self.prior_use_categorical is True. + """ + assert self.prior_use_categorical + return self._gumbel_temperature diff --git a/aloha-devel/robomimic/scripts/config_gen/act_gen.py b/aloha-devel/robomimic/scripts/config_gen/act_gen.py new file mode 100644 index 0000000000000000000000000000000000000000..8962941d7e32dcf6cb195ea6a1ac53517fcaf58a --- /dev/null +++ b/aloha-devel/robomimic/scripts/config_gen/act_gen.py @@ -0,0 +1,131 @@ +from robomimic.scripts.config_gen.helper import * + +def make_generator_helper(args): + algo_name_short = "act" + generator = get_generator( + algo_name="act", + config_file=os.path.join(base_path, 'robomimic/exps/templates/act.json'), + args=args, + algo_name_short=algo_name_short, + pt=True, + ) + if args.ckpt_mode is None: + args.ckpt_mode = "off" + + + generator.add_param( + key="train.num_epochs", + name="", + group=-1, + values=[1000], + ) + + generator.add_param( + key="train.batch_size", + name="", + group=-1, + values=[64], + ) + + generator.add_param( + key="train.max_grad_norm", + name="", + group=-1, + values=[100.0], + ) + + if args.env == "r2d2": + generator.add_param( + key="train.data", + name="ds", + group=2, + values=[ + [{"path": p} for p in scan_datasets("~/Downloads/example_pen_in_cup", postfix="trajectory_im128.h5")], + ], + value_names=[ + "pen-in-cup", + ], + ) + generator.add_param( + key="train.action_keys", + name="ac_keys", + group=-1, + values=[ + [ + "action/abs_pos", + "action/abs_rot_6d", + "action/gripper_position", + ], + ], + value_names=[ + "abs", + ], + ) + elif args.env == "kitchen": + raise NotImplementedError + elif args.env == "square": + generator.add_param( + key="train.data", + name="ds", + group=2, + values=[ + [ + {"path": "TODO.hdf5"}, # replace with your own path + ], + ], + value_names=[ + "square", + ], + ) + + # update env config to use absolute action control + generator.add_param( + key="experiment.env_meta_update_dict", + name="", + group=-1, + values=[ + {"env_kwargs": {"controller_configs": {"control_delta": False}}} + ], + ) + + generator.add_param( + key="train.action_keys", + name="ac_keys", + group=-1, + values=[ + [ + "action_dict/abs_pos", + "action_dict/abs_rot_6d", + "action_dict/gripper", + # "actions", + ], + ], + value_names=[ + "abs", + ], + ) + + + else: + raise ValueError + + generator.add_param( + key="train.output_dir", + name="", + group=-1, + values=[ + "~/expdata/{env}/{mod}/{algo_name_short}".format( + env=args.env, + mod=args.mod, + algo_name_short=algo_name_short, + ) + ], + ) + + return generator + +if __name__ == "__main__": + parser = get_argparser() + + args = parser.parse_args() + make_generator(args, make_generator_helper) diff --git a/aloha-devel/robomimic/scripts/config_gen/bc_rnn_gen.py b/aloha-devel/robomimic/scripts/config_gen/bc_rnn_gen.py new file mode 100644 index 0000000000000000000000000000000000000000..d687006f76bf48052088e1dfd67428468ce6139d --- /dev/null +++ b/aloha-devel/robomimic/scripts/config_gen/bc_rnn_gen.py @@ -0,0 +1,116 @@ +from robomimic.scripts.config_gen.helper import * + +def make_generator_helper(args): + algo_name_short = "bc_rnn" + + generator = get_generator( + algo_name="bc", + config_file=os.path.join(base_path, 'robomimic/exps/templates/bc.json'), + args=args, + algo_name_short=algo_name_short, + pt=True, + ) + if args.ckpt_mode is None: + args.ckpt_mode = "off" + + if args.env == "r2d2": + raise NotImplementedError + elif args.env == "kitchen": + generator.add_param( + key="train.data", + name="ds", + group=2, + values=[ + [{"path": "~/datasets/kitchen/prior/human_demos/pnp_table_to_cab/bowls/20230816_im84.hdf5", "filter_key": "100_demos"}], + # [{"path": "~/datasets/kitchen/prior/human_demos/pnp_table_to_cab/all/20230806_im84.hdf5", "filter_key": "100_demos"}], + # [{"path": "~/datasets/kitchen/prior/mimicgen/pnp_table_to_cab/viraj_mg_2023-08-10-20-31-14/demo_im84.hdf5", "filter_key": "100_demos"}], + # [{"path": "~/datasets/kitchen/prior/mimicgen/pnp_table_to_cab/viraj_mg_2023-08-10-20-31-14/demo_im84.hdf5", "filter_key": "1000_demos"}], + ], + value_names=[ + "bowls-human-100", + # "human-100", + # "mg-100", + # "mg-1000", + ], + ) + else: + raise ValueError + + # change default settings: rnn, predict 10 steps into future + generator.add_param( + key="algo.rnn.enabled", + name="", + group=-1, + values=[True], + hidename=True, + ) + generator.add_param( + key="train.seq_length", + name="", + group=-1, + values=[10], + hidename=True, + ) + generator.add_param( + key="algo.rnn.horizon", + name="", + group=-1, + values=[10], + hidename=True, + ) + if args.mod == "im": + generator.add_param( + key="algo.rnn.hidden_dim", + name="", + group=-1, + values=[1000], + hidename=True, + ) + + generator.add_param( + key="algo.gmm.enabled", + name="gmm", + group=130801, + values=[True], + ) + generator.add_param( + key="algo.gmm.min_std", + name="mindstd", + group=271314, + values=[ + 0.03, + #0.0001, + ], + hidename=True, + ) + generator.add_param( + key="train.max_grad_norm", + name="maxgradnorm", + group=18371, + values=[ + # None, + 100.0, + ], + hidename=True, + ) + + generator.add_param( + key="train.output_dir", + name="", + group=-1, + values=[ + "~/expdata/{env}/{mod}/{algo_name_short}".format( + env=args.env, + mod=args.mod, + algo_name_short=algo_name_short, + ) + ], + ) + + return generator + +if __name__ == "__main__": + parser = get_argparser() + + args = parser.parse_args() + make_generator(args, make_generator_helper) diff --git a/aloha-devel/robomimic/scripts/config_gen/helper.py b/aloha-devel/robomimic/scripts/config_gen/helper.py new file mode 100644 index 0000000000000000000000000000000000000000..48a3af07c0436f463f1a0b545162cb03967ea456 --- /dev/null +++ b/aloha-devel/robomimic/scripts/config_gen/helper.py @@ -0,0 +1,954 @@ +import argparse +import os +import time +import datetime + +import robomimic +import robomimic.utils.hyperparam_utils as HyperparamUtils + +base_path = os.path.abspath(os.path.join(os.path.dirname(robomimic.__file__), os.pardir)) + +def scan_datasets(folder, postfix=".h5"): + dataset_paths = [] + for root, dirs, files in os.walk(os.path.expanduser(folder)): + for f in files: + if f.endswith(postfix): + dataset_paths.append(os.path.join(root, f)) + return dataset_paths + + +def get_generator(algo_name, config_file, args, algo_name_short=None, pt=False): + if args.wandb_proj_name is None: + strings = [ + algo_name_short if (algo_name_short is not None) else algo_name, + args.name, + args.env, + args.mod, + ] + args.wandb_proj_name = '_'.join([str(s) for s in strings if s is not None]) + + if args.script is not None: + generated_config_dir = os.path.join(os.path.dirname(args.script), "json") + else: + curr_time = datetime.datetime.fromtimestamp(time.time()).strftime('%m-%d-%y-%H-%M-%S') + generated_config_dir=os.path.join( + '~/', 'tmp/autogen_configs/ril', algo_name, args.env, args.mod, args.name, curr_time, "json", + ) + + generator = HyperparamUtils.ConfigGenerator( + base_config_file=config_file, + generated_config_dir=generated_config_dir, + wandb_proj_name=args.wandb_proj_name, + script_file=args.script, + ) + + args.algo_name = algo_name + args.pt = pt + + return generator + + +def set_env_settings(generator, args): + if args.env in ["r2d2"]: + assert args.mod == "im" + generator.add_param( + key="experiment.rollout.enabled", + name="", + group=-1, + values=[ + False + ], + ) + generator.add_param( + key="experiment.save.every_n_epochs", + name="", + group=-1, + values=[50], + ) + generator.add_param( + key="experiment.mse.enabled", + name="", + group=-1, + values=[True], + ), + generator.add_param( + key="experiment.mse.every_n_epochs", + name="", + group=-1, + values=[50], + ), + generator.add_param( + key="experiment.mse.on_save_ckpt", + name="", + group=-1, + values=[True], + ), + generator.add_param( + key="experiment.mse.num_samples", + name="", + group=-1, + values=[20], + ), + generator.add_param( + key="experiment.mse.visualize", + name="", + group=-1, + values=[True], + ), + if "observation.modalities.obs.low_dim" not in generator.parameters: + generator.add_param( + key="observation.modalities.obs.low_dim", + name="", + group=-1, + values=[ + ["robot_state/cartesian_position", "robot_state/gripper_position"] + ], + ) + if "observation.modalities.obs.rgb" not in generator.parameters: + generator.add_param( + key="observation.modalities.obs.rgb", + name="", + group=-1, + values=[ + [ + "camera/image/hand_camera_left_image", + "camera/image/varied_camera_1_left_image", "camera/image/varied_camera_2_left_image" # uncomment to use all 3 cameras + ] + ], + ) + generator.add_param( + key="observation.encoder.rgb.obs_randomizer_class", + name="obsrand", + group=-1, + values=[ + # "CropRandomizer", # crop only + # "ColorRandomizer", # jitter only + ["ColorRandomizer", "CropRandomizer"], # jitter, followed by crop + ], + hidename=True, + ) + generator.add_param( + key="observation.encoder.rgb.obs_randomizer_kwargs", + name="obsrandargs", + group=-1, + values=[ + # {"crop_height": 116, "crop_width": 116, "num_crops": 1, "pos_enc": False}, # crop only + # {}, # jitter only + [{}, {"crop_height": 116, "crop_width": 116, "num_crops": 1, "pos_enc": False}], # jitter, followed by crop + ], + hidename=True, + ) + if ("observation.encoder.rgb.obs_randomizer_kwargs" not in generator.parameters) and \ + ("observation.encoder.rgb.obs_randomizer_kwargs.crop_height" not in generator.parameters): + generator.add_param( + key="observation.encoder.rgb.obs_randomizer_kwargs.crop_height", + name="", + group=-1, + values=[ + 116 + ], + ) + generator.add_param( + key="observation.encoder.rgb.obs_randomizer_kwargs.crop_width", + name="", + group=-1, + values=[ + 116 + ], + ) + # remove spatial softmax by default for r2d2 dataset + generator.add_param( + key="observation.encoder.rgb.core_kwargs.pool_class", + name="", + group=-1, + values=[ + None + ], + ) + generator.add_param( + key="observation.encoder.rgb.core_kwargs.pool_kwargs", + name="", + group=-1, + values=[ + None + ], + ) + + # specify dataset type is r2d2 rather than default robomimic + generator.add_param( + key="train.data_format", + name="", + group=-1, + values=[ + "r2d2" + ], + ) + + # here, we list how each action key should be treated (normalized etc) + generator.add_param( + key="train.action_config", + name="", + group=-1, + values=[ + { + "action/cartesian_position":{ + "normalization": "min_max", + }, + "action/abs_pos":{ + "normalization": "min_max", + }, + "action/abs_rot_6d":{ + "normalization": "min_max", + "format": "rot_6d", + "convert_at_runtime": "rot_euler", + }, + "action/abs_rot_euler":{ + "normalization": "min_max", + "format": "rot_euler", + }, + "action/gripper_position":{ + "normalization": "min_max", + }, + "action/cartesian_velocity":{ + "normalization": None, + }, + "action/rel_pos":{ + "normalization": None, + }, + "action/rel_rot_6d":{ + "format": "rot_6d", + "normalization": None, + "convert_at_runtime": "rot_euler", + }, + "action/rel_rot_euler":{ + "format": "rot_euler", + "normalization": None, + }, + "action/gripper_velocity":{ + "normalization": None, + }, + } + ], + ) + generator.add_param( + key="train.dataset_keys", + name="", + group=-1, + values=[[]], + ) + if "train.action_keys" not in generator.parameters: + generator.add_param( + key="train.action_keys", + name="ac_keys", + group=-1, + values=[ + [ + "action/rel_pos", + "action/rel_rot_euler", + "action/gripper_velocity", + ], + ], + value_names=[ + "rel", + ], + ) + # observation key groups to swap + generator.add_param( + key="train.shuffled_obs_key_groups", + name="", + group=-1, + values=[[[ + ( + "camera/image/varied_camera_1_left_image", + "camera/image/varied_camera_1_right_image", + "camera/extrinsics/varied_camera_1_left", + "camera/extrinsics/varied_camera_1_right", + ), + ( + "camera/image/varied_camera_2_left_image", + "camera/image/varied_camera_2_right_image", + "camera/extrinsics/varied_camera_2_left", + "camera/extrinsics/varied_camera_2_right", + ), + ]]], + ) + elif args.env == "kitchen": + generator.add_param( + key="train.action_config", + name="", + group=-1, + values=[ + { + "actions":{ + "normalization": None, + }, + "action_dict/abs_pos": { + "normalization": "min_max" + }, + "action_dict/abs_rot_axis_angle": { + "normalization": "min_max", + "format": "rot_axis_angle" + }, + "action_dict/abs_rot_6d": { + "normalization": None, + "format": "rot_6d" + }, + "action_dict/rel_pos": { + "normalization": None, + }, + "action_dict/rel_rot_axis_angle": { + "normalization": None, + "format": "rot_axis_angle" + }, + "action_dict/rel_rot_6d": { + "normalization": None, + "format": "rot_6d" + }, + "action_dict/gripper": { + "normalization": None, + }, + "action_dict/base_mode": { + "normalization": None, + } + } + ], + ) + + if args.mod == 'im': + generator.add_param( + key="observation.modalities.obs.low_dim", + name="", + group=-1, + values=[ + ["robot0_eef_pos", + "robot0_eef_quat", + "robot0_base_pos", + "robot0_gripper_qpos"] + ], + ) + generator.add_param( + key="observation.modalities.obs.rgb", + name="", + group=-1, + values=[ + ["robot0_agentview_left_image", + "robot0_agentview_right_image", + "robot0_eye_in_hand_image"] + ], + ) + else: + generator.add_param( + key="observation.modalities.obs.low_dim", + name="", + group=-1, + values=[ + ["robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "robot0_base_pos", + "object", + ] + ], + ) + elif args.env in ['square', 'lift', 'place_close']: + # # set videos off + # args.no_video = True + + generator.add_param( + key="train.action_config", + name="", + group=-1, + values=[ + { + "actions":{ + "normalization": None, + }, + "action_dict/abs_pos": { + "normalization": "min_max" + }, + "action_dict/abs_rot_axis_angle": { + "normalization": "min_max", + "format": "rot_axis_angle" + }, + "action_dict/abs_rot_6d": { + "normalization": None, + "format": "rot_6d" + }, + "action_dict/rel_pos": { + "normalization": None, + }, + "action_dict/rel_rot_axis_angle": { + "normalization": None, + "format": "rot_axis_angle" + }, + "action_dict/rel_rot_6d": { + "normalization": None, + "format": "rot_6d" + }, + "action_dict/gripper": { + "normalization": None, + } + } + ], + ) + + if args.mod == 'im': + generator.add_param( + key="observation.modalities.obs.low_dim", + name="", + group=-1, + values=[ + ["robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos"] + ], + ) + generator.add_param( + key="observation.modalities.obs.rgb", + name="", + group=-1, + values=[ + ["agentview_image", + "robot0_eye_in_hand_image"] + ], + ) + else: + generator.add_param( + key="observation.modalities.obs.low_dim", + name="", + group=-1, + values=[ + ["robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object"] + ], + ) + elif args.env == 'transport': + # set videos off + args.no_video = True + + # TODO: fix 2 robot case + generator.add_param( + key="train.action_config", + name="", + group=-1, + values=[ + { + "actions":{ + "normalization": None, + }, + "action_dict/abs_pos": { + "normalization": "min_max" + }, + "action_dict/abs_rot_axis_angle": { + "normalization": "min_max", + "format": "rot_axis_angle" + }, + "action_dict/abs_rot_6d": { + "normalization": None, + "format": "rot_6d" + }, + "action_dict/rel_pos": { + "normalization": None, + }, + "action_dict/rel_rot_axis_angle": { + "normalization": None, + "format": "rot_axis_angle" + }, + "action_dict/rel_rot_6d": { + "normalization": None, + "format": "rot_6d" + }, + "action_dict/gripper": { + "normalization": None, + } + } + ], + ) + + if args.mod == 'im': + generator.add_param( + key="observation.modalities.obs.low_dim", + name="", + group=-1, + values=[ + ["robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "robot1_eef_pos", + "robot1_eef_quat", + "robot1_gripper_qpos"] + ], + ) + generator.add_param( + key="observation.modalities.obs.rgb", + name="", + group=-1, + values=[ + ["shouldercamera0_image", + "robot0_eye_in_hand_image", + "shouldercamera1_image", + "robot1_eye_in_hand_image"] + ], + ) + else: + generator.add_param( + key="observation.modalities.obs.low_dim", + name="", + group=-1, + values=[ + ["robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "robot1_eef_pos", + "robot1_eef_quat", + "robot1_gripper_qpos", + "object"] + ], + ) + + generator.add_param( + key="experiment.rollout.horizon", + name="", + group=-1, + values=[700], + ) + elif args.env == 'tool_hang': + # set videos off + args.no_video = True + + generator.add_param( + key="train.action_config", + name="", + group=-1, + values=[ + { + "actions":{ + "normalization": None, + }, + "action_dict/abs_pos": { + "normalization": "min_max" + }, + "action_dict/abs_rot_axis_angle": { + "normalization": "min_max", + "format": "rot_axis_angle" + }, + "action_dict/abs_rot_6d": { + "normalization": None, + "format": "rot_6d" + }, + "action_dict/rel_pos": { + "normalization": None, + }, + "action_dict/rel_rot_axis_angle": { + "normalization": None, + "format": "rot_axis_angle" + }, + "action_dict/rel_rot_6d": { + "normalization": None, + "format": "rot_6d" + }, + "action_dict/gripper": { + "normalization": None, + } + } + ], + ) + + if args.mod == 'im': + generator.add_param( + key="observation.modalities.obs.low_dim", + name="", + group=-1, + values=[ + ["robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos"] + ], + ) + generator.add_param( + key="observation.modalities.obs.rgb", + name="", + group=-1, + values=[ + ["sideview_image", + "robot0_eye_in_hand_image"] + ], + ) + generator.add_param( + key="observation.encoder.rgb.obs_randomizer_kwargs.crop_height", + name="", + group=-1, + values=[ + 216 + ], + ) + generator.add_param( + key="observation.encoder.rgb.obs_randomizer_kwargs.crop_width", + name="", + group=-1, + values=[ + 216 + ], + ) + generator.add_param( + key="observation.encoder.rgb2.obs_randomizer_kwargs.crop_height", + name="", + group=-1, + values=[ + 216 + ], + ) + generator.add_param( + key="observation.encoder.rgb2.obs_randomizer_kwargs.crop_width", + name="", + group=-1, + values=[ + 216 + ], + ) + else: + generator.add_param( + key="observation.modalities.obs.low_dim", + name="", + group=-1, + values=[ + ["robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object"] + ], + ) + + generator.add_param( + key="experiment.rollout.horizon", + name="", + group=-1, + values=[700], + ) + else: + raise ValueError + + +def set_mod_settings(generator, args): + if args.mod == 'ld': + if "experiment.save.epochs" not in generator.parameters: + generator.add_param( + key="experiment.save.epochs", + name="", + group=-1, + values=[ + [2000] + ], + ) + elif args.mod == 'im': + if "experiment.save.every_n_epochs" not in generator.parameters: + generator.add_param( + key="experiment.save.every_n_epochs", + name="", + group=-1, + values=[40], + ) + + generator.add_param( + key="experiment.epoch_every_n_steps", + name="", + group=-1, + values=[500], + ) + if "train.num_data_workers" not in generator.parameters: + generator.add_param( + key="train.num_data_workers", + name="", + group=-1, + values=[4], + ) + generator.add_param( + key="train.hdf5_cache_mode", + name="", + group=-1, + values=["low_dim"], + ) + if "train.batch_size" not in generator.parameters: + generator.add_param( + key="train.batch_size", + name="", + group=-1, + values=[16], + ) + if "train.num_epochs" not in generator.parameters: + generator.add_param( + key="train.num_epochs", + name="", + group=-1, + values=[600], + ) + if "experiment.rollout.rate" not in generator.parameters: + generator.add_param( + key="experiment.rollout.rate", + name="", + group=-1, + values=[40], + ) + + +def set_debug_mode(generator, args): + if not args.debug: + return + + generator.add_param( + key="experiment.mse.every_n_epochs", + name="", + group=-1, + values=[2], + value_names=[""], + ) + generator.add_param( + key="experiment.mse.visualize", + name="", + group=-1, + values=[True], + value_names=[""], + ) + generator.add_param( + key="experiment.rollout.n", + name="", + group=-1, + values=[2], + value_names=[""], + ) + generator.add_param( + key="experiment.rollout.horizon", + name="", + group=-1, + values=[30], + value_names=[""], + ) + generator.add_param( + key="experiment.rollout.rate", + name="", + group=-1, + values=[2], + value_names=[""], + ) + generator.add_param( + key="experiment.epoch_every_n_steps", + name="", + group=-1, + values=[2], + value_names=[""], + ) + generator.add_param( + key="experiment.save.every_n_epochs", + name="", + group=-1, + values=[2], + value_names=[""], + ) + generator.add_param( + key="experiment.validation_epoch_every_n_steps", + name="", + group=-1, + values=[2], + value_names=[""], + ) + generator.add_param( + key="train.num_epochs", + name="", + group=-1, + values=[2], + value_names=[""], + ) + if args.name is None: + generator.add_param( + key="experiment.name", + name="", + group=-1, + values=["debug"], + value_names=[""], + ) + generator.add_param( + key="experiment.save.enabled", + name="", + group=-1, + values=[False], + value_names=[""], + ) + generator.add_param( + key="train.hdf5_cache_mode", + name="", + group=-1, + values=["low_dim"], + value_names=[""], + ) + generator.add_param( + key="train.num_data_workers", + name="", + group=-1, + values=[3], + ) + + +def set_output_dir(generator, args): + assert args.name is not None + + vals = generator.parameters["train.output_dir"].values + + for i in range(len(vals)): + vals[i] = os.path.join(vals[i], args.name) + + +def set_wandb_mode(generator, args): + generator.add_param( + key="experiment.logging.log_wandb", + name="", + group=-1, + values=[not args.no_wandb], + ) + + +def set_num_seeds(generator, args): + if args.n_seeds is not None and "train.seed" not in generator.parameters: + generator.add_param( + key="train.seed", + name="seed", + group=-10, + values=[i + 1 for i in range(args.n_seeds)], + prepend=True, + ) + + +def get_argparser(): + parser = argparse.ArgumentParser() + + parser.add_argument( + "--name", + type=str, + ) + + parser.add_argument( + "--env", + type=str, + default='r2d2', + ) + + parser.add_argument( + '--mod', + type=str, + choices=['ld', 'im'], + default='im', + ) + + parser.add_argument( + "--ckpt_mode", + type=str, + choices=["off", "all", "best_only"], + default=None, + ) + + parser.add_argument( + "--script", + type=str, + default=None + ) + + parser.add_argument( + "--wandb_proj_name", + type=str, + default=None + ) + + parser.add_argument( + "--debug", + action="store_true", + ) + + parser.add_argument( + '--no_video', + action='store_true' + ) + + parser.add_argument( + "--tmplog", + action="store_true", + ) + + parser.add_argument( + "--nr", + type=int, + default=-1 + ) + + parser.add_argument( + "--no_wandb", + action="store_true", + ) + + parser.add_argument( + "--n_seeds", + type=int, + default=None + ) + + parser.add_argument( + "--num_cmd_groups", + type=int, + default=None + ) + + return parser + + +def make_generator(args, make_generator_helper): + if args.tmplog or args.debug and args.name is None: + args.name = "debug" + else: + time_str = datetime.datetime.fromtimestamp(time.time()).strftime('%m-%d-') + args.name = time_str + str(args.name) + + if args.debug or args.tmplog: + args.no_wandb = True + + if args.wandb_proj_name is not None: + # prepend data to wandb name + # time_str = datetime.datetime.fromtimestamp(time.time()).strftime('%m-%d-') + # args.wandb_proj_name = time_str + args.wandb_proj_name + pass + + if (args.debug or args.tmplog) and (args.wandb_proj_name is None): + args.wandb_proj_name = 'debug' + + if not args.debug: + assert args.name is not None + + # make config generator + generator = make_generator_helper(args) + + if args.ckpt_mode is None: + if args.pt: + args.ckpt_mode = "all" + else: + args.ckpt_mode = "best_only" + + set_env_settings(generator, args) + set_mod_settings(generator, args) + set_output_dir(generator, args) + set_num_seeds(generator, args) + set_wandb_mode(generator, args) + + # set the debug settings last, to override previous setting changes + set_debug_mode(generator, args) + + """ misc settings """ + generator.add_param( + key="experiment.validate", + name="", + group=-1, + values=[ + False, + ], + ) + + # generate jsons and script + generator.generate(override_base_name=True) diff --git a/aloha-devel/robomimic/scripts/conversion/convert_d4rl.py b/aloha-devel/robomimic/scripts/conversion/convert_d4rl.py new file mode 100644 index 0000000000000000000000000000000000000000..99fc1d93c53709f6f149d43457c582cea77e853c --- /dev/null +++ b/aloha-devel/robomimic/scripts/conversion/convert_d4rl.py @@ -0,0 +1,143 @@ +""" +Helper script to convert D4RL data into an hdf5 compatible with this repository. +Takes a folder path and a D4RL env name. This script downloads the corresponding +raw D4RL dataset into a "d4rl" subfolder, and then makes a converted dataset +in the "d4rl/converted" subfolder. + +This script has been tested on the follwing commits: + + https://github.com/rail-berkeley/d4rl/tree/9b68f31bab6a8546edfb28ff0bd9d5916c62fd1f + https://github.com/rail-berkeley/d4rl/tree/26adf732efafdad864b3df2287e7b778ee4f7f63 + +Args: + env (str): d4rl env name, which specifies the dataset to download and convert + folder (str): specify folder to download raw d4rl datasets and converted d4rl datasets to. + A `d4rl` subfolder will be created in this folder with the raw d4rl dataset, and + a `d4rl/converted` subfolder will be created in this folder with the converted + datasets (if they do not already exist). Defaults to the datasets folder at + the top-level of the repository. + +Example usage: + + # downloads to default path at robomimic/datasets/d4rl + python convert_d4rl.py --env walker2d-medium-expert-v2 + + # download to custom path + python convert_d4rl.py --env walker2d-medium-expert-v2 --folder /path/to/folder +""" + +import os +import h5py +import json +import argparse +import numpy as np + +import gym +import d4rl +import robomimic +from robomimic.envs.env_gym import EnvGym +from robomimic.utils.log_utils import custom_tqdm + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--env", + type=str, + help="d4rl env name, which specifies the dataset to download and convert", + ) + parser.add_argument( + "--folder", + type=str, + default=None, + help="specify folder to download raw d4rl datasets and converted d4rl datasets to.\ + A `d4rl` subfolder will be created in this folder with the raw d4rl dataset, and\ + a `d4rl/converted` subfolder will be created in this folder with the converted\ + datasets (if they do not already exist). Defaults to the datasets folder at\ + the top-level of the repository.", + ) + args = parser.parse_args() + + base_folder = args.folder + if base_folder is None: + base_folder = os.path.join(robomimic.__path__[0], "../datasets") + base_folder = os.path.join(base_folder, "d4rl") + + # get dataset + d4rl.set_dataset_path(base_folder) + env = gym.make(args.env) + ds = env.env.get_dataset() + env.close() + + # env + env = EnvGym(args.env) + + # output file + write_folder = os.path.join(base_folder, "converted") + if not os.path.exists(write_folder): + os.makedirs(write_folder) + output_path = os.path.join(base_folder, "converted", "{}.hdf5".format(args.env.replace("-", "_"))) + f_sars = h5py.File(output_path, "w") + f_sars_grp = f_sars.create_group("data") + + # code to split D4RL data into trajectories + # (modified from https://github.com/aviralkumar2907/d4rl_evaluations/blob/bear_intergrate/bear/examples/bear_hdf5_d4rl.py#L18) + all_obs = ds['observations'] + all_act = ds['actions'] + N = all_obs.shape[0] + + obs = all_obs[:N-1] + actions = all_act[:N-1] + next_obs = all_obs[1:] + rewards = np.squeeze(ds['rewards'][:N-1]) + dones = np.squeeze(ds['terminals'][:N-1]).astype(np.int32) + + assert 'timeouts' in ds + timeouts = ds['timeouts'][:] + + ctr = 0 + total_samples = 0 + num_traj = 0 + traj = dict(obs=[], next_obs=[], actions=[], rewards=[], dones=[]) + + print("\nConverting hdf5...") + for idx in custom_tqdm(range(obs.shape[0])): + + # add transition + traj["obs"].append(obs[idx]) + traj["actions"].append(actions[idx]) + traj["rewards"].append(rewards[idx]) + traj["next_obs"].append(next_obs[idx]) + traj["dones"].append(dones[idx]) + ctr += 1 + + # if hit timeout or done is True, end the current trajectory and start a new trajectory + if timeouts[idx] or dones[idx]: + + # replace next obs with copy of current obs for final timestep, and make sure done is true + traj["next_obs"][-1] = np.array(obs[idx]) + traj["dones"][-1] = 1 + + # store trajectory + ep_data_grp = f_sars_grp.create_group("demo_{}".format(num_traj)) + ep_data_grp.create_dataset("obs/flat", data=np.array(traj["obs"])) + ep_data_grp.create_dataset("next_obs/flat", data=np.array(traj["next_obs"])) + ep_data_grp.create_dataset("actions", data=np.array(traj["actions"])) + ep_data_grp.create_dataset("rewards", data=np.array(traj["rewards"])) + ep_data_grp.create_dataset("dones", data=np.array(traj["dones"])) + ep_data_grp.attrs["num_samples"] = len(traj["actions"]) + total_samples += len(traj["actions"]) + num_traj += 1 + + # reset + ctr = 0 + traj = dict(obs=[], next_obs=[], actions=[], rewards=[], dones=[]) + + print("\nExcluding {} samples at end of file due to no trajectory truncation.".format(len(traj["actions"]))) + print("Wrote {} trajectories to new converted hdf5 at {}\n".format(num_traj, output_path)) + + # metadata + f_sars_grp.attrs["total"] = total_samples + f_sars_grp.attrs["env_args"] = json.dumps(env.serialize(), indent=4) + + f_sars.close() + diff --git a/aloha-devel/robomimic/scripts/conversion/convert_r2d2.py b/aloha-devel/robomimic/scripts/conversion/convert_r2d2.py new file mode 100644 index 0000000000000000000000000000000000000000..b7c7d1a4dcaa063dc1b7971dc0b12a3bf482df82 --- /dev/null +++ b/aloha-devel/robomimic/scripts/conversion/convert_r2d2.py @@ -0,0 +1,291 @@ +""" +Add image information to existing r2d2 hdf5 file +""" +import h5py +import os +import numpy as np +import glob +from tqdm import tqdm +import argparse +import shutil +import torch + +""" +Follow instructions here to setup zed: +https://www.stereolabs.com/docs/installation/linux/ +""" +import pyzed.sl as sl + +import robomimic.utils.torch_utils as TorchUtils + +from r2d2.camera_utils.wrappers.recorded_multi_camera_wrapper import RecordedMultiCameraWrapper +from r2d2.trajectory_utils.trajectory_reader import TrajectoryReader +from r2d2.camera_utils.info import camera_type_to_string_dict + +from r2d2.camera_utils.camera_readers.zed_camera import ZedCamera, standard_params + +def get_cam_instrinsics(svo_path): + """ + utility function to get camera intrinsics + """ + intrinsics = {} + + return intrinsics + +def convert_dataset(path, args): + recording_folderpath = os.path.join(os.path.dirname(path), "recordings", "MP4") + camera_kwargs = dict( + hand_camera=dict(image=True, concatenate_images=False, resolution=(args.imsize, args.imsize), resize_func="cv2"), + varied_camera=dict(image=True, concatenate_images=False, resolution=(args.imsize, args.imsize), resize_func="cv2"), + ) + camera_reader = RecordedMultiCameraWrapper(recording_folderpath, camera_kwargs) + + output_path = os.path.join(os.path.dirname(path), "trajectory_im{}.h5".format(args.imsize)) + # if os.path.exists(output_path): + # # dataset already exists, skip + # f = h5py.File(output_path) + # if "observation/camera/image/hand_camera_image" in f.keys(): + # return + # f.close() + + shutil.copyfile(path, output_path) + f = h5py.File(output_path, "a") + + demo_len = f["action"]["cartesian_position"].shape[0] + + if "camera" not in f["observation"]: + f["observation"].create_group("camera").create_group("image") + image_grp = f["observation/camera/image"] + + """ + Extract camera type and keys. Examples of what they should look like: + camera_type_dict = { + '17225336': 'hand_camera', + '24013089': 'varied_camera', + '25047636': 'varied_camera' + } + CAM_NAME_TO_KEY_MAPPING = { + "hand_camera_left_image": "17225336_left", + "hand_camera_right_image": "17225336_right", + "varied_camera_1_left_image": "24013089_left", + "varied_camera_1_right_image": "24013089_right", + "varied_camera_2_left_image": "25047636_left", + "varied_camera_2_right_image": "25047636_right", + } + """ + + CAM_ID_TO_TYPE = {} + hand_cam_ids = [] + varied_cam_ids = [] + for k in f["observation"]["camera_type"]: + cam_type = camera_type_to_string_dict[f["observation"]["camera_type"][k][0]] + CAM_ID_TO_TYPE[k] = cam_type + if cam_type == "hand_camera": + hand_cam_ids.append(k) + elif cam_type == "varied_camera": + varied_cam_ids.append(k) + else: + raise ValueError + + # sort the camera ids: important to maintain consistency of cams between train and eval! + hand_cam_ids = sorted(hand_cam_ids) + varied_cam_ids = sorted(varied_cam_ids) + + IMAGE_NAME_TO_CAM_KEY_MAPPING = {} + IMAGE_NAME_TO_CAM_KEY_MAPPING["hand_camera_left_image"] = "{}_left".format(hand_cam_ids[0]) + IMAGE_NAME_TO_CAM_KEY_MAPPING["hand_camera_right_image"] = "{}_right".format(hand_cam_ids[0]) + + # set up mapping for varied cameras + for i in range(len(varied_cam_ids)): + for side in ["left", "right"]: + cam_name = "varied_camera_{}_{}_image".format(i+1, side) + cam_key = "{}_{}".format(varied_cam_ids[i], side) + IMAGE_NAME_TO_CAM_KEY_MAPPING[cam_name] = cam_key + + cam_data = {cam_name: [] for cam_name in IMAGE_NAME_TO_CAM_KEY_MAPPING.keys()} + traj_reader = TrajectoryReader(path, read_images=False) + + for index in range(demo_len): + + timestep = traj_reader.read_timestep(index=index) + timestamp_dict = timestep["observation"]["timestamp"]["cameras"] + + timestamp_dict = {} + camera_obs = camera_reader.read_cameras( + index=index, camera_type_dict=CAM_ID_TO_TYPE, timestamp_dict=timestamp_dict + ) + for cam_name in IMAGE_NAME_TO_CAM_KEY_MAPPING.keys(): + if camera_obs is None: + im = np.zeros((args.imsize, args.imsize, 3)) + else: + im_key = IMAGE_NAME_TO_CAM_KEY_MAPPING[cam_name] + im = camera_obs["image"][im_key] + + # perform bgr_to_rgb operation + im = im[:,:,::-1] + + cam_data[cam_name].append(im) + + for cam_name in cam_data.keys(): + cam_data[cam_name] = np.array(cam_data[cam_name]).astype(np.uint8) + if cam_name in image_grp: + del image_grp[cam_name] + image_grp.create_dataset(cam_name, data=cam_data[cam_name], compression="gzip") + + # extract camera extrinsics data + if "extrinsics" not in f["observation/camera"]: + f["observation/camera"].create_group("extrinsics") + extrinsics_grp = f["observation/camera/extrinsics"] + for raw_key in f["observation/camera_extrinsics"].keys(): + cam_key = "_".join(raw_key.split("_")[:2]) + # reverse search for image name + im_name = None + for (k, v) in IMAGE_NAME_TO_CAM_KEY_MAPPING.items(): + if v == cam_key: + im_name = k + break + if im_name is None: # sometimes the raw_key doesn't correspond to any camera we have images for + continue + extr_name = "_".join(im_name.split("_")[:-2] + raw_key.split("_")[1:]) + data = f["observation/camera_extrinsics"][raw_key] + extrinsics_grp.create_dataset(extr_name, data=data) + + svo_path = os.path.join(os.path.dirname(path), "recordings", "SVO") + cam_reader_svo = RecordedMultiCameraWrapper(svo_path, camera_kwargs) + if "intrinsics" not in f["observation/camera"]: + f["observation/camera"].create_group("intrinsics") + intrinsics_grp = f["observation/camera/intrinsics"] + for cam_id, svo_reader in cam_reader_svo.camera_dict.items(): + cam = svo_reader._cam + calib_params = cam.get_camera_information().camera_configuration.calibration_parameters + for (posftix, params)in zip( + ["_left", "_right"], + [calib_params.left_cam, calib_params.right_cam] + ): + # get name to store intrinsics under + cam_key = cam_id + posftix + # reverse search for image name + im_name = None + for (k, v) in IMAGE_NAME_TO_CAM_KEY_MAPPING.items(): + if v == cam_key: + im_name = k + break + if im_name is None: # sometimes the raw_key doesn't correspond to any camera we have images for + continue + intr_name = "_".join(im_name.split("_")[:-1]) + + if intr_name not in intrinsics_grp: + intrinsics_grp.create_group(intr_name) + cam_intr_grp = intrinsics_grp[intr_name] + + # these lines are copied from _process_intrinsics function in svo_reader.py + cam_intrinsics = { + "camera_matrix": np.array([[params.fx, 0, params.cx], [0, params.fy, params.cy], [0, 0, 1]]), + "dist_coeffs": np.array(list(params.disto)), + } + # batchify across trajectory + for k in cam_intrinsics: + data = np.repeat(cam_intrinsics[k][None], demo_len, axis=0) + cam_intr_grp.create_dataset(k, data=data) + + # extract action key data + action_dict_group = f["action"] + for in_ac_key in ["cartesian_position", "cartesian_velocity"]: + in_action = action_dict_group[in_ac_key][:] + in_pos = in_action[:,:3].astype(np.float64) + in_rot = in_action[:,3:6].astype(np.float64) # in euler format + rot_ = torch.from_numpy(in_rot) + rot_6d = TorchUtils.euler_angles_to_rot_6d( + rot_, convention="XYZ", + ) + rot_6d = rot_6d.numpy().astype(np.float64) + + if in_ac_key == "cartesian_position": + prefix = "abs_" + elif in_ac_key == "cartesian_velocity": + prefix = "rel_" + else: + raise ValueError + + this_action_dict = { + prefix + 'pos': in_pos, + prefix + 'rot_euler': in_rot, + prefix + 'rot_6d': rot_6d, + } + for key, data in this_action_dict.items(): + if key in action_dict_group: + del action_dict_group[key] + action_dict_group.create_dataset(key, data=data) + + # ensure all action keys are batched (ie., are not 0-dimensional) + for k in action_dict_group: + if isinstance(action_dict_group[k], h5py.Dataset) and len(action_dict_group[k].shape) == 1: + reshaped_values = np.reshape(action_dict_group[k][:], (-1, 1)) + del action_dict_group[k] + action_dict_group.create_dataset(k, data=reshaped_values) + + # post-processing: remove timesteps where robot movement is disabled + movement_enabled = f["observation/controller_info/movement_enabled"][:] + timesteps_to_remove = np.where(movement_enabled == False)[0] + + if not args.keep_idle_timesteps: + remove_timesteps(f, timesteps_to_remove) + + f.close() + +def remove_timesteps(f, timesteps_to_remove): + total_timesteps = f["action/cartesian_position"].shape[0] + + def remove_timesteps_for_group(g): + for k in g: + if isinstance(g[k], h5py._hl.dataset.Dataset): + if g[k].shape[0] != total_timesteps: + print("skipping {}".format(k)) + continue + new_dataset = np.delete(g[k], timesteps_to_remove, axis=0) + del g[k] + g.create_dataset(k, data=new_dataset) + elif isinstance(g[k], h5py._hl.group.Group): + remove_timesteps_for_group(g[k]) + else: + raise NotImplementedError + + for k in f: + remove_timesteps_for_group(f[k]) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--folder", + type=str, + help="folder containing hdf5's to add camera images to", + default="~/datasets/r2d2/success" + ) + + parser.add_argument( + "--imsize", + type=int, + default=128, + help="image size (w and h)", + ) + + parser.add_argument( + "--keep_idle_timesteps", + action="store_true", + help="override the default behavior of truncating idle timesteps", + ) + + args = parser.parse_args() + + datasets = [] + for root, dirs, files in os.walk(os.path.expanduser(args.folder)): + for f in files: + if f == "trajectory.h5": + datasets.append(os.path.join(root, f)) + + print("converting datasets...") + for d in tqdm(datasets): + d = os.path.expanduser(d) + convert_dataset(d, args) diff --git a/aloha-devel/robomimic/scripts/conversion/set_dataset_attr.py b/aloha-devel/robomimic/scripts/conversion/set_dataset_attr.py new file mode 100644 index 0000000000000000000000000000000000000000..4f148d08f0bc9ef6336c919ee7cf4e639aded8b5 --- /dev/null +++ b/aloha-devel/robomimic/scripts/conversion/set_dataset_attr.py @@ -0,0 +1,98 @@ +""" +Example: +python robomimic/scripts/set_dataset_attr.py --glob 'datasets/**/*_abs.hdf5' --env_args env_kwargs.controller_configs.control_delta=false absolute_actions=true +""" +import argparse +import pathlib +import json +import sys +import tqdm +import h5py + +def update_env_args_dict(env_args_dict: dict, key: tuple, value): + if key is None: + return env_args_dict + elif len(key) == 0: + return env_args_dict + elif len(key) == 1: + env_args_dict[key[0]] = value + return env_args_dict + else: + this_key = key[0] + if this_key not in env_args_dict: + env_args_dict[this_key] = dict() + update_env_args_dict(env_args_dict[this_key], key[1:], value) + return env_args_dict + +def main(): + parser = argparse.ArgumentParser() + + parser.add_argument( + "--glob", + type=str, + required=True + ) + + parser.add_argument( + "--env_args", + type=str, + default=None + ) + + parser.add_argument( + 'attrs', + nargs='*' + ) + + args = parser.parse_args() + + # parse attrs to set + # format: key=value + # values are parsed with json + attrs_dict = dict() + for attr_arg in args.attrs: + key, svalue = attr_arg.split("=") + value = json.loads(svalue) + attrs_dict[key] = value + + # parse env_args update + env_args_key = None + env_args_value = None + if args.env_args is not None: + key, svalue = args.env_args.split('=') + env_args_key = key.split('.') + env_args_value = json.loads(svalue) + + # find files + file_paths = list(pathlib.Path.cwd().glob(args.glob)) + + # confirm with the user + print("Found matching files:") + for f in file_paths: + print(f) + print("Are you sure to modify these files with the following attributes:") + print(json.dumps(attrs_dict, indent=2)) + if env_args_key is not None: + print("env_args."+'.'.join(env_args_key)+'='+str(env_args_value)) + result = input("[y/n]?") + if 'y' not in result: + sys.exit(0) + + # execute + for file_path in tqdm.tqdm(file_paths): + with h5py.File(str(file_path), mode='r+') as file: + # update env_args + if env_args_key is not None: + env_args = file['data'].attrs['env_args'] + env_args_dict = json.loads(env_args) + env_args_dict = update_env_args_dict( + env_args_dict=env_args_dict, + key=env_args_key, value=env_args_value) + env_args = json.dumps(env_args_dict) + file['data'].attrs['env_args'] = env_args + + # update other attrs + file['data'].attrs.update(attrs_dict) + +if __name__ == "__main__": + main() diff --git a/aloha-devel/robomimic/scripts/dataset_states_to_obs_multicore.py b/aloha-devel/robomimic/scripts/dataset_states_to_obs_multicore.py new file mode 100644 index 0000000000000000000000000000000000000000..689b2244fb07c5db39bed08088b107c64b7ba2e0 --- /dev/null +++ b/aloha-devel/robomimic/scripts/dataset_states_to_obs_multicore.py @@ -0,0 +1,808 @@ +""" +Script to extract observations from low-dimensional simulation states in a robosuite dataset. + +Args: + dataset (str): path to input hdf5 dataset + + output_name (str): name of output hdf5 dataset + + n (int): if provided, stop after n trajectories are processed + + shaped (bool): if flag is set, use dense rewards + + camera_names (str or [str]): camera name(s) to use for image observations. + Leave out to not use image observations. + + camera_height (int): height of image observation. + + camera_width (int): width of image observation + + done_mode (int): how to write done signal. If 0, done is 1 whenever s' is a success state. + If 1, done is 1 at the end of each trajectory. If 2, both. + + copy_rewards (bool): if provided, copy rewards from source file instead of inferring them + + copy_dones (bool): if provided, copy dones from source file instead of inferring them + +Example usage: + + # extract low-dimensional observations + python dataset_states_to_obs.py --dataset /path/to/demo.hdf5 --output_name low_dim.hdf5 --done_mode 2 + + # extract 84x84 image observations + python dataset_states_to_obs.py --dataset /path/to/demo.hdf5 --output_name image.hdf5 \ + --done_mode 2 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 + + # (space saving option) extract 84x84 image observations with compression and without + # extracting next obs (not needed for pure imitation learning algos) + python dataset_states_to_obs.py --dataset /path/to/demo.hdf5 --output_name image.hdf5 \ + --done_mode 2 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 \ + --compress --exclude-next-obs + + # use dense rewards, and only annotate the end of trajectories with done signal + python dataset_states_to_obs.py --dataset /path/to/demo.hdf5 --output_name image_dense_done_1.hdf5 \ + --done_mode 1 --dense --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 +""" +import os +import json +import h5py +import argparse +import numpy as np +from copy import deepcopy +import multiprocessing +import queue +import time + +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.file_utils as FileUtils +import robomimic.utils.env_utils as EnvUtils +from robomimic.envs.env_base import EnvBase + + +""" + These methods: extract_datagen_info_from_trajectory and extract_datagen_info_from_trajectory_real_robot + are copied over from mimicgen/dataset_states_to_args as importing that file caused environment creation issues +""" +def extract_datagen_info_from_trajectory( + env, + initial_state, + states, + actions, +): + """ + Helper function to extract observations, rewards, and dones along a trajectory using + the simulator environment. + + Args: + env (instance of EnvBase): environment + initial_state (dict): initial simulation state to load + states (np.array): array of simulation states to load to extract information + actions (np.array): array of actions + """ + assert isinstance(env, EnvBase) + assert len(states) == actions.shape[0] + + # load the initial state + env.reset() + env.reset_to(initial_state) + + all_datagen_infos = [] + traj_len = len(states) + for t in range(traj_len): + # reset to state + env.reset_to({"states" : states[t]}) + # env.base_env.gym.fetch_results(env.base_env.sim, True) + + # extract datagen info + datagen_info = env.base_env.get_datagen_info(action=actions[t]) + all_datagen_infos.append(datagen_info) + + # convert list of dict to dict of list for obs dictionaries (for convenient writes to hdf5 dataset) + all_datagen_infos = TensorUtils.list_of_flat_dict_to_dict_of_list(all_datagen_infos) + for k in all_datagen_infos: + # list to numpy array + all_datagen_infos[k] = np.array(all_datagen_infos[k]) + + return all_datagen_infos + + +def extract_datagen_info_from_trajectory_real_robot( + env_meta, + observations, + actions, +): + """ + Real robot version of helper function to extract datagen info. On the real robot, we assume + we stored "action-free" datagen-info directly in the observations, and we will use the actions + to compute the target poses needed for datagen offline. + + Args: + env_meta (dict): dictionary containing environment metadata from dataset + observations (list): list of observations for this trajectory (each will be a dict) + actions (np.array): array of actions + """ + assert len(observations) == actions.shape[0] + traj_len = actions.shape[0] + + # use static method of appropriate mimicgen base robot env to convert controller pose in obs + action to target pose + if EnvUtils.is_real_robot_gprs_env(env_meta=env_meta): + from mimicgen.envs.real_gprs.base import MG_Real_GPRS_Env + static_method = MG_Real_GPRS_Env.action_to_pose_target_stateless + # TODO: remove hardcode of action scaling here + max_dpos = np.array([0.08, 0.08, 0.08]) + max_drot = np.array([0.5, 0.5, 0.5]) + elif EnvUtils.is_real_robot_env(env_meta=env_meta): + from mimicgen.envs.real.base import MG_Real_Env + static_method = MG_Real_Env.action_to_pose_target_stateless + action_scale = np.array(env_meta["env_kwargs"]["action_scale"]).reshape(-1) + max_dpos = action_scale[:3] + max_drot = action_scale[3:6] + else: + raise Exception("env meta must be real robot type") + + all_datagen_infos = [] + for t in range(traj_len): + # first copy action-free datagen info from observation + datagen_info = dict() + obs = observations[t] + for k in obs: + if k.startswith("datagen_"): + datagen_info[k[8:]] = np.array(obs[k]) + + # use action and controller pose in observation to compute target pose + datagen_info["target_pos"], datagen_info["target_rot"] = static_method( + action=actions[t], + start_pos=obs["datagen_eef_pos"], + start_rot=obs["datagen_eef_rot"], + max_dpos=max_dpos, + max_drot=max_drot, + ) + + all_datagen_infos.append(datagen_info) + + # convert list of dict to dict of list for obs dictionaries (for convenient writes to hdf5 dataset) + all_datagen_infos = TensorUtils.list_of_flat_dict_to_dict_of_list(all_datagen_infos) + for k in all_datagen_infos: + # list to numpy array + all_datagen_infos[k] = np.array(all_datagen_infos[k]) + + return all_datagen_infos + +""" End of dataset_states_to_args copy over """ + +def extract_trajectory( + env, + initial_state, + states, + actions, + actions_abs, + done_mode, +): + """ + Helper function to extract observations, rewards, and dones along a trajectory using + the simulator environment. + + Args: + env (instance of EnvBase): environment + initial_state (dict): initial simulation state to load + states (np.array): array of simulation states to load to extract information + actions (np.array): array of actions + done_mode (int): how to write done signal. If 0, done is 1 whenever s' is a + success state. If 1, done is 1 at the end of each trajectory. + If 2, do both. + """ + assert isinstance(env, EnvBase) + assert states.shape[0] == actions.shape[0] + + # load the initial state + env.reset() + obs = env.reset_to(initial_state) + + traj = dict( + obs=[], + next_obs=[], + rewards=[], + dones=[], + actions=np.array(actions), + states=np.array(states), + initial_state_dict=initial_state, + ) + if actions_abs is not None: + traj["actions_abs"] = np.array(actions_abs) + + traj_len = states.shape[0] + # iteration variable @t is over "next obs" indices + for t in range(1, traj_len + 1): + + # get next observation + if t == traj_len: + # play final action to get next observation for last timestep + next_obs, _, _, _ = env.step(actions[t - 1]) + else: + # reset to simulator state to get observation + next_obs = env.reset_to({"states" : states[t]}) + + # infer reward signal + # note: our tasks use reward r(s'), reward AFTER transition, so this is + # the reward for the current timestep + r = env.get_reward() + + # infer done signal + done = False + if (done_mode == 1) or (done_mode == 2): + # done = 1 at end of trajectory + done = done or (t == traj_len) + if (done_mode == 0) or (done_mode == 2): + # done = 1 when s' is task success state + done = done or env.is_success()["task"] + done = int(done) + + # collect transition + traj["obs"].append(obs) + traj["next_obs"].append(next_obs) + traj["rewards"].append(r) + traj["dones"].append(done) + + # update for next iter + obs = deepcopy(next_obs) + + # convert list of dict to dict of list for obs dictionaries (for convenient writes to hdf5 dataset) + traj["obs"] = TensorUtils.list_of_flat_dict_to_dict_of_list(traj["obs"]) + traj["next_obs"] = TensorUtils.list_of_flat_dict_to_dict_of_list(traj["next_obs"]) + + # list to numpy array + for k in traj: + if k == "initial_state_dict": + continue + if isinstance(traj[k], dict): + for kp in traj[k]: + traj[k][kp] = np.array(traj[k][kp]) + else: + traj[k] = np.array(traj[k]) + + return traj + + +""" The process that writes over the generated files to memory """ +def write_traj_to_file(args, output_path, total_samples, total_run, processes, is_robosuite_env, mul_queue): + f = h5py.File(args.dataset, "r") + f_out = h5py.File(output_path, "w") + data_grp = f_out.create_group("data") + start_time = time.time() + num_processed = 0 + + try: + while((total_run.value < (processes)) or not mul_queue.empty()): + if not mul_queue.empty(): + num_processed = num_processed + 1 + item = mul_queue.get() + ep = item[0] + traj = item[1] + datagen_info = item[2] + process_num = item[3] + try: + ep_data_grp = data_grp.create_group(ep) + ep_data_grp.create_dataset("actions", data=np.array(traj["actions"])) + ep_data_grp.create_dataset("states", data=np.array(traj["states"])) + ep_data_grp.create_dataset("rewards", data=np.array(traj["rewards"])) + ep_data_grp.create_dataset("dones", data=np.array(traj["dones"])) + if "actions_abs" in traj: + ep_data_grp.create_dataset("actions_abs", data=np.array(traj["actions_abs"])) + for k in traj["obs"]: + if args.compress: + ep_data_grp.create_dataset("obs/{}".format(k), data=np.array(traj["obs"][k]), compression="gzip") + else: + ep_data_grp.create_dataset("obs/{}".format(k), data=np.array(traj["obs"][k])) + if not args.exclude_next_obs: + if args.compress: + ep_data_grp.create_dataset("next_obs/{}".format(k), data=np.array(traj["next_obs"][k]), compression="gzip") + else: + ep_data_grp.create_dataset("next_obs/{}".format(k), data=np.array(traj["next_obs"][k])) + + for k in datagen_info: + ep_data_grp.create_dataset("datagen_info/{}".format(k), data=np.array(datagen_info[k])) + + # copy action dict (if applicable) + if "data/{}/action_dict".format(ep) in f: + action_dict = f["data/{}/action_dict".format(ep)] + for k in action_dict: + ep_data_grp.create_dataset("action_dict/{}".format(k), data=np.array(action_dict[k][()])) + + # episode metadata + if is_robosuite_env: + ep_data_grp.attrs["model_file"] = traj["initial_state_dict"]["model"] # model xml for this episode + if "ep_info" in f["data/{}".format(ep)].attrs: + ep_data_grp.attrs["ep_info"] = f["data/{}".format(ep)].attrs["ep_info"] + ep_data_grp.attrs["num_samples"] = traj["actions"].shape[0] # number of transitions in this episode + + total_samples.value += traj["actions"].shape[0] + except Exception as e: + print("++"*50) + print(f"Error at Process {process_num} on episode {ep} with \n\n {e}") + print("++"*50) + raise Exception("Write out to file has failed") + print("ep {}: wrote {} transitions to group {} at process {} with {} finished".format(num_processed, ep_data_grp.attrs["num_samples"], ep, process_num, total_run.value)) + except KeyboardInterrupt: + print("Control C pressed. Closing File and ending \n\n\n\n\n\n\n") + + + if "mask" in f: + f.copy("mask", f_out) + + # global metadata + data_grp.attrs["total"] = total_samples.value + env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=args.dataset) + env = EnvUtils.create_env_for_data_processing( + env_meta=env_meta, + camera_names=args.camera_names, + camera_height=args.camera_height, + camera_width=args.camera_width, + reward_shaping=args.shaped, + ) + print("total processes end {}".format(total_run.value)) + data_grp.attrs["env_args"] = json.dumps(env.serialize(), indent=4) # environment info + print("Wrote {} trajectories to {}".format(total_samples.value, output_path)) + + f_out.close() + f.close() + print("Writing has finished") + + end_time = time.time() + + # Calculate the elapsed time + elapsed_time = end_time - start_time + + print(f"Time elapsed: {elapsed_time:.2f} seconds") + return + +# runs multiple trajectory. If there has been an unrecoverable error, the system puts the current work back into the queue and exits +def extract_multiple_trajectories(process_num, current_work_array, work_queue, lock, args2, num_finished, mul_queue): + try: + extract_multiple_trajectories_with_error(process_num, current_work_array, work_queue, lock, args2, mul_queue) + except Exception as e: + work_queue.put(current_work_array[process_num]) + print("*>*"*50) + print(e) + + num_finished.value = num_finished.value + 1 + + +def retrieve_new_index(process_num, current_work_array, work_queue, lock): + with lock: + if work_queue.empty(): + return -1 + try: + tmp = work_queue.get(False) + current_work_array[process_num] = tmp + return tmp + except queue.Empty: + return -1 + +def extract_multiple_trajectories_with_error(process_num, current_work_array, work_queue, lock, args, mul_queue): + # create environment to use for data processing + + env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=args.dataset) + env = EnvUtils.create_env_for_data_processing( + env_meta=env_meta, + camera_names=args.camera_names, + camera_height=args.camera_height, + camera_width=args.camera_width, + reward_shaping=args.shaped, + ) + + print("==== Using environment with the following metadata ====") + print(json.dumps(env.serialize(), indent=4)) + print("") + + # some operations for playback are robosuite-specific, so determine if this environment is a robosuite env + is_robosuite_env = EnvUtils.is_robosuite_env(env_meta) + + if args.real: + is_robosuite_env = False + is_simpler_env = False + is_factory_env = False + + else: + # some operations are env-type-specific + is_simpler_env = False #EnvUtils.is_simpler_env(env_meta) + is_factory_env = False #EnvUtils.is_factory_env(env_meta) + + # list of all demonstration episodes (sorted in increasing number order) + f = h5py.File(args.dataset, "r") + demos = list(f["data"].keys()) + inds = np.argsort([int(elem[5:]) for elem in demos]) + demos = [demos[i] for i in inds] + + # maybe reduce the number of demonstrations to playback + if args.n is not None: + demos = demos[:args.n] + + ind = retrieve_new_index(process_num, current_work_array, work_queue, lock) + while (not work_queue.empty()) and (ind != -1): + try: + # print("Running {} index".format(ind)) + ep = demos[ind] + + # prepare initial state to reload from + states = f["data/{}/states".format(ep)][()] + initial_state = dict(states=states[0]) + if is_robosuite_env: + initial_state["model"] = f["data/{}".format(ep)].attrs["model_file"] + + # extract obs, rewards, dones + actions = f["data/{}/actions".format(ep)][()] + if "data/{}/actions_abs".format(ep) in f: + actions_abs = f["data/{}/actions_abs".format(ep)][()] + else: + actions_abs = None + + + traj = extract_trajectory( + env=env, + initial_state=initial_state, + states=states, + actions=actions, + actions_abs=actions_abs, + done_mode=args.done_mode, + ) + + # maybe copy reward or done signal from source file + if args.copy_rewards: + traj["rewards"] = f["data/{}/rewards".format(ep)][()] + if args.copy_dones: + traj["dones"] = f["data/{}/dones".format(ep)][()] + + + ep_grp = f["data/{}".format(ep)] + + if args.real: + traj_len = ep_grp["actions"].shape[0] + obs = [] + obs_grp = ep_grp["obs"] + for i in range(traj_len): + obs.append( + { k : np.array(obs_grp[k][i]) for k in obs_grp } + ) + datagen_info = extract_datagen_info_from_trajectory_real_robot( + env_meta=env_meta, + observations=obs, + actions=ep_grp["actions"][()], + ) + else: + # prepare states to reload from + if is_simpler_env or is_factory_env: + # states are dictionaries - make list of dictionaries + traj_len = ep_grp["actions"].shape[0] + states = [] + state_grp = ep_grp["states"] + for i in range(traj_len): + states.append( + { k : np.array(state_grp[k][i]) for k in state_grp } + ) + else: + states = ep_grp["states"][()] + initial_state = dict(states=states[0]) + if is_robosuite_env: + initial_state["model"] = ep_grp.attrs["model_file"] + + # extract datagen info + actions = ep_grp["actions"][()] + datagen_info = extract_datagen_info_from_trajectory( + env=env, + initial_state=initial_state, + states=states, + actions=actions, + ) + + # store transitions + + # IMPORTANT: keep name of group the same as source file, to make sure that filter keys are + # consistent as well + # print("ADD TO QUEUE {} of index {}".format(process_num, ind)) + mul_queue.put([ep, traj, datagen_info, process_num]) + + ind = retrieve_new_index(process_num, current_work_array, work_queue, lock) + except Exception as e: + print("_"*50) + print(process_num) + print("Error {} {}".format(ind, e)) + print("_"*50) + env = EnvUtils.create_env_for_data_processing( #when it errors, it like blows up the environment for some reason + env_meta=env_meta, + camera_names=args.camera_names, + camera_height=args.camera_height, + camera_width=args.camera_width, + reward_shaping=args.shaped, + ) + + f.close() + print("Process {} finished".format(process_num)) + +def dataset_states_to_obs_multiprocessing(args): + # create environment to use for data processing + + # output file in same directory as input file + output_name = args.output_name + if output_name is None: + if len(args.camera_names) == 0: + output_name = os.path.basename(args.dataset)[:-5] + "_ld.hdf5" + else: + output_name = os.path.basename(args.dataset)[:-5] + "_im{}.hdf5".format(args.camera_width) + + output_path = os.path.join(os.path.dirname(args.dataset), output_name) + + print("input file: {}".format(args.dataset)) + print("output file: {}".format(output_path)) + + + f = h5py.File(args.dataset, "r") + demos = list(f["data"].keys()) + inds = np.argsort([int(elem[5:]) for elem in demos]) + demos = [demos[i] for i in inds] + + if args.n is not None: + demos = demos[:args.n] + + num_demos = len(demos) + f.close() + + + env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=args.dataset) + is_robosuite_env = EnvUtils.is_robosuite_env(env_meta) + num_processes = 8 + + index = multiprocessing.Value('i', 0) + lock = multiprocessing.Lock() + total_samples_shared = multiprocessing.Value('i', 0) + num_finished = multiprocessing.Value('i', 0) + mul_queue = multiprocessing.Queue() + work_queue = multiprocessing.Queue() + for index in range(num_demos): + work_queue.put(index) + current_work_array = multiprocessing.Array('i', num_processes) + processes = [] + for i in range(num_processes): + process = multiprocessing.Process(target=extract_multiple_trajectories, args=(i, current_work_array, work_queue, lock, args, num_finished, mul_queue)) + processes.append(process) + + process1 = multiprocessing.Process(target=write_traj_to_file, args=(args, output_path, total_samples_shared, num_finished, num_processes, is_robosuite_env, mul_queue)) + processes.append(process1) + + for process in processes: + process.start() + + for process in processes: + process.join() + + print("Finished Multiprocessing") + return + +def dataset_states_to_obs(args): + # create environment to use for data processing + env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=args.dataset) + env = EnvUtils.create_env_for_data_processing( + env_meta=env_meta, + camera_names=args.camera_names, + camera_height=args.camera_height, + camera_width=args.camera_width, + reward_shaping=args.shaped, + ) + + print("==== Using environment with the following metadata ====") + print(json.dumps(env.serialize(), indent=4)) + print("") + + # some operations for playback are robosuite-specific, so determine if this environment is a robosuite env + is_robosuite_env = EnvUtils.is_robosuite_env(env_meta) + + # list of all demonstration episodes (sorted in increasing number order) + f = h5py.File(args.dataset, "r") + demos = list(f["data"].keys()) + inds = np.argsort([int(elem[5:]) for elem in demos]) + demos = [demos[i] for i in inds] + + # maybe reduce the number of demonstrations to playback + if args.n is not None: + demos = demos[:args.n] + + # output file in same directory as input file + output_name = args.output_name + if output_name is None: + if len(args.camera_names) == 0: + output_name = os.path.basename(args.dataset)[:-5] + "_ld.hdf5" + else: + output_name = os.path.basename(args.dataset)[:-5] + "_im{}.hdf5".format(args.camera_width) + + output_path = os.path.join(os.path.dirname(args.dataset), output_name) + f_out = h5py.File(output_path, "w") + data_grp = f_out.create_group("data") + print("input file: {}".format(args.dataset)) + print("output file: {}".format(output_path)) + + total_samples = 0 + for ind in range(len(demos)): + # for ind in range(1005): + ep = demos[ind] + + # prepare initial state to reload from + states = f["data/{}/states".format(ep)][()] + initial_state = dict(states=states[0]) + if is_robosuite_env: + initial_state["model"] = f["data/{}".format(ep)].attrs["model_file"] + + # extract obs, rewards, dones + actions = f["data/{}/actions".format(ep)][()] + if "data/{}/actions_abs".format(ep) in f: + actions_abs = f["data/{}/actions_abs".format(ep)][()] + else: + actions_abs = None + traj = extract_trajectory( + env=env, + initial_state=initial_state, + states=states, + actions=actions, + actions_abs=actions_abs, + done_mode=args.done_mode, + ) + + # maybe copy reward or done signal from source file + if args.copy_rewards: + traj["rewards"] = f["data/{}/rewards".format(ep)][()] + if args.copy_dones: + traj["dones"] = f["data/{}/dones".format(ep)][()] + + # store transitions + + # IMPORTANT: keep name of group the same as source file, to make sure that filter keys are + # consistent as well + ep_data_grp = data_grp.create_group(ep) + ep_data_grp.create_dataset("actions", data=np.array(traj["actions"])) + ep_data_grp.create_dataset("states", data=np.array(traj["states"])) + ep_data_grp.create_dataset("rewards", data=np.array(traj["rewards"])) + ep_data_grp.create_dataset("dones", data=np.array(traj["dones"])) + if "actions_abs" in traj: + ep_data_grp.create_dataset("actions_abs", data=np.array(traj["actions_abs"])) + for k in traj["obs"]: + if args.compress: + ep_data_grp.create_dataset("obs/{}".format(k), data=np.array(traj["obs"][k]), compression="gzip") + else: + ep_data_grp.create_dataset("obs/{}".format(k), data=np.array(traj["obs"][k])) + if not args.exclude_next_obs: + if args.compress: + ep_data_grp.create_dataset("next_obs/{}".format(k), data=np.array(traj["next_obs"][k]), compression="gzip") + else: + ep_data_grp.create_dataset("next_obs/{}".format(k), data=np.array(traj["next_obs"][k])) + + # copy action dict (if applicable) + if "data/{}/action_dict".format(ep) in f: + action_dict = f["data/{}/action_dict".format(ep)] + for k in action_dict: + ep_data_grp.create_dataset("action_dict/{}".format(k), data=np.array(action_dict[k][()])) + + # episode metadata + if is_robosuite_env: + ep_data_grp.attrs["model_file"] = traj["initial_state_dict"]["model"] # model xml for this episode + if "ep_info" in f["data/{}".format(ep)].attrs: + ep_data_grp.attrs["ep_info"] = f["data/{}".format(ep)].attrs["ep_info"] + ep_data_grp.attrs["num_samples"] = traj["actions"].shape[0] # number of transitions in this episode + total_samples += traj["actions"].shape[0] + print("ep {}: wrote {} transitions to group {}".format(ind, ep_data_grp.attrs["num_samples"], ep)) + + + # copy over all filter keys that exist in the original hdf5 + if "mask" in f: + f.copy("mask", f_out) + + # global metadata + data_grp.attrs["total"] = total_samples + data_grp.attrs["env_args"] = json.dumps(env.serialize(), indent=4) # environment info + print("Wrote {} trajectories to {}".format(len(demos), output_path)) + + f.close() + f_out.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--dataset", + type=str, + required=True, + help="path to input hdf5 dataset", + ) + # name of hdf5 to write - it will be in the same directory as @dataset + parser.add_argument( + "--output_name", + type=str, + help="name of output hdf5 dataset", + ) + + # specify number of demos to process - useful for debugging conversion with a handful + # of trajectories + parser.add_argument( + "--n", + type=int, + default=None, + help="(optional) stop after n trajectories are processed", + ) + + # flag for reward shaping + parser.add_argument( + "--shaped", + action='store_true', + help="(optional) use shaped rewards", + ) + + # camera names to use for observations + parser.add_argument( + "--camera_names", + type=str, + nargs='+', + default=[], + help="(optional) camera name(s) to use for image observations. Leave out to not use image observations.", + ) + + parser.add_argument( + "--camera_height", + type=int, + default=84, + help="(optional) height of image observations", + ) + + parser.add_argument( + "--camera_width", + type=int, + default=84, + help="(optional) width of image observations", + ) + + # specifies how the "done" signal is written. If "0", then the "done" signal is 1 wherever + # the transition (s, a, s') has s' in a task completion state. If "1", the "done" signal + # is one at the end of every trajectory. If "2", the "done" signal is 1 at task completion + # states for successful trajectories and 1 at the end of all trajectories. + parser.add_argument( + "--done_mode", + type=int, + default=0, + help="how to write done signal. If 0, done is 1 whenever s' is a success state.\ + If 1, done is 1 at the end of each trajectory. If 2, both.", + ) + + # flag for copying rewards from source file instead of re-writing them + parser.add_argument( + "--copy_rewards", + action='store_true', + help="(optional) copy rewards from source file instead of inferring them", + ) + + # flag for copying dones from source file instead of re-writing them + parser.add_argument( + "--copy_dones", + action='store_true', + help="(optional) copy dones from source file instead of inferring them", + ) + + # flag to exclude next obs in dataset + parser.add_argument( + "--exclude-next-obs", + action='store_true', + help="(optional) exclude next obs in dataset", + ) + + # flag to compress observations with gzip option in hdf5 + parser.add_argument( + "--compress", + action='store_true', + help="(optional) compress observations with gzip option in hdf5", + ) + + # real robot + parser.add_argument( + "--real", + action='store_true', + help="specify this if using real robot dataset", + ) + + args = parser.parse_args() + # dataset_states_to_obs(args) + dataset_states_to_obs_multiprocessing(args) diff --git a/aloha-devel/robomimic/scripts/download_datasets.py b/aloha-devel/robomimic/scripts/download_datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..caf3a280a14aec6f3c39157e9f9d84dd2a2486c4 --- /dev/null +++ b/aloha-devel/robomimic/scripts/download_datasets.py @@ -0,0 +1,163 @@ +""" +Script to download datasets packaged with the repository. By default, all +datasets will be stored at robomimic/datasets, unless the @download_dir +argument is supplied. We recommend using the default, as most examples that +use these datasets assume that they can be found there. + +The @tasks, @dataset_types, and @hdf5_types arguments can all be supplied +to choose which datasets to download. + +Args: + download_dir (str): Base download directory. Created if it doesn't exist. + Defaults to datasets folder in repository - only pass in if you would + like to override the location. + + tasks (list): Tasks to download datasets for. Defaults to lift task. Pass 'all' to + download all tasks (sim + real) 'sim' to download all sim tasks, 'real' to + download all real tasks, or directly specify the list of tasks. + + dataset_types (list): Dataset types to download datasets for (e.g. ph, mh, mg). + Defaults to ph. Pass 'all' to download datasets for all available dataset + types per task, or directly specify the list of dataset types. + + hdf5_types (list): hdf5 types to download datasets for (e.g. raw, low_dim, image). + Defaults to low_dim. Pass 'all' to download datasets for all available hdf5 + types per task and dataset, or directly specify the list of hdf5 types. + +Example usage: + + # default behavior - just download lift proficient-human low-dim dataset + python download_datasets.py + + # download low-dim proficient-human datasets for all simulation tasks + # (do a dry run first to see which datasets would be downloaded) + python download_datasets.py --tasks sim --dataset_types ph --hdf5_types low_dim --dry_run + python download_datasets.py --tasks sim --dataset_types ph --hdf5_types low_dim + + # download all low-dim and image multi-human datasets for the can and square tasks + python download_datasets.py --tasks can square --dataset_types mh --hdf5_types low_dim image + + # download the sparse reward machine-generated low-dim datasets + python download_datasets.py --tasks all --dataset_types mg --hdf5_types low_dim_sparse + + # download all real robot datasets + python download_datasets.py --tasks real +""" +import os +import argparse + +import robomimic +import robomimic.utils.file_utils as FileUtils +from robomimic import DATASET_REGISTRY + +ALL_TASKS = ["lift", "can", "square", "transport", "tool_hang", "lift_real", "can_real", "tool_hang_real"] +ALL_DATASET_TYPES = ["ph", "mh", "mg", "paired"] +ALL_HDF5_TYPES = ["raw", "low_dim", "image", "low_dim_sparse", "low_dim_dense", "image_sparse", "image_dense"] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + # directory to download datasets to + parser.add_argument( + "--download_dir", + type=str, + default=None, + help="Base download directory. Created if it doesn't exist. Defaults to datasets folder in repository.", + ) + + # tasks to download datasets for + parser.add_argument( + "--tasks", + type=str, + nargs='+', + default=["lift"], + help="Tasks to download datasets for. Defaults to lift task. Pass 'all' to download all tasks (sim + real)\ + 'sim' to download all sim tasks, 'real' to download all real tasks, or directly specify the list of\ + tasks.", + ) + + # dataset types to download datasets for + parser.add_argument( + "--dataset_types", + type=str, + nargs='+', + default=["ph"], + help="Dataset types to download datasets for (e.g. ph, mh, mg). Defaults to ph. Pass 'all' to download \ + datasets for all available dataset types per task, or directly specify the list of dataset types.", + ) + + # hdf5 types to download datasets for + parser.add_argument( + "--hdf5_types", + type=str, + nargs='+', + default=["low_dim"], + help="hdf5 types to download datasets for (e.g. raw, low_dim, image). Defaults to raw. Pass 'all' \ + to download datasets for all available hdf5 types per task and dataset, or directly specify the list\ + of hdf5 types.", + ) + + # dry run - don't actually download datasets, but print which datasets would be downloaded + parser.add_argument( + "--dry_run", + action='store_true', + help="set this flag to do a dry run to only print which datasets would be downloaded" + ) + + args = parser.parse_args() + + # set default base directory for downloads + default_base_dir = args.download_dir + if default_base_dir is None: + default_base_dir = os.path.join(robomimic.__path__[0], "../datasets") + + # load args + download_tasks = args.tasks + if "all" in download_tasks: + assert len(download_tasks) == 1, "all should be only tasks argument but got: {}".format(args.tasks) + download_tasks = ALL_TASKS + elif "sim" in download_tasks: + assert len(download_tasks) == 1, "sim should be only tasks argument but got: {}".format(args.tasks) + download_tasks = [task for task in ALL_TASKS if "real" not in task] + elif "real" in download_tasks: + assert len(download_tasks) == 1, "real should be only tasks argument but got: {}".format(args.tasks) + download_tasks = [task for task in ALL_TASKS if "real" in task] + + download_dataset_types = args.dataset_types + if "all" in download_dataset_types: + assert len(download_dataset_types) == 1, "all should be only dataset_types argument but got: {}".format(args.dataset_types) + download_dataset_types = ALL_DATASET_TYPES + + download_hdf5_types = args.hdf5_types + if "all" in download_hdf5_types: + assert len(download_hdf5_types) == 1, "all should be only hdf5_types argument but got: {}".format(args.hdf5_types) + download_hdf5_types = ALL_HDF5_TYPES + + # download requested datasets + for task in DATASET_REGISTRY: + if task in download_tasks: + for dataset_type in DATASET_REGISTRY[task]: + if dataset_type in download_dataset_types: + for hdf5_type in DATASET_REGISTRY[task][dataset_type]: + if hdf5_type in download_hdf5_types: + download_dir = os.path.abspath(os.path.join(default_base_dir, task, dataset_type)) + print("\nDownloading dataset:\n task: {}\n dataset type: {}\n hdf5 type: {}\n download path: {}" + .format(task, dataset_type, hdf5_type, download_dir)) + url = DATASET_REGISTRY[task][dataset_type][hdf5_type]["url"] + if url is None: + print( + "Skipping {}-{}-{}, no url for dataset exists.".format(task, dataset_type, hdf5_type) + + " Create this dataset locally by running the appropriate command from robomimic/scripts/extract_obs_from_raw_datasets.sh." + ) + continue + if args.dry_run: + print("\ndry run: skip download") + else: + # Make sure path exists and create if it doesn't + os.makedirs(download_dir, exist_ok=True) + FileUtils.download_url( + url=DATASET_REGISTRY[task][dataset_type][hdf5_type]["url"], + download_dir=download_dir, + ) + print("") diff --git a/aloha-devel/robomimic/scripts/extract_obs_from_raw_datasets.sh b/aloha-devel/robomimic/scripts/extract_obs_from_raw_datasets.sh new file mode 100644 index 0000000000000000000000000000000000000000..00fc78f8bf08df5339e79c65019db683dfac6e59 --- /dev/null +++ b/aloha-devel/robomimic/scripts/extract_obs_from_raw_datasets.sh @@ -0,0 +1,140 @@ +#!/bin/bash + +# This script holds the commands that were used to go from raw robosuite demo.hdf5 files +# to our processed low-dim and image hdf5 files. + +BASE_DATASET_DIR="../../datasets" +echo "Using base dataset directory: $BASE_DATASET_DIR" + + +### NOTE: we use done-mode 0 for MG (dones on task success) ### + + +### mg ### + + +# lift - mg, sparse +python dataset_states_to_obs.py --done_mode 0 \ +--dataset $BASE_DATASET_DIR/lift/mg/demo_v141.hdf5 \ +--output_name low_dim_sparse_v141.hdf5 +python dataset_states_to_obs.py --done_mode 0 \ +--dataset $BASE_DATASET_DIR/lift/mg/demo_v141.hdf5 \ +--output_name image_sparse_v141.hdf5 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 + +# lift - mg, dense +python dataset_states_to_obs.py --done_mode 0 --shaped \ +--dataset $BASE_DATASET_DIR/lift/mg/demo_v141.hdf5 \ +--output_name low_dim_dense_v141.hdf5 +python dataset_states_to_obs.py --done_mode 0 --shaped \ +--dataset $BASE_DATASET_DIR/lift/mg/demo_v141.hdf5 \ +--output_name image_dense_v141.hdf5 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 + +# can - mg, sparse +python dataset_states_to_obs.py --done_mode 0 \ +--dataset $BASE_DATASET_DIR/can/mg/demo_v141.hdf5 \ +--output_name low_dim_sparse_v141.hdf5 +python dataset_states_to_obs.py --done_mode 0 \ +--dataset $BASE_DATASET_DIR/can/mg/demo_v141.hdf5 \ +--output_name image_sparse_v141.hdf5 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 + +# can - mg, dense +python dataset_states_to_obs.py --done_mode 0 --shaped \ +--dataset $BASE_DATASET_DIR/can/mg/demo_v141.hdf5 \ +--output_name low_dim_dense_v141.hdf5 +python dataset_states_to_obs.py --done_mode 0 --shaped \ +--dataset $BASE_DATASET_DIR/can/mg/demo_v141.hdf5 \ +--output_name image_dense_v141.hdf5 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 + + +### NOTE: we use done-mode 2 for PH / MH (dones on task success and end of trajectory) ### + + +### ph ### + + +# lift - ph +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/lift/ph/demo_v141.hdf5 \ +--output_name low_dim_v141.hdf5 +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/lift/ph/demo_v141.hdf5 \ +--output_name image_v141.hdf5 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 + +# can - ph +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/can/ph/demo_v141.hdf5 \ +--output_name low_dim_v141.hdf5 +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/can/ph/demo_v141.hdf5 \ +--output_name image_v141.hdf5 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 + +# square - ph +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/square/ph/demo_v141.hdf5 \ +--output_name low_dim_v141.hdf5 +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/square/ph/demo_v141.hdf5 \ +--output_name image_v141.hdf5 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 + +# transport - ph +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/transport/ph/demo_v141.hdf5 \ +--output_name low_dim_v141.hdf5 +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/transport/ph/demo_v141.hdf5 \ +--output_name image_v141.hdf5 --camera_names shouldercamera0 shouldercamera1 robot0_eye_in_hand robot1_eye_in_hand --camera_height 84 --camera_width 84 + +# tool hang - ph +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/tool_hang/ph/demo_v141.hdf5 \ +--output_name low_dim_v141.hdf5 +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/tool_hang/ph/demo_v141.hdf5 \ +--output_name image_v141.hdf5 --camera_names sideview robot0_eye_in_hand --camera_height 240 --camera_width 240 + + +### mh ### + + +# lift - mh +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/lift/mh/demo_v141.hdf5 \ +--output_name low_dim_v141.hdf5 +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/lift/mh/demo_v141.hdf5 \ +--output_name image_v141.hdf5 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 + +# can - mh +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/can/mh/demo_v141.hdf5 \ +--output_name low_dim_v141.hdf5 +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/can/mh/demo_v141.hdf5 \ +--output_name image_v141.hdf5 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 + +# square - mh +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/square/mh/demo_v141.hdf5 \ +--output_name low_dim_v141.hdf5 +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/square/mh/demo_v141.hdf5 \ +--output_name image_v141.hdf5 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 + +# transport - mh +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/transport/mh/demo_v141.hdf5 \ +--output_name low_dim_v141.hdf5 +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/transport/mh/demo_v141.hdf5 \ +--output_name image_v141.hdf5 --camera_names shouldercamera0 shouldercamera1 robot0_eye_in_hand robot1_eye_in_hand --camera_height 84 --camera_width 84 + + +### can-paired ### + + +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/can/paired/demo_v141.hdf5 \ +--output_name low_dim_v141.hdf5 +python dataset_states_to_obs.py --done_mode 2 \ +--dataset $BASE_DATASET_DIR/can/paired/demo_v141.hdf5 \ +--output_name image_v141.hdf5 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 diff --git a/aloha-devel/robomimic/scripts/generate_config_templates.py b/aloha-devel/robomimic/scripts/generate_config_templates.py new file mode 100644 index 0000000000000000000000000000000000000000..56e1d8710c124cd418850bf25e016873ed88c49d --- /dev/null +++ b/aloha-devel/robomimic/scripts/generate_config_templates.py @@ -0,0 +1,28 @@ +""" +Helpful script to generate example config files for each algorithm. These should be re-generated +when new config options are added, or when default settings in the config classes are modified. +""" +import os +import json + +import robomimic +from robomimic.config import get_all_registered_configs + + +def main(): + # store template config jsons in this directory + target_dir = os.path.join(robomimic.__path__[0], "exps/templates/") + + # iterate through registered algorithm config classes + all_configs = get_all_registered_configs() + for algo_name in all_configs: + # make config class for this algorithm + c = all_configs[algo_name]() + assert algo_name == c.algo_name + # dump to json + json_path = os.path.join(target_dir, "{}.json".format(algo_name)) + c.dump(filename=json_path) + + +if __name__ == '__main__': + main() diff --git a/aloha-devel/robomimic/scripts/generate_paper_configs.py b/aloha-devel/robomimic/scripts/generate_paper_configs.py new file mode 100644 index 0000000000000000000000000000000000000000..52ed7d5b15a25def7da7a02c7c0e135772f269a0 --- /dev/null +++ b/aloha-devel/robomimic/scripts/generate_paper_configs.py @@ -0,0 +1,1369 @@ +""" +Helper script to generate jsons for reproducing paper experiments. + +Args: + config_dir (str): Directory where generated configs will be placed. + Defaults to 'paper' subfolder in exps folder of repository + + dataset_dir (str): Base dataset directory where released datasets can be + found on disk. Defaults to datasets folder in repository. + + output_dir (str): Base output directory for all training runs that will be + written to generated configs. + +Example usage: + # Assume datasets alredy exist in robomimic/../datasets folder. Configs will be generated under robomimic/exps/paper + python generate_paper_configs.py --output_dir /tmp/experiment_results + + # Specify where datasets exist, and specify where configs should be generated. + python generate_paper_configs.py --config_dir /tmp/configs --dataset_dir /tmp/datasets --output_dir /tmp/experiment_results +""" +import os +import argparse +import robomimic +from robomimic import DATASET_REGISTRY +from robomimic.config import Config, BCConfig, BCQConfig, CQLConfig, HBCConfig, IRISConfig, config_factory + + +def modify_config_for_default_low_dim_exp(config): + """ + Modifies a Config object with experiment, training, and observation settings that + were used across all low-dimensional experiments by default. + + Args: + config (Config instance): config to modify + """ + + with config.experiment.values_unlocked(): + # save model during every evaluation (every 50 epochs) + config.experiment.save.enabled = True + config.experiment.save.every_n_epochs = 50 + + # every epoch is 100 gradient steps, and validation epoch is 10 gradient steps + config.experiment.epoch_every_n_steps = 100 + config.experiment.validation_epoch_every_n_steps = 10 + + # do 50 evaluation rollouts every 50 epochs + # NOTE: horizon will generally get set depending on the task and dataset type + config.experiment.rollout.enabled = True + config.experiment.rollout.n = 50 + config.experiment.rollout.horizon = 400 + config.experiment.rollout.rate = 50 + config.experiment.rollout.warmstart = 0 + config.experiment.rollout.terminate_on_success = True + + with config.train.values_unlocked(): + # assume entire dataset can fit in memory + config.train.num_data_workers = 0 + config.train.hdf5_cache_mode = "all" + + # batch size 100 and 2000 training epochs + config.train.batch_size = 100 + config.train.num_epochs = 2000 + + with config.observation.values_unlocked(): + # default observation is eef pose, gripper finger position, and object information, + # all of which are low-dim. + default_low_dim_obs = [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object", + ] + # handle hierarchical observation configs + if config.algo_name == "hbc": + configs_to_set = [ + config.observation.actor.modalities.obs, + config.observation.planner.modalities.obs, + config.observation.planner.modalities.subgoal, + ] + elif config.algo_name == "iris": + configs_to_set = [ + config.observation.actor.modalities.obs, + config.observation.value_planner.planner.modalities.obs, + config.observation.value_planner.planner.modalities.subgoal, + config.observation.value_planner.value.modalities.obs, + ] + else: + configs_to_set = [config.observation.modalities.obs] + # set all observations / subgoals to use the correct low-dim modalities + for cfg in configs_to_set: + cfg.low_dim = list(default_low_dim_obs) + cfg.rgb = [] + + return config + + +def modify_config_for_default_image_exp(config): + """ + Modifies a Config object with experiment, training, and observation settings that + were used across all image experiments by default. + + Args: + config (Config instance): config to modify + """ + assert config.algo_name not in ["hbc", "iris"], "no image training for HBC and IRIS" + + with config.experiment.values_unlocked(): + # save model during every evaluation (every 20 epochs) + config.experiment.save.enabled = True + config.experiment.save.every_n_epochs = 20 + + # every epoch is 500 gradient steps, and validation epoch is 50 gradient steps + config.experiment.epoch_every_n_steps = 500 + config.experiment.validation_epoch_every_n_steps = 50 + + # do 50 evaluation rollouts every 20 epochs + # NOTE: horizon will generally get set depending on the task and dataset type + config.experiment.rollout.enabled = True + config.experiment.rollout.n = 50 + config.experiment.rollout.horizon = 400 + config.experiment.rollout.rate = 20 + config.experiment.rollout.warmstart = 0 + config.experiment.rollout.terminate_on_success = True + + with config.train.values_unlocked(): + # only cache low-dim info, and use 2 data workers to increase fetch speed for image obs + config.train.num_data_workers = 2 + config.train.hdf5_cache_mode = "low_dim" + + # batch size 16 and 600 training epochs + config.train.batch_size = 16 + config.train.num_epochs = 600 + + + with config.observation.values_unlocked(): + # default low-dim observation is eef pose, gripper finger position + # default image observation is external camera and wrist camera + config.observation.modalities.obs.low_dim = [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + ] + config.observation.modalities.obs.rgb = [ + "agentview_image", + "robot0_eye_in_hand_image", + ] + config.observation.modalities.goal.low_dim = [] + config.observation.modalities.goal.rgb = [] + + # default image encoder architecture is ResNet with spatial softmax + config.observation.encoder.rgb.core_class = "VisualCore" + config.observation.encoder.rgb.core_kwargs.feature_dimension = 64 + config.observation.encoder.rgb.core_kwargs.backbone_class = 'ResNet18Conv' # ResNet backbone for image observations (unused if no image observations) + config.observation.encoder.rgb.core_kwargs.backbone_kwargs.pretrained = False # kwargs for visual core + config.observation.encoder.rgb.core_kwargs.backbone_kwargs.input_coord_conv = False + config.observation.encoder.rgb.core_kwargs.pool_class = "SpatialSoftmax" # Alternate options are "SpatialMeanPool" or None (no pooling) + config.observation.encoder.rgb.core_kwargs.pool_kwargs.num_kp = 32 # Default arguments for "SpatialSoftmax" + config.observation.encoder.rgb.core_kwargs.pool_kwargs.learnable_temperature = False # Default arguments for "SpatialSoftmax" + config.observation.encoder.rgb.core_kwargs.pool_kwargs.temperature = 1.0 # Default arguments for "SpatialSoftmax" + config.observation.encoder.rgb.core_kwargs.pool_kwargs.noise_std = 0.0 + + # observation randomizer class - set to None to use no randomization, or 'CropRandomizer' to use crop randomization + config.observation.encoder.rgb.obs_randomizer_class = "CropRandomizer" + + # kwargs for observation randomizers (for the CropRandomizer, this is size and number of crops) + config.observation.encoder.rgb.obs_randomizer_kwargs.crop_height = 76 + config.observation.encoder.rgb.obs_randomizer_kwargs.crop_width = 76 + config.observation.encoder.rgb.obs_randomizer_kwargs.num_crops = 1 + config.observation.encoder.rgb.obs_randomizer_kwargs.pos_enc = False + + return config + + +def modify_config_for_dataset(config, task_name, dataset_type, hdf5_type, base_dataset_dir, filter_key=None): + """ + Modifies a Config object with experiment, training, and observation settings to + correspond to experiment settings for the dataset collected on @task_name with + dataset source @dataset_type (e.g. ph, mh, mg), and hdf5 type @hdf5_type (e.g. low_dim + or image). + + Args: + config (Config instance): config to modify + + task_name (str): identify task that dataset was collected on + + dataset_type (str): dataset type for this dataset (e.g. ph, mh, mg). + + hdf5_type (str): hdf5 type for this dataset (e.g. raw, low_dim, image). + + base_dataset_dir (str): path to directory where datasets are on disk. + Directory structure is expected to be consistent with the output + of @make_dataset_dirs in the download_datasets.py script. + + filter_key (str): if not None, use the provided filter key to select a subset of the + provided dataset + """ + assert task_name in DATASET_REGISTRY, \ + "task {} not found in dataset registry!".format(task_name) + assert dataset_type in DATASET_REGISTRY[task_name], \ + "dataset type {} not found for task {} in dataset registry!".format(dataset_type, task_name) + assert hdf5_type in DATASET_REGISTRY[task_name][dataset_type], \ + "hdf5 type {} not found for dataset type {} and task {} in dataset registry!".format(hdf5_type, dataset_type, task_name) + + is_real_dataset = "real" in task_name + if is_real_dataset: + assert config.algo_name == "bc", "we only ran BC-RNN on real robot" + else: + assert hdf5_type != "raw", "cannot train on raw demonstrations" + + with config.experiment.values_unlocked(): + + # look up rollout evaluation horizon in registry and set it + config.experiment.rollout.horizon = DATASET_REGISTRY[task_name][dataset_type][hdf5_type]["horizon"] + + if dataset_type == "mg": + # machine-generated datasets did not use validation + config.experiment.validate = False + else: + # all other datasets used validation + config.experiment.validate = True + + if is_real_dataset: + # no evaluation rollouts for real robot training + config.experiment.rollout.enabled = False + + with config.train.values_unlocked(): + # set dataset path and possibly filter keys + url = DATASET_REGISTRY[task_name][dataset_type][hdf5_type]["url"] + if url is None: + # infer file_name + if task_name in ["lift", "can", "square", "tool_hang", "transport"]: + file_name = "{}_v141.hdf5".format(hdf5_type) + elif task_name in ["lift_real", "can_real", "tool_hang_real"]: + file_name = "{}.hdf5".format(hdf5_type) + else: + raise ValueError("Unknown dataset type") + else: + file_name = url.split("/")[-1] + config.train.data = os.path.join(base_dataset_dir, task_name, dataset_type, file_name) + config.train.hdf5_filter_key = None if filter_key is None else filter_key + config.train.hdf5_validation_filter_key = None + if config.experiment.validate: + # set train and valid keys for validation + config.train.hdf5_filter_key = "train" if filter_key is None else "{}_train".format(filter_key) + config.train.hdf5_validation_filter_key = "valid" if filter_key is None else "{}_valid".format(filter_key) + + with config.observation.values_unlocked(): + # maybe modify observation names and randomization sizes (since image size might be different) + + if is_real_dataset: + # modify observation names for real robot datasets + config.observation.modalities.obs.low_dim = [ + "ee_pose", + "gripper_position", + ] + + if task_name == "tool_hang_real": + # side and wrist camera + config.observation.modalities.obs.rgb = [ + "image_side", + "image_wrist", + ] + # 240x240 images -> crops should be 216x216 + config.observation.encoder.rgb.obs_randomizer_kwargs.crop_height = 216 + config.observation.encoder.rgb.obs_randomizer_kwargs.crop_width = 216 + else: + # front and wrist camera + config.observation.modalities.obs.rgb = [ + "image", + "image_wrist", + ] + # 120x120 images -> crops should be 108x108 + config.observation.encoder.rgb.obs_randomizer_kwargs.crop_height = 108 + config.observation.encoder.rgb.obs_randomizer_kwargs.crop_width = 108 + + elif hdf5_type in ["image", "image_sparse", "image_dense"]: + if task_name == "transport": + # robot proprioception per arm + config.observation.modalities.obs.low_dim = [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "robot1_eef_pos", + "robot1_eef_quat", + "robot1_gripper_qpos", + ] + + # shoulder and wrist cameras per arm + config.observation.modalities.obs.rgb = [ + "shouldercamera0_image", + "robot0_eye_in_hand_image", + "shouldercamera1_image", + "robot1_eye_in_hand_image", + ] + elif task_name == "tool_hang": + # side and wrist camera + config.observation.modalities.obs.rgb = [ + "sideview_image", + "robot0_eye_in_hand_image", + ] + # 240x240 images -> crops should be 216x216 + config.observation.encoder.rgb.obs_randomizer_kwargs.crop_height = 216 + config.observation.encoder.rgb.obs_randomizer_kwargs.crop_width = 216 + + elif hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"]: + if task_name == "transport": + # robot proprioception per arm + default_low_dim_obs = [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "robot1_eef_pos", + "robot1_eef_quat", + "robot1_gripper_qpos", + "object", + ] + # handle hierarchical observation configs + if config.algo_name == "hbc": + configs_to_set = [ + config.observation.actor.modalities.obs, + config.observation.planner.modalities.obs, + config.observation.planner.modalities.subgoal, + ] + elif config.algo_name == "iris": + configs_to_set = [ + config.observation.actor.modalities.obs, + config.observation.value_planner.planner.modalities.obs, + config.observation.value_planner.planner.modalities.subgoal, + config.observation.value_planner.value.modalities.obs, + ] + else: + configs_to_set = [config.observation.modalities.obs] + # set all observations / subgoals to use the correct low-dim modalities + for obs_key_config in configs_to_set: + obs_key_config.low_dim = list(default_low_dim_obs) + obs_key_config.rgb = [] + + return config + + +def modify_bc_config_for_dataset(config, task_name, dataset_type, hdf5_type): + """ + Modifies a BCConfig object for training on a particular kind of dataset. This function + just sets algorithm hyperparameters in the algo config depending on the kind of + dataset. + + Args: + config (BCConfig instance): config to modify + + task_name (str): identify task that dataset was collected on. Only used to distinguish + between simulation and real-world, for an assert statement + + dataset_type (str): dataset type for this dataset (e.g. ph, mh, mg, paired). + + hdf5_type (str): hdf5 type for this dataset (e.g. raw, low_dim, image). + """ + assert isinstance(config, BCConfig), "must be BCConfig" + assert config.algo_name == "bc", "must be BCConfig" + assert dataset_type in ["ph", "mh", "mg", "paired"], "invalid dataset type" + is_real_dataset = "real" in task_name + if not is_real_dataset: + assert hdf5_type != "raw", "cannot train on raw demonstrations" + + with config.algo.values_unlocked(): + # base parameters that may get modified + config.algo.optim_params.policy.learning_rate.initial = 1e-4 # learning rate 1e-4 + config.algo.actor_layer_dims = (1024, 1024) # MLP size (1024, 1024) + config.algo.gmm.enabled = True # enable GMM + + if dataset_type == "mg": + # machine-generated datasets don't use GMM + config.algo.gmm.enabled = False # disable GMM + if hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"]: + # low-dim mg uses LR 1e-3 + config.algo.optim_params.policy.learning_rate.initial = 1e-3 # learning rate 1e-3 + + return config + + +def modify_bc_rnn_config_for_dataset(config, task_name, dataset_type, hdf5_type): + """ + Modifies a BCConfig object for training on a particular kind of dataset. This function + just sets algorithm hyperparameters in the algo config depending on the kind of + dataset. + + Args: + config (BCConfig instance): config to modify + + task_name (str): identify task that dataset was collected on. Only used to distinguish + between simulation and real-world, for an assert statement + + dataset_type (str): dataset type for this dataset (e.g. ph, mh, mg, paired). + + hdf5_type (str): hdf5 type for this dataset (e.g. raw, low_dim, image). + """ + assert isinstance(config, BCConfig), "must be BCConfig" + assert config.algo_name == "bc", "must be BCConfig" + assert dataset_type in ["ph", "mh", "mg", "paired"], "invalid dataset type" + is_real_dataset = "real" in task_name + if not is_real_dataset: + assert hdf5_type != "raw", "cannot train on raw demonstrations" + + with config.train.values_unlocked(): + # make sure RNN is enabled with sequence length 10 + config.train.seq_length = 10 + + with config.algo.values_unlocked(): + # make sure RNN is enabled with sequence length 10 + config.algo.rnn.enabled = True + config.algo.rnn.horizon = 10 + + # base parameters that may get modified + config.algo.optim_params.policy.learning_rate.initial = 1e-4 # learning rate 1e-4 + config.algo.actor_layer_dims = () # no MLP layers between rnn layer and output + config.algo.gmm.enabled = True # enable GMM + config.algo.rnn.hidden_dim = 400 # rnn dim 400 + + if dataset_type == "mg": + # update hyperparams for machine-generated datasets + config.algo.gmm.enabled = False # disable GMM + if hdf5_type not in ["low_dim", "low_dim_sparse", "low_dim_dense"]: + # image datasets use RNN dim 1000 + config.algo.rnn.hidden_dim = 1000 # rnn dim 1000 + else: + # update hyperparams for all other dataset types (ph, mh, paired) + if hdf5_type not in ["low_dim", "low_dim_sparse", "low_dim_dense"]: + # image datasets use RNN dim 1000 + config.algo.rnn.hidden_dim = 1000 # rnn dim 1000 + + return config + + +def modify_bcq_config_for_dataset(config, task_name, dataset_type, hdf5_type): + """ + Modifies a BCQConfig object for training on a particular kind of dataset. This function + just sets algorithm hyperparameters in the algo config depending on the kind of + dataset. + + Args: + config (BCQConfig instance): config to modify + + task_name (str): identify task that dataset was collected on. Only used to distinguish + between simulation and real-world, for an assert statement + + dataset_type (str): dataset type for this dataset (e.g. ph, mh, mg, paired). + + hdf5_type (str): hdf5 type for this dataset (e.g. raw, low_dim, image). + """ + assert isinstance(config, BCQConfig), "must be BCQConfig" + assert config.algo_name == "bcq", "must be BCQConfig" + assert dataset_type in ["ph", "mh", "mg", "paired"], "invalid dataset type" + is_real_dataset = "real" in task_name + assert not is_real_dataset, "we only ran BC-RNN on real robot" + if not is_real_dataset: + assert hdf5_type != "raw", "cannot train on raw demonstrations" + + with config.algo.values_unlocked(): + # base parameters that may get modified further + config.algo.optim_params.critic.learning_rate.initial = 1e-4 # all learning rates 1e-3 + config.algo.optim_params.action_sampler.learning_rate.initial = 1e-4 + config.algo.optim_params.actor.learning_rate.initial = 1e-3 + config.algo.actor.enabled = False # disable actor by default + config.algo.action_sampler.vae.enabled = True # use VAE action sampler + config.algo.action_sampler.gmm.enabled = False + config.algo.action_sampler.vae.kl_weight = 0.05 # beta 0.05 for VAE + config.algo.action_sampler.vae.latent_dim = 14 # latent dim 14 + config.algo.action_sampler.vae.prior.learn = False # N(0, 1) prior + config.algo.critic.layer_dims = (300, 400) # all MLP sizes at (300, 400) + config.algo.action_sampler.vae.encoder_layer_dims = (300, 400) + config.algo.action_sampler.vae.decoder_layer_dims = (300, 400) + config.algo.actor.layer_dims = (300, 400) + config.algo.target_tau = 5e-4 # tau 5e-4 + config.algo.discount = 0.99 # discount 0.99 + config.algo.critic.num_action_samples = 10 # number of action sampler samples at train and test + config.algo.critic.num_action_samples_rollout = 100 + + if dataset_type == "mg": + # update hyperparams for machine-generated datasets + config.algo.optim_params.critic.learning_rate.initial = 1e-3 # all learning rates 1e-3 + config.algo.optim_params.action_sampler.learning_rate.initial = 1e-3 + config.algo.optim_params.actor.learning_rate.initial = 1e-3 + config.algo.action_sampler.vae.kl_weight = 0.5 # beta 0.5 for VAE + config.algo.target_tau = 5e-3 # tau 5e-3 + + if hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"]: + # enable actor only on low-dim + config.algo.actor.enabled = True + else: + # make some modifications where needed for human datasets + if hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"]: + if dataset_type in ["mh", "paired"]: + # low-dim, MH had higher layer sizes + config.algo.critic.layer_dims = (1024, 1024) + config.algo.action_sampler.vae.encoder_layer_dims = (1024, 1024) + config.algo.action_sampler.vae.decoder_layer_dims = (1024, 1024) + config.algo.action_sampler.vae.prior_layer_dims = (1024, 1024) + + config.algo.action_sampler.vae.kl_weight = 0.5 + + # use learned GMM prior for MH dataset + config.algo.action_sampler.vae.prior.learn = True + config.algo.action_sampler.vae.prior.is_conditioned = True + config.algo.action_sampler.vae.prior.use_gmm = True + config.algo.action_sampler.vae.prior.gmm_learn_weights = True + else: + if dataset_type == "ph": + # image, PH used higher critic LR of 1e-3 + config.algo.optim_params.critic.learning_rate.initial = 1e-3 + # image datasets used bigger VAE + config.algo.action_sampler.vae.encoder_layer_dims = (1024, 1024) + config.algo.action_sampler.vae.decoder_layer_dims = (1024, 1024) + if dataset_type in ["mh", "paired"]: + # image, MH also had bigger critic + config.algo.critic.layer_dims = (1024, 1024) + + return config + + +def modify_cql_config_for_dataset(config, task_name, dataset_type, hdf5_type): + """ + Modifies a CQLConfig object for training on a particular kind of dataset. This function + just sets algorithm hyperparameters in the algo config depending on the kind of + dataset. + + Args: + config (CQLConfig instance): config to modify + + task_name (str): identify task that dataset was collected on. Only used to distinguish + between simulation and real-world, for an assert statement + + dataset_type (str): dataset type for this dataset (e.g. ph, mh, mg, paired). + + hdf5_type (str): hdf5 type for this dataset (e.g. raw, low_dim, image). + """ + assert isinstance(config, CQLConfig), "must be CQLConfig" + assert config.algo_name == "cql", "must be CQLConfig" + assert dataset_type in ["ph", "mh", "mg", "paired"], "invalid dataset type" + is_real_dataset = "real" in task_name + assert not is_real_dataset, "we only ran BC-RNN on real robot" + if not is_real_dataset: + assert hdf5_type != "raw", "cannot train on raw demonstrations" + + with config.train.values_unlocked(): + # CQL uses batch size 1024 (for low-dim) and 8 (for image) + if hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"]: + config.train.batch_size = 1024 + else: + config.train.batch_size = 8 + + with config.algo.values_unlocked(): + # base parameters that may get modified further + config.algo.optim_params.critic.learning_rate.initial = 1e-3 # learning rates + config.algo.optim_params.actor.learning_rate.initial = 3e-4 + config.algo.actor.target_entropy = "default" # use automatic entropy tuning to default target value + config.algo.critic.deterministic_backup = True # deterministic Q-backup + config.algo.critic.target_q_gap = 5.0 # use Lagrange, with threshold 5.0 + config.algo.critic.min_q_weight = 1.0 + config.algo.target_tau = 5e-3 # tau 5e-3 + config.algo.discount = 0.99 # discount 0.99 + config.algo.critic.layer_dims = (300, 400) # all MLP sizes at (300, 400) + config.algo.actor.layer_dims = (300, 400) + + if hdf5_type not in ["low_dim", "low_dim_sparse", "low_dim_dense"]: + # update policy LR to 1e-4 for image runs + config.algo.optim_params.actor.learning_rate.initial = 1e-4 + + return config + + +def modify_hbc_config_for_dataset(config, task_name, dataset_type, hdf5_type): + """ + Modifies a HBCConfig object for training on a particular kind of dataset. This function + just sets algorithm hyperparameters in the algo config depending on the kind of + dataset. + + Args: + config (HBCConfig instance): config to modify + + task_name (str): identify task that dataset was collected on. Only used to distinguish + between simulation and real-world, for an assert statement + + dataset_type (str): dataset type for this dataset (e.g. ph, mh, mg, paired). + + hdf5_type (str): hdf5 type for this dataset (e.g. raw, low_dim, image). + """ + assert isinstance(config, HBCConfig), "must be HBCConfig" + assert config.algo_name == "hbc", "must be HBCConfig" + assert dataset_type in ["ph", "mh", "mg", "paired"], "invalid dataset type" + assert hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"], "HBC only runs on low-dim" + is_real_dataset = "real" in task_name + assert not is_real_dataset, "we only ran BC-RNN on real robot" + + with config.algo.values_unlocked(): + # base parameters that may get modified further + config.algo.actor.optim_params.policy.learning_rate.initial = 1e-3 # learning rates + config.algo.planner.optim_params.goal_network.learning_rate.initial = 1e-3 + + config.algo.planner.vae.enabled = True # goal VAE settings + config.algo.planner.vae.kl_weight = 5e-4 # beta 5e-4 + config.algo.planner.vae.latent_dim = 16 # latent dim 16 + config.algo.planner.vae.prior.learn = True # learn GMM prior with 10 modes + config.algo.planner.vae.prior.is_conditioned = True + config.algo.planner.vae.prior.use_gmm = True + config.algo.planner.vae.prior.gmm_learn_weights = True + config.algo.planner.vae.prior.gmm_num_modes = 10 + config.algo.planner.vae.encoder_layer_dims = (1024, 1024) # VAE network sizes + config.algo.planner.vae.decoder_layer_dims = (1024, 1024) + config.algo.planner.vae.prior_layer_dims = (1024, 1024) + + config.algo.actor.rnn.hidden_dim = 400 # actor RNN dim + config.algo.actor.actor_layer_dims = () # no MLP layers between rnn layer and output + + if dataset_type == "mg": + # update hyperparams for machine-generated datasets + config.algo.actor.rnn.hidden_dim = 100 + config.algo.actor.actor_layer_dims = (1024, 1024) + + return config + + +def modify_iris_config_for_dataset(config, task_name, dataset_type, hdf5_type): + """ + Modifies a IRISConfig object for training on a particular kind of dataset. This function + just sets algorithm hyperparameters in the algo config depending on the kind of + dataset. + + Args: + config (IRISConfig instance): config to modify + + task_name (str): identify task that dataset was collected on. Only used to distinguish + between simulation and real-world, for an assert statement + + dataset_type (str): dataset type for this dataset (e.g. ph, mh, mg, paired). + + hdf5_type (str): hdf5 type for this dataset (e.g. raw, low_dim, image). + """ + assert isinstance(config, IRISConfig), "must be IRISConfig" + assert config.algo_name == "iris", "must be IRISConfig" + assert dataset_type in ["ph", "mh", "mg", "paired"], "invalid dataset type" + assert hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"], "IRIS only runs on low-dim" + is_real_dataset = "real" in task_name + assert not is_real_dataset, "we only ran BC-RNN on real robot" + + with config.algo.values_unlocked(): + # base parameters that may get modified further + config.algo.actor.optim_params.policy.learning_rate.initial = 1e-3 # learning rates + config.algo.value_planner.planner.optim_params.goal_network.learning_rate.initial = 1e-3 + config.algo.value_planner.value.optim_params.critic.learning_rate.initial = 1e-3 + config.algo.value_planner.value.optim_params.action_sampler.learning_rate.initial = 1e-4 + + config.algo.value_planner.planner.vae.enabled = True # goal VAE settings + config.algo.value_planner.planner.vae.kl_weight = 5e-4 # beta 5e-4 + config.algo.value_planner.planner.vae.latent_dim = 14 # latent dim 14 + config.algo.value_planner.planner.vae.prior.learn = True # learn GMM prior with 10 modes + config.algo.value_planner.planner.vae.prior.is_conditioned = True + config.algo.value_planner.planner.vae.prior.use_gmm = True + config.algo.value_planner.planner.vae.prior.gmm_learn_weights = True + config.algo.value_planner.planner.vae.prior.gmm_num_modes = 10 + config.algo.value_planner.planner.vae.encoder_layer_dims = (1024, 1024) # VAE network sizes + config.algo.value_planner.planner.vae.decoder_layer_dims = (1024, 1024) + config.algo.value_planner.planner.vae.prior_layer_dims = (1024, 1024) + + config.algo.value_planner.value.target_tau = 5e-4 # Value tau + config.algo.value_planner.value.action_sampler.vae.kl_weight = 0.5 # Value KL + config.algo.value_planner.value.action_sampler.vae.latent_dim = 16 + config.algo.value_planner.value.action_sampler.actor_layer_dims = (300, 400) + + config.algo.actor.rnn.hidden_dim = 400 # actor RNN dim + config.algo.actor.actor_layer_dims = () # no MLP layers between rnn layer and output + + if dataset_type in ["mh", "paired"]: + # value LR 1e-4, KL weight is 0.05 for multi-human datasets + config.algo.value_planner.value.optim_params.critic.learning_rate.initial = 1e-4 + config.algo.value_planner.value.action_sampler.vae.kl_weight = 0.05 + + if dataset_type in ["mg"]: + # Enable value actor and set larger target tau + config.algo.value_planner.value.actor.enabled = True + config.algo.value_planner.value.optim_params.actor.learning_rate.initial = 1e-3 + config.algo.value_planner.value.target_tau = 5e-3 + + return config + + +def generate_experiment_config( + base_exp_name, + base_config_dir, + base_dataset_dir, + base_output_dir, + algo_name, + algo_config_modifier, + task_name, + dataset_type, + hdf5_type, + filter_key=None, + additional_name=None, + additional_config_modifier=None, +): + """ + Helper function to generate a config for a particular experiment. + + Args: + base_exp_name (str): name that identifies this set of experiments + + base_config_dir (str): base directory to place generated configs + + base_dataset_dir (str): path to directory where datasets are on disk. + Directory structure is expected to be consistent with the output + of @make_dataset_dirs in the download_datasets.py script. + + base_output_dir (str): directory to save training results to. If None, will use the directory + from the default algorithm configs. + + algo_name (str): identifies the algorithm - one of ["bc", "bc_rnn", "bcq", "cql", hbc", "iris"] + + algo_config_modifier (function): function to modify config to add algo hyperparameter + settings, given the task, dataset, and hdf5 types. + + task_name (str): identify task that dataset was collected on. Only used to distinguish + between simulation and real-world, for an assert statement + + dataset_type (str): dataset type for this dataset (e.g. ph, mh, mg, paired). + + hdf5_type (str): hdf5 type for this dataset (e.g. raw, low_dim, image). + + filter_key (str): if not None, use the provided filter key to select a subset of the + provided dataset + + additional_name (str): if provided, will add this name to the generated experiment name, and + the name of the generated config json + + additional_config_modifier (function): if provided, run this last function on the config + to make final modifications before generating the json. + """ + if "real" not in task_name: + assert hdf5_type != "raw", "cannot train on raw demonstrations" + + # decide whether to use low-dim or image training defaults + modifier_for_obs = modify_config_for_default_image_exp + if hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"]: + modifier_for_obs = modify_config_for_default_low_dim_exp + + algo_config_name = "bc" if algo_name == "bc_rnn" else algo_name + config = config_factory(algo_name=algo_config_name) + # turn into default config for observation modalities (e.g.: low-dim or rgb) + config = modifier_for_obs(config) + # add in config based on the dataset + config = modify_config_for_dataset( + config=config, + task_name=task_name, + dataset_type=dataset_type, + hdf5_type=hdf5_type, + base_dataset_dir=base_dataset_dir, + filter_key=filter_key, + ) + # add in algo hypers based on dataset + config = algo_config_modifier( + config=config, + task_name=task_name, + dataset_type=dataset_type, + hdf5_type=hdf5_type, + ) + if additional_config_modifier is not None: + # use additional config modifier if provided + config = additional_config_modifier(config) + + # account for filter key in experiment naming and directory naming + filter_key_str = "_{}".format(filter_key) if filter_key is not None else "" + dataset_type_dir = "{}/{}".format(dataset_type, filter_key) if filter_key is not None else dataset_type + + # account for @additional_name + additional_name_str = "_{}".format(additional_name) if additional_name is not None else "" + json_name = "{}{}".format(algo_name, additional_name_str) + + # set experiment name + with config.experiment.values_unlocked(): + config.experiment.name = "{}_{}_{}_{}{}_{}{}".format(base_exp_name, algo_name, task_name, dataset_type, filter_key_str, hdf5_type, additional_name_str) + # set output folder + with config.train.values_unlocked(): + if base_output_dir is None: + base_output_dir = config.train.output_dir + config.train.output_dir = os.path.join(base_output_dir, base_exp_name, algo_name, task_name, dataset_type_dir, hdf5_type, "trained_models") + + # save config to json file + dir_to_save = os.path.join(base_config_dir, base_exp_name, task_name, dataset_type_dir, hdf5_type) + os.makedirs(dir_to_save, exist_ok=True) + json_path = os.path.join(dir_to_save, "{}.json".format(json_name)) + config.dump(filename=json_path) + + return config, json_path + + +def generate_core_configs( + base_config_dir, + base_dataset_dir, + base_output_dir, + algo_to_config_modifier, +): + """ + Helper function to generate all configs for core set of experiments. + + Args: + base_config_dir (str): base directory to place generated configs + + base_dataset_dir (str): path to directory where datasets are on disk. + Directory structure is expected to be consistent with the output + of @make_dataset_dirs in the download_datasets.py script. + + base_output_dir (str): directory to save training results to. If None, will use the directory + from the default algorithm configs. + + algo_to_config_modifier (dict): dictionary that maps algo name to a function that modifies configs + to add algo hyperparameter settings, given the task, dataset, and hdf5 types. + """ + core_json_paths = Config() # use for convenient nested dict + for task in DATASET_REGISTRY: + for dataset_type in DATASET_REGISTRY[task]: + for hdf5_type in DATASET_REGISTRY[task][dataset_type]: + # if not real robot dataset, skip raw hdf5 + is_real_dataset = ("real" in task) + if not is_real_dataset and hdf5_type == "raw": + continue + + # get list of algorithms to generate configs for, for this hdf5 dataset + algos_to_generate = ["bc", "bc_rnn", "bcq", "cql", "hbc", "iris"] + if hdf5_type not in ["low_dim", "low_dim_sparse", "low_dim_dense"]: + # no hbc or iris for image runs + algos_to_generate = algos_to_generate[:-2] + if is_real_dataset: + # we only ran BC-RNN on real robot + algos_to_generate = ["bc_rnn"] + + for algo_name in algos_to_generate: + + # generate config for this experiment + config, json_path = generate_experiment_config( + base_exp_name="core", + base_config_dir=base_config_dir, + base_dataset_dir=base_dataset_dir, + base_output_dir=base_output_dir, + algo_name=algo_name, + algo_config_modifier=algo_to_config_modifier[algo_name], + task_name=task, + dataset_type=dataset_type, + hdf5_type=hdf5_type, + ) + + # save json path into dict + core_json_paths[task][dataset_type][hdf5_type][algo_name] = json_path + + return core_json_paths + + +def generate_subopt_configs( + base_config_dir, + base_dataset_dir, + base_output_dir, + algo_to_config_modifier, +): + """ + Helper function to generate all configs for the suboptimal human subsets of the multi-human datasets. + Note that while the paper includes the results on the can-paired dataset along with results on these + datasets, the configs for runs on the can-paired dataset is in the "core" set of runs. + + Args: + base_config_dir (str): base directory to place generated configs + + base_dataset_dir (str): path to directory where datasets are on disk. + Directory structure is expected to be consistent with the output + of @make_dataset_dirs in the download_datasets.py script. + + base_output_dir (str): directory to save training results to. If None, will use the directory + from the default algorithm configs. + + algo_to_config_modifier (dict): dictionary that maps algo name to a function that modifies configs + to add algo hyperparameter settings, given the task, dataset, and hdf5 types. + """ + subopt_json_paths = Config() # use for convenient nested dict + for task in ["lift", "can", "square", "transport"]: + # only generate configs for multi-human data subsets + for dataset_type in ["mh"]: + # only low-dim / image + for hdf5_type in ["low_dim", "image"]: + + # get list of algorithms to generate configs for, for this hdf5 dataset + algos_to_generate = ["bc", "bc_rnn", "bcq", "cql", "hbc", "iris"] + if hdf5_type == "image": + # no hbc or iris for image runs + algos_to_generate = algos_to_generate[:-2] + + for algo_name in algos_to_generate: + + for fk in ["worse", "okay", "better", "worse_okay", "worse_better", "okay_better"]: + + # generate config for this experiment + config, json_path = generate_experiment_config( + base_exp_name="subopt", + base_config_dir=base_config_dir, + base_dataset_dir=base_dataset_dir, + base_output_dir=base_output_dir, + algo_name=algo_name, + algo_config_modifier=algo_to_config_modifier[algo_name], + task_name=task, + dataset_type=dataset_type, + hdf5_type=hdf5_type, + filter_key=fk, + ) + + # save json path into dict + dataset_type_dir = "{}/{}".format(dataset_type, fk) + subopt_json_paths[task][dataset_type_dir][hdf5_type][algo_name] = json_path + + return subopt_json_paths + + +def generate_dataset_size_configs( + base_config_dir, + base_dataset_dir, + base_output_dir, + algo_to_config_modifier, +): + """ + Helper function to generate all configs for the dataset size ablation experiments, where BC-RNN models + were trained on 20% and 50% dataset sizes. + + Args: + base_config_dir (str): base directory to place generated configs + + base_dataset_dir (str): path to directory where datasets are on disk. + Directory structure is expected to be consistent with the output + of @make_dataset_dirs in the download_datasets.py script. + + base_output_dir (str): directory to save training results to. If None, will use the directory + from the default algorithm configs. + + algo_to_config_modifier (dict): dictionary that maps algo name to a function that modifies configs + to add algo hyperparameter settings, given the task, dataset, and hdf5 types. + """ + size_ablation_json_paths = Config() # use for convenient nested dict + for task in ["lift", "can", "square", "transport"]: + for dataset_type in ["ph", "mh"]: + for hdf5_type in ["low_dim", "image"]: + + # only bc-rnn + algo_name = "bc_rnn" + for fk in ["20_percent", "50_percent"]: + + # generate config for this experiment + config, json_path = generate_experiment_config( + base_exp_name="dataset_size", + base_config_dir=base_config_dir, + base_dataset_dir=base_dataset_dir, + base_output_dir=base_output_dir, + algo_name=algo_name, + algo_config_modifier=algo_to_config_modifier[algo_name], + task_name=task, + dataset_type=dataset_type, + hdf5_type=hdf5_type, + filter_key=fk, + ) + + # save json path into dict + dataset_type_dir = "{}/{}".format(dataset_type, fk) + size_ablation_json_paths[task][dataset_type_dir][hdf5_type][algo_name] = json_path + + return size_ablation_json_paths + + +def generate_obs_ablation_configs( + base_config_dir, + base_dataset_dir, + base_output_dir, + algo_to_config_modifier, +): + """ + Helper function to generate all configs for the observation ablation experiments, where BC and BC-RNN models + were trained on different versions of low-dim and image observations. + + Args: + base_config_dir (str): base directory to place generated configs + + base_dataset_dir (str): path to directory where datasets are on disk. + Directory structure is expected to be consistent with the output + of @make_dataset_dirs in the download_datasets.py script. + + base_output_dir (str): directory to save training results to. If None, will use the directory + from the default algorithm configs. + + algo_to_config_modifier (dict): dictionary that maps algo name to a function that modifies configs + to add algo hyperparameter settings, given the task, dataset, and hdf5 types. + """ + + # observation config modifiers for these experiments + def add_eef_vel(config): + with config.observation.values_unlocked(): + old_low_dim_mods = list(config.observation.modalities.obs.low_dim) + old_low_dim_mods.extend(["robot0_eef_vel_lin", "robot0_eef_vel_ang", "robot0_gripper_qvel"]) + if "robot1_eef_pos" in old_low_dim_mods: + old_low_dim_mods.extend(["robot1_eef_vel_lin", "robot1_eef_vel_ang", "robot1_gripper_qvel"]) + config.observation.modalities.obs.low_dim = old_low_dim_mods + return config + + def add_proprio(config): + with config.observation.values_unlocked(): + old_low_dim_mods = list(config.observation.modalities.obs.low_dim) + old_low_dim_mods.extend(["robot0_joint_pos_cos", "robot0_joint_pos_sin", "robot0_joint_vel"]) + if "robot1_eef_pos" in old_low_dim_mods: + old_low_dim_mods.extend(["robot1_joint_pos_cos", "robot1_joint_pos_sin", "robot1_joint_vel"]) + config.observation.modalities.obs.low_dim = old_low_dim_mods + return config + + def remove_wrist(config): + with config.observation.values_unlocked(): + old_image_mods = list(config.observation.modalities.obs.rgb) + config.observation.modalities.obs.rgb = [m for m in old_image_mods if "eye_in_hand" not in m] + return config + + def remove_rand(config): + with config.observation.values_unlocked(): + config.observation.encoder.rgb.obs_randomizer_class = None + return config + + obs_ablation_json_paths = Config() # use for convenient nested dict + for task in ["square", "transport"]: + for dataset_type in ["ph", "mh"]: + for hdf5_type in ["low_dim", "image"]: + + # observation modifiers to apply + if hdf5_type == "low_dim": + obs_modifiers = [add_eef_vel, add_proprio] + else: + obs_modifiers = [add_eef_vel, add_proprio, remove_wrist, remove_rand] + + # only bc and bc-rnn + algos_to_generate = ["bc", "bc_rnn"] + for algo_name in algos_to_generate: + for obs_modifier in obs_modifiers: + # generate config for this experiment + config, json_path = generate_experiment_config( + base_exp_name="obs_ablation", + base_config_dir=base_config_dir, + base_dataset_dir=base_dataset_dir, + base_output_dir=base_output_dir, + algo_name=algo_name, + algo_config_modifier=algo_to_config_modifier[algo_name], + task_name=task, + dataset_type=dataset_type, + hdf5_type=hdf5_type, + additional_name=obs_modifier.__name__, + additional_config_modifier=obs_modifier, + ) + + # save json path into dict + algo_name_str = "{}_{}".format(algo_name, obs_modifier.__name__) + obs_ablation_json_paths[task][dataset_type][hdf5_type][algo_name_str] = json_path + + return obs_ablation_json_paths + + +def generate_hyper_ablation_configs( + base_config_dir, + base_dataset_dir, + base_output_dir, + algo_to_config_modifier, +): + """ + Helper function to generate all configs for the hyperparameter sensitivity experiments, + where BC-RNN models were trained on different ablations. + + Args: + base_config_dir (str): base directory to place generated configs + + base_dataset_dir (str): path to directory where datasets are on disk. + Directory structure is expected to be consistent with the output + of @make_dataset_dirs in the download_datasets.py script. + + base_output_dir (str): directory to save training results to. If None, will use the directory + from the default algorithm configs. + + algo_to_config_modifier (dict): dictionary that maps algo name to a function that modifies configs + to add algo hyperparameter settings, given the task, dataset, and hdf5 types. + """ + + # observation config modifiers for these experiments + def change_lr(config): + with config.algo.values_unlocked(): + config.algo.optim_params.policy.learning_rate.initial = 1e-3 + return config + + def change_gmm(config): + with config.algo.values_unlocked(): + config.algo.gmm.enabled = False + return config + + def change_mlp(config): + with config.algo.values_unlocked(): + config.algo.actor_layer_dims = (1024, 1024) + return config + + def change_conv(config): + with config.observation.values_unlocked(): + config.observation.encoder.rgb.core_class = 'ShallowConv' + config.observation.encoder.rgb.core_kwargs = Config() + return config + + def change_rnnd_low_dim(config): + with config.algo.values_unlocked(): + config.algo.rnn.hidden_dim = 100 + return config + + def change_rnnd_image(config): + with config.algo.values_unlocked(): + config.algo.rnn.hidden_dim = 400 + return config + + hyper_ablation_json_paths = Config() # use for convenient nested dict + for task in ["square", "transport"]: + for dataset_type in ["ph", "mh"]: + for hdf5_type in ["low_dim", "image"]: + + # observation modifiers to apply + if hdf5_type == "low_dim": + hyper_modifiers = [change_lr, change_gmm, change_mlp, change_rnnd_low_dim] + else: + hyper_modifiers = [change_lr, change_gmm, change_conv, change_rnnd_image] + + # only bc and bc-rnn + algo_name = "bc_rnn" + for hyper_modifier in hyper_modifiers: + # generate config for this experiment + config, json_path = generate_experiment_config( + base_exp_name="hyper_ablation", + base_config_dir=base_config_dir, + base_dataset_dir=base_dataset_dir, + base_output_dir=base_output_dir, + algo_name=algo_name, + algo_config_modifier=algo_to_config_modifier[algo_name], + task_name=task, + dataset_type=dataset_type, + hdf5_type=hdf5_type, + additional_name=hyper_modifier.__name__, + additional_config_modifier=hyper_modifier, + ) + + # save json path into dict + algo_name_str = "{}_{}".format(algo_name, hyper_modifier.__name__) + hyper_ablation_json_paths[task][dataset_type][hdf5_type][algo_name_str] = json_path + + return hyper_ablation_json_paths + + +def generate_d4rl_configs( + base_config_dir, + base_dataset_dir, + base_output_dir, + algo_to_config_modifier, +): + """ + Helper function to generate all configs for reproducing BCQ, CQL, and TD3-BC runs on some D4RL + environments. + + Args: + base_config_dir (str): base directory to place generated configs + + base_dataset_dir (str): path to directory where datasets are on disk. + Directory structure is expected to be consistent with the output + of @make_dataset_dirs in the download_datasets.py script. + + base_output_dir (str): directory to save training results to. If None, will use the directory + from the default algorithm configs. + + algo_to_config_modifier (dict): dictionary that maps algo name to a function that modifies configs + to add algo hyperparameter settings, given the task, dataset, and hdf5 types. + """ + + def bcq_algo_config_modifier(config): + with config.algo.values_unlocked(): + # all LRs 1e-3, enable actor + config.algo.optim_params.critic.learning_rate.initial = 1e-3 + config.algo.optim_params.action_sampler.learning_rate.initial = 1e-3 + config.algo.optim_params.actor.learning_rate.initial = 1e-3 + config.algo.actor.enabled = True + config.algo.action_sampler.vae.kl_weight = 0.5 + return config + + def cql_algo_config_modifier(config): + with config.algo.values_unlocked(): + # taken from TD3-BC settings described in their paper + config.algo.optim_params.critic.learning_rate.initial = 3e-4 + config.algo.optim_params.actor.learning_rate.initial = 3e-5 + config.algo.actor.bc_start_steps = 40000 # pre-training steps for actor + config.algo.critic.target_q_gap = None # no Lagrange, and fixed weight of 10.0 + config.algo.critic.cql_weight = 10.0 + config.algo.critic.min_q_weight = 1.0 + config.algo.critic.deterministic_backup = True # deterministic backup (no entropy in Q-target) + config.algo.actor.layer_dims = (256, 256, 256) # MLP sizes + config.algo.critic.layer_dims = (256, 256, 256) + return config + + def iql_algo_config_modifier(config): + with config.algo.values_unlocked(): + # taken from IQL settings described in their paper + config.algo.target_tau = 0.005 + config.algo.vf_quantile = 0.7 + config.algo.adv.beta = 3.0 + config.algo.optim_params.critic.learning_rate.initial = 3e-4 + config.algo.optim_params.vf.learning_rate.initial = 3e-4 + config.algo.optim_params.actor.learning_rate.initial = 3e-4 + config.algo.actor.layer_dims = (256, 256, 256) # MLP sizes + config.algo.critic.layer_dims = (256, 256, 256) + return config + + d4rl_tasks = [ + # "halfcheetah-random-v2", + # "hopper-random-v2", + # "walker2d-random-v2", + "halfcheetah-medium-v2", + "hopper-medium-v2", + "walker2d-medium-v2", + "halfcheetah-expert-v2", + "hopper-expert-v2", + "walker2d-expert-v2", + "halfcheetah-medium-expert-v2", + "hopper-medium-expert-v2", + "walker2d-medium-expert-v2", + # "halfcheetah-medium-replay-v2", + # "hopper-medium-replay-v2", + # "walker2d-medium-replay-v2", + ] + d4rl_json_paths = Config() # use for convenient nested dict + for task_name in d4rl_tasks: + for algo_name in ["bcq", "cql", "td3_bc", "iql"]: + config = config_factory(algo_name=algo_name) + + # hack: copy experiment and train sections from td3-bc, since that has defaults for training with D4RL + if algo_name != "td3_bc": + ref_config = config_factory(algo_name="td3_bc") + with config.values_unlocked(): + config.experiment = ref_config.experiment + config.train = ref_config.train + config.observation = ref_config.observation + config.train.hdf5_normalize_obs = False # only TD3-BC uses observation normalization + + # modify algo section for d4rl defaults + if algo_name == "bcq": + config = bcq_algo_config_modifier(config) + elif algo_name == "cql": + config = cql_algo_config_modifier(config) + elif algo_name == "iql": + config = iql_algo_config_modifier(config) + + # set experiment name + with config.experiment.values_unlocked(): + config.experiment.name = "{}_{}_{}".format("d4rl", algo_name, task_name) + # set output folder and dataset + with config.train.values_unlocked(): + if base_output_dir is None: + base_output_dir_for_algo = "../{}_trained_models".format(algo_name) + else: + base_output_dir_for_algo = base_output_dir + config.train.output_dir = os.path.join(base_output_dir_for_algo, "d4rl", algo_name, task_name, "trained_models") + config.train.data = os.path.join(base_dataset_dir, "d4rl", "converted", + "{}.hdf5".format(task_name.replace("-", "_"))) + + # save config to json file + dir_to_save = os.path.join(base_config_dir, "d4rl", task_name) + os.makedirs(dir_to_save, exist_ok=True) + json_path = os.path.join(dir_to_save, "{}.json".format(algo_name)) + config.dump(filename=json_path) + + # save json path into dict + d4rl_json_paths[task_name][""][""][algo_name] = json_path + + return d4rl_json_paths + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + # Directory where generated configs will be placed + parser.add_argument( + "--config_dir", + type=str, + default=None, + help="Directory where generated configs will be placed. Defaults to 'paper' subfolder in exps folder of repository", + ) + + # directory where released datasets are located + parser.add_argument( + "--dataset_dir", + type=str, + default=None, + help="Base dataset directory for released datasets. Defaults to datasets folder in repository.", + ) + + # output directory for training runs (will be written to configs) + parser.add_argument( + "--output_dir", + type=str, + default=None, + help="Base output directory for all training runs that will be written to generated configs.", + ) + + args = parser.parse_args() + + # read args + generated_configs_base_dir = args.config_dir + if generated_configs_base_dir is None: + generated_configs_base_dir = os.path.join(robomimic.__path__[0], "exps/paper") + + datasets_base_dir = args.dataset_dir + if datasets_base_dir is None: + datasets_base_dir = os.path.join(robomimic.__path__[0], "../datasets") + + output_base_dir = args.output_dir + + # algo to modifier + algo_to_modifier = dict( + bc=modify_bc_config_for_dataset, + bc_rnn=modify_bc_rnn_config_for_dataset, + bcq=modify_bcq_config_for_dataset, + cql=modify_cql_config_for_dataset, + hbc=modify_hbc_config_for_dataset, + iris=modify_iris_config_for_dataset, + ) + + # exp name to config generator + exp_name_to_generator = dict( + core=generate_core_configs, + subopt=generate_subopt_configs, + dataset_size=generate_dataset_size_configs, + obs_ablation=generate_obs_ablation_configs, + hyper_ablation=generate_hyper_ablation_configs, + d4rl=generate_d4rl_configs, + ) + + # generate configs for each experiment name + config_json_paths = Config() # use for convenient nested dict + for exp_name in exp_name_to_generator: + config_json_paths[exp_name] = exp_name_to_generator[exp_name]( + base_config_dir=generated_configs_base_dir, + base_dataset_dir=datasets_base_dir, + base_output_dir=output_base_dir, + algo_to_config_modifier=algo_to_modifier, + ) + + # write output shell scripts + for exp_name in config_json_paths: + shell_path = os.path.join(generated_configs_base_dir, "{}.sh".format(exp_name)) + with open(shell_path, "w") as f: + f.write("#!/bin/bash\n\n") + f.write("# " + "=" * 10 + exp_name + "=" * 10 + "\n") + train_script_loc = os.path.join(robomimic.__path__[0], "scripts/train.py") + + for task in config_json_paths[exp_name]: + for dataset_type in config_json_paths[exp_name][task]: + for hdf5_type in config_json_paths[exp_name][task][dataset_type]: + f.write("\n") + f.write("# task: {}\n".format(task)) + if len(dataset_type) > 0: + f.write("# dataset type: {}\n".format(dataset_type)) + if len(hdf5_type) > 0: + f.write("# hdf5 type: {}\n".format(hdf5_type)) + for algo_name in config_json_paths[exp_name][task][dataset_type][hdf5_type]: + # f.write("# {}\n".format(algo_name)) + exp_json_path = config_json_paths[exp_name][task][dataset_type][hdf5_type][algo_name] + cmd = "python {} --config {}\n".format(train_script_loc, exp_json_path) + f.write(cmd) + f.write("\n") diff --git a/aloha-devel/robomimic/scripts/get_dataset_info.py b/aloha-devel/robomimic/scripts/get_dataset_info.py new file mode 100644 index 0000000000000000000000000000000000000000..9349ed8aaa15823b662d582ef116d78067c3f559 --- /dev/null +++ b/aloha-devel/robomimic/scripts/get_dataset_info.py @@ -0,0 +1,134 @@ +""" +Helper script to report dataset information. By default, will print trajectory length statistics, +the maximum and minimum action element in the dataset, filter keys present, environment +metadata, and the structure of the first demonstration. If --verbose is passed, it will +report the exact demo keys under each filter key, and the structure of all demonstrations +(not just the first one). + +Args: + dataset (str): path to hdf5 dataset + + filter_key (str): if provided, report statistics on the subset of trajectories + in the file that correspond to this filter key + + verbose (bool): if flag is provided, print more details, like the structure of all + demonstrations (not just the first one) + +Example usage: + + # run script on example hdf5 packaged with repository + python get_dataset_info.py --dataset ../../tests/assets/test.hdf5 + + # run script only on validation data + python get_dataset_info.py --dataset ../../tests/assets/test.hdf5 --filter_key valid +""" +import h5py +import json +import argparse +import numpy as np + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--dataset", + type=str, + help="path to hdf5 dataset", + ) + parser.add_argument( + "--filter_key", + type=str, + default=None, + help="(optional) if provided, report statistics on the subset of trajectories \ + in the file that correspond to this filter key", + ) + parser.add_argument( + "--verbose", + action='store_true', + help="verbose output", + ) + args = parser.parse_args() + + # extract demonstration list from file + filter_key = args.filter_key + all_filter_keys = None + f = h5py.File(args.dataset, "r") + if filter_key is not None: + # use the demonstrations from the filter key instead + print("NOTE: using filter key {}".format(filter_key)) + demos = sorted([elem.decode("utf-8") for elem in np.array(f["mask/{}".format(filter_key)])]) + else: + # use all demonstrations + demos = sorted(list(f["data"].keys())) + + # extract filter key information + if "mask" in f: + all_filter_keys = {} + for fk in f["mask"]: + fk_demos = sorted([elem.decode("utf-8") for elem in np.array(f["mask/{}".format(fk)])]) + all_filter_keys[fk] = fk_demos + + # put demonstration list in increasing episode order + inds = np.argsort([int(elem[5:]) for elem in demos]) + demos = [demos[i] for i in inds] + + # extract length of each trajectory in the file + traj_lengths = [] + action_min = np.inf + action_max = -np.inf + for ep in demos: + traj_lengths.append(f["data/{}/actions".format(ep)].shape[0]) + action_min = min(action_min, np.min(f["data/{}/actions".format(ep)][()])) + action_max = max(action_max, np.max(f["data/{}/actions".format(ep)][()])) + traj_lengths = np.array(traj_lengths) + + # report statistics on the data + print("") + print("total transitions: {}".format(np.sum(traj_lengths))) + print("total trajectories: {}".format(traj_lengths.shape[0])) + print("traj length mean: {}".format(np.mean(traj_lengths))) + print("traj length std: {}".format(np.std(traj_lengths))) + print("traj length min: {}".format(np.min(traj_lengths))) + print("traj length max: {}".format(np.max(traj_lengths))) + print("action min: {}".format(action_min)) + print("action max: {}".format(action_max)) + print("") + print("==== Filter Keys ====") + if all_filter_keys is not None: + for fk in all_filter_keys: + print("filter key {} with {} demos".format(fk, len(all_filter_keys[fk]))) + else: + print("no filter keys") + print("") + if args.verbose: + if all_filter_keys is not None: + print("==== Filter Key Contents ====") + for fk in all_filter_keys: + print("filter_key {} with {} demos: {}".format(fk, len(all_filter_keys[fk]), all_filter_keys[fk])) + print("") + env_meta = json.loads(f["data"].attrs["env_args"]) + print("==== Env Meta ====") + print(json.dumps(env_meta, indent=4)) + print("") + + print("==== Dataset Structure ====") + for ep in demos: + print("episode {} with {} transitions".format(ep, f["data/{}".format(ep)].attrs["num_samples"])) + for k in f["data/{}".format(ep)]: + if k in ["obs", "next_obs"]: + print(" key: {}".format(k)) + for obs_k in f["data/{}/{}".format(ep, k)]: + shape = f["data/{}/{}/{}".format(ep, k, obs_k)].shape + print(" observation key {} with shape {}".format(obs_k, shape)) + elif isinstance(f["data/{}/{}".format(ep, k)], h5py.Dataset): + key_shape = f["data/{}/{}".format(ep, k)].shape + print(" key: {} with shape {}".format(k, key_shape)) + + if not args.verbose: + break + + f.close() + + # maybe display error message + print("") + if (action_min < -1.) or (action_max > 1.): + raise Exception("Dataset should have actions in [-1., 1.] but got bounds [{}, {}]".format(action_min, action_max)) diff --git a/aloha-devel/robomimic/scripts/hyperparam_helper.py b/aloha-devel/robomimic/scripts/hyperparam_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..870c739ecbf4a751b6e62c363555d451cfa68ae2 --- /dev/null +++ b/aloha-devel/robomimic/scripts/hyperparam_helper.py @@ -0,0 +1,141 @@ +""" +A useful script for generating json files and shell scripts for conducting parameter scans. +The script takes a path to a base json file as an argument and a shell file name. +It generates a set of new json files in the same folder as the base json file, and +a shell file script that contains commands to run for each experiment. + +Instructions: + +(1) Start with a base json that specifies a complete set of parameters for a single + run. This only needs to include parameters you want to sweep over, and parameters + that are different from the defaults. You can set this file path by either + passing it as an argument (e.g. --config /path/to/base.json) or by directly + setting the config file in @make_generator. The new experiment jsons will be put + into the same directory as the base json. + +(2) Decide on what json parameters you would like to sweep over, and fill those in as + keys in @make_generator below, taking note of the hierarchical key + formatting using "/" or ".". Fill in corresponding values for each - these will + be used in creating the experiment names, and for determining the range + of values to sweep. Parameters that should be sweeped together should + be assigned the same group number. + +(3) Set the output script name by either passing it as an argument (e.g. --script /path/to/script.sh) + or by directly setting the script file in @make_generator. The script to run all experiments + will be created at the specified path. + +Args: + config (str): path to a base config json file that will be modified to generate config jsons. + The jsons will be generated in the same folder as this file. + + script (str): path to output script that contains commands to run the generated training runs + +Example usage: + + # assumes that /tmp/gen_configs/base.json has already been created (see quickstart section of docs for an example) + python hyperparam_helper.py --config /tmp/gen_configs/base.json --script /tmp/gen_configs/out.sh +""" +import argparse + +import robomimic +import robomimic.utils.hyperparam_utils as HyperparamUtils + + +def make_generator(config_file, script_file): + """ + Implement this function to setup your own hyperparameter scan! + """ + generator = HyperparamUtils.ConfigGenerator( + base_config_file=config_file, script_file=script_file + ) + + # use RNN with horizon 10 + generator.add_param( + key="algo.rnn.enabled", + name="", + group=0, + values=[True], + ) + generator.add_param( + key="train.seq_length", + name="", + group=0, + values=[10], + ) + generator.add_param( + key="algo.rnn.horizon", + name="", + group=0, + values=[10], + ) + + # LR - 1e-3, 1e-4 + generator.add_param( + key="algo.optim_params.policy.learning_rate.initial", + name="plr", + group=1, + values=[1e-3, 1e-4], + ) + + # GMM y / n + generator.add_param( + key="algo.gmm.enabled", + name="gmm", + group=2, + values=[True, False], + value_names=["t", "f"], + ) + + # RNN dim 400 + MLP dims (1024, 1024) vs. RNN dim 1000 + empty MLP dims () + generator.add_param( + key="algo.rnn.hidden_dim", + name="rnnd", + group=3, + values=[ + 400, + 1000, + ], + ) + generator.add_param( + key="algo.actor_layer_dims", + name="mlp", + group=3, + values=[ + [1024, 1024], + [], + ], + value_names=["1024", "0"], + ) + + return generator + + +def main(args): + + # make config generator + generator = make_generator(config_file=args.config, script_file=args.script) + + # generate jsons and script + generator.generate() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + # Path to base json config - will override any defaults. + parser.add_argument( + "--config", + type=str, + help="path to base config json that will be modified to generate jsons. The jsons will\ + be generated in the same folder as this file.", + ) + + # Script name to generate - will override any defaults + parser.add_argument( + "--script", + type=str, + help="path to output script that contains commands to run the generated training runs", + ) + + args = parser.parse_args() + main(args) diff --git a/aloha-devel/robomimic/scripts/playback_dataset.py b/aloha-devel/robomimic/scripts/playback_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..1858d99fccc2ba17fe097e63d5d909251b885f59 --- /dev/null +++ b/aloha-devel/robomimic/scripts/playback_dataset.py @@ -0,0 +1,372 @@ +""" +A script to visualize dataset trajectories by loading the simulation states +one by one or loading the first state and playing actions back open-loop. +The script can generate videos as well, by rendering simulation frames +during playback. The videos can also be generated using the image observations +in the dataset (this is useful for real-robot datasets) by using the +--use-obs argument. + +Args: + dataset (str): path to hdf5 dataset + + filter_key (str): if provided, use the subset of trajectories + in the file that correspond to this filter key + + n (int): if provided, stop after n trajectories are processed + + use-obs (bool): if flag is provided, visualize trajectories with dataset + image observations instead of simulator + + use-actions (bool): if flag is provided, use open-loop action playback + instead of loading sim states + + render (bool): if flag is provided, use on-screen rendering during playback + + video_path (str): if provided, render trajectories to this video file path + + video_skip (int): render frames to a video every @video_skip steps + + render_image_names (str or [str]): camera name(s) / image observation(s) to + use for rendering on-screen or to video + + first (bool): if flag is provided, use first frame of each episode for playback + instead of the entire episode. Useful for visualizing task initializations. + +Example usage below: + + # force simulation states one by one, and render agentview and wrist view cameras to video + python playback_dataset.py --dataset /path/to/dataset.hdf5 \ + --render_image_names agentview robot0_eye_in_hand \ + --video_path /tmp/playback_dataset.mp4 + + # playback the actions in the dataset, and render agentview camera during playback to video + python playback_dataset.py --dataset /path/to/dataset.hdf5 \ + --use-actions --render_image_names agentview \ + --video_path /tmp/playback_dataset_with_actions.mp4 + + # use the observations stored in the dataset to render videos of the dataset trajectories + python playback_dataset.py --dataset /path/to/dataset.hdf5 \ + --use-obs --render_image_names agentview_image \ + --video_path /tmp/obs_trajectory.mp4 + + # visualize initial states in the demonstration data + python playback_dataset.py --dataset /path/to/dataset.hdf5 \ + --first --render_image_names agentview \ + --video_path /tmp/dataset_task_inits.mp4 +""" + +import os +import json +import h5py +import argparse +import imageio +import numpy as np +import random + +import robomimic +import robomimic.utils.obs_utils as ObsUtils +import robomimic.utils.env_utils as EnvUtils +import robomimic.utils.file_utils as FileUtils +from robomimic.envs.env_base import EnvBase, EnvType + + +# Define default cameras to use for each env type +DEFAULT_CAMERAS = { + EnvType.ROBOSUITE_TYPE: ["agentview"], + EnvType.IG_MOMART_TYPE: ["rgb"], + EnvType.GYM_TYPE: ValueError("No camera names supported for gym type env!"), +} + + +def playback_trajectory_with_env( + env, + initial_state, + states, + actions=None, + render=False, + video_writer=None, + video_skip=5, + camera_names=None, + first=False, +): + """ + Helper function to playback a single trajectory using the simulator environment. + If @actions are not None, it will play them open-loop after loading the initial state. + Otherwise, @states are loaded one by one. + + Args: + env (instance of EnvBase): environment + initial_state (dict): initial simulation state to load + states (np.array): array of simulation states to load + actions (np.array): if provided, play actions back open-loop instead of using @states + render (bool): if True, render on-screen + video_writer (imageio writer): video writer + video_skip (int): determines rate at which environment frames are written to video + camera_names (list): determines which camera(s) are used for rendering. Pass more than + one to output a video with multiple camera views concatenated horizontally. + first (bool): if True, only use the first frame of each episode. + """ + assert isinstance(env, EnvBase) + + write_video = (video_writer is not None) + video_count = 0 + assert not (render and write_video) + + # load the initial state + ## this reset call doesn't seem necessary. + ## seems ok to remove but haven't fully tested it. + ## removing for now + # env.reset() + env.reset_to(initial_state) + + traj_len = states.shape[0] + action_playback = (actions is not None) + if action_playback: + assert states.shape[0] == actions.shape[0] + + for i in range(traj_len): + if action_playback: + env.step(actions[i]) + if i < traj_len - 1: + # check whether the actions deterministically lead to the same recorded states + state_playback = env.get_state()["states"] + if not np.all(np.equal(states[i + 1], state_playback)): + err = np.linalg.norm(states[i + 1] - state_playback) + print("warning: playback diverged by {} at step {}".format(err, i)) + else: + env.reset_to({"states" : states[i]}) + + # on-screen render + if render: + env.render(mode="human", camera_name=camera_names[0]) + + # video render + if write_video: + if video_count % video_skip == 0: + video_img = [] + for cam_name in camera_names: + video_img.append(env.render(mode="rgb_array", height=512, width=512, camera_name=cam_name)) + video_img = np.concatenate(video_img, axis=1) # concatenate horizontally + video_writer.append_data(video_img) + video_count += 1 + + if first: + break + + +def playback_trajectory_with_obs( + traj_grp, + video_writer, + video_skip=5, + image_names=None, + first=False, +): + """ + This function reads all "rgb" observations in the dataset trajectory and + writes them into a video. + + Args: + traj_grp (hdf5 file group): hdf5 group which corresponds to the dataset trajectory to playback + video_writer (imageio writer): video writer + video_skip (int): determines rate at which environment frames are written to video + image_names (list): determines which image observations are used for rendering. Pass more than + one to output a video with multiple image observations concatenated horizontally. + first (bool): if True, only use the first frame of each episode. + """ + assert image_names is not None, "error: must specify at least one image observation to use in @image_names" + video_count = 0 + + traj_len = traj_grp["actions"].shape[0] + for i in range(traj_len): + if video_count % video_skip == 0: + # concatenate image obs together + im = [traj_grp["obs/{}".format(k)][i] for k in image_names] + frame = np.concatenate(im, axis=1) + video_writer.append_data(frame) + video_count += 1 + + if first: + break + + +def playback_dataset(args): + # some arg checking + write_video = (args.video_path is not None) + assert not (args.render and write_video) # either on-screen or video but not both + + # Auto-fill camera rendering info if not specified + if args.render_image_names is None: + # We fill in the automatic values + env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=args.dataset) + env_type = EnvUtils.get_env_type(env_meta=env_meta) + args.render_image_names = DEFAULT_CAMERAS[env_type] + + if args.render: + # on-screen rendering can only support one camera + assert len(args.render_image_names) == 1 + + if args.use_obs: + assert write_video, "playback with observations can only write to video" + assert not args.use_actions, "playback with observations is offline and does not support action playback" + + # create environment only if not playing back with observations + if not args.use_obs: + # need to make sure ObsUtils knows which observations are images, but it doesn't matter + # for playback since observations are unused. Pass a dummy spec here. + dummy_spec = dict( + obs=dict( + low_dim=["robot0_eef_pos"], + rgb=[], + ), + ) + ObsUtils.initialize_obs_utils_with_obs_specs(obs_modality_specs=dummy_spec) + + env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=args.dataset) + env = EnvUtils.create_env_from_metadata(env_meta=env_meta, render=args.render, render_offscreen=write_video) + + # some operations for playback are robosuite-specific, so determine if this environment is a robosuite env + is_robosuite_env = EnvUtils.is_robosuite_env(env_meta) + + f = h5py.File(args.dataset, "r") + + # list of all demonstration episodes (sorted in increasing number order) + if args.filter_key is not None: + print("using filter key: {}".format(args.filter_key)) + demos = [elem.decode("utf-8") for elem in np.array(f["mask/{}".format(args.filter_key)])] + else: + demos = list(f["data"].keys()) + inds = np.argsort([int(elem[5:]) for elem in demos]) + demos = [demos[i] for i in inds] + + # maybe reduce the number of demonstrations to playback + if args.n is not None: + # if not args.dont_shuffle_demos: + # random.shuffle(demos) + random.shuffle(demos) + demos = demos[:args.n] + + # maybe dump video + video_writer = None + if write_video: + video_writer = imageio.get_writer(args.video_path, fps=20) + + for ind in range(len(demos)): + ep = demos[ind] + print("Playing back episode: {}".format(ep)) + + if args.use_obs: + playback_trajectory_with_obs( + traj_grp=f["data/{}".format(ep)], + video_writer=video_writer, + video_skip=args.video_skip, + image_names=args.render_image_names, + first=args.first, + ) + continue + + # prepare initial state to reload from + states = f["data/{}/states".format(ep)][()] + initial_state = dict(states=states[0]) + if is_robosuite_env: + initial_state["model"] = f["data/{}".format(ep)].attrs["model_file"] + initial_state["ep_meta"] = f["data/{}".format(ep)].attrs.get("ep_meta", None) + + # supply actions if using open-loop action playback + actions = None + if args.use_actions: + actions = f["data/{}/actions".format(ep)][()] + + playback_trajectory_with_env( + env=env, + initial_state=initial_state, + states=states, actions=actions, + render=args.render, + video_writer=video_writer, + video_skip=args.video_skip, + camera_names=args.render_image_names, + first=args.first, + ) + + f.close() + if write_video: + video_writer.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--dataset", + type=str, + help="path to hdf5 dataset", + ) + parser.add_argument( + "--filter_key", + type=str, + default=None, + help="(optional) filter key, to select a subset of trajectories in the file", + ) + + # number of trajectories to playback. If omitted, playback all of them. + parser.add_argument( + "--n", + type=int, + default=None, + help="(optional) stop after n trajectories are played", + ) + + # Use image observations instead of doing playback using the simulator env. + parser.add_argument( + "--use-obs", + action='store_true', + help="visualize trajectories with dataset image observations instead of simulator", + ) + + # Playback stored dataset actions open-loop instead of loading from simulation states. + parser.add_argument( + "--use-actions", + action='store_true', + help="use open-loop action playback instead of loading sim states", + ) + + # Whether to render playback to screen + parser.add_argument( + "--render", + action='store_true', + help="on-screen rendering", + ) + + # Dump a video of the dataset playback to the specified path + parser.add_argument( + "--video_path", + type=str, + default=None, + help="(optional) render trajectories to this video file path", + ) + + # How often to write video frames during the playback + parser.add_argument( + "--video_skip", + type=int, + default=5, + help="render frames to video every n steps", + ) + + # camera names to render, or image observations to use for writing to video + parser.add_argument( + "--render_image_names", + type=str, + nargs='+', + default=None, + help="(optional) camera name(s) / image observation(s) to use for rendering on-screen or to video. Default is" + "None, which corresponds to a predefined camera for each env type", + ) + + # Only use the first frame of each episode + parser.add_argument( + "--first", + action='store_true', + help="use first frame of each episode", + ) + + args = parser.parse_args() + playback_dataset(args) diff --git a/aloha-devel/robomimic/scripts/run_trained_agent.py b/aloha-devel/robomimic/scripts/run_trained_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..2ec09dfe981db7ac46941b5dc0da355041835281 --- /dev/null +++ b/aloha-devel/robomimic/scripts/run_trained_agent.py @@ -0,0 +1,375 @@ +""" +The main script for evaluating a policy in an environment. + +Args: + agent (str): path to saved checkpoint pth file + + horizon (int): if provided, override maximum horizon of rollout from the one + in the checkpoint + + env (str): if provided, override name of env from the one in the checkpoint, + and use it for rollouts + + render (bool): if flag is provided, use on-screen rendering during rollouts + + video_path (str): if provided, render trajectories to this video file path + + video_skip (int): render frames to a video every @video_skip steps + + camera_names (str or [str]): camera name(s) to use for rendering on-screen or to video + + dataset_path (str): if provided, an hdf5 file will be written at this path with the + rollout data + + dataset_obs (bool): if flag is provided, and @dataset_path is provided, include + possible high-dimensional observations in output dataset hdf5 file (by default, + observations are excluded and only simulator states are saved). + + seed (int): if provided, set seed for rollouts + +Example usage: + + # Evaluate a policy with 50 rollouts of maximum horizon 400 and save the rollouts to a video. + # Visualize the agentview and wrist cameras during the rollout. + + python run_trained_agent.py --agent /path/to/model.pth \ + --n_rollouts 50 --horizon 400 --seed 0 \ + --video_path /path/to/output.mp4 \ + --camera_names agentview robot0_eye_in_hand + + # Write the 50 agent rollouts to a new dataset hdf5. + + python run_trained_agent.py --agent /path/to/model.pth \ + --n_rollouts 50 --horizon 400 --seed 0 \ + --dataset_path /path/to/output.hdf5 --dataset_obs + + # Write the 50 agent rollouts to a new dataset hdf5, but exclude the dataset observations + # since they might be high-dimensional (they can be extracted again using the + # dataset_states_to_obs.py script). + + python run_trained_agent.py --agent /path/to/model.pth \ + --n_rollouts 50 --horizon 400 --seed 0 \ + --dataset_path /path/to/output.hdf5 +""" +import argparse +import json +import h5py +import imageio +import numpy as np +from copy import deepcopy + +import torch + +import robomimic +import robomimic.utils.file_utils as FileUtils +import robomimic.utils.torch_utils as TorchUtils +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.obs_utils as ObsUtils +from robomimic.envs.env_base import EnvBase +from robomimic.envs.wrappers import EnvWrapper +from robomimic.algo import RolloutPolicy + + +def rollout(policy, env, horizon, render=False, video_writer=None, video_skip=5, return_obs=False, camera_names=None): + """ + Helper function to carry out rollouts. Supports on-screen rendering, off-screen rendering to a video, + and returns the rollout trajectory. + + Args: + policy (instance of RolloutPolicy): policy loaded from a checkpoint + env (instance of EnvBase): env loaded from a checkpoint or demonstration metadata + horizon (int): maximum horizon for the rollout + render (bool): whether to render rollout on-screen + video_writer (imageio writer): if provided, use to write rollout to video + video_skip (int): how often to write video frames + return_obs (bool): if True, return possibly high-dimensional observations along the trajectoryu. + They are excluded by default because the low-dimensional simulation states should be a minimal + representation of the environment. + camera_names (list): determines which camera(s) are used for rendering. Pass more than + one to output a video with multiple camera views concatenated horizontally. + + Returns: + stats (dict): some statistics for the rollout - such as return, horizon, and task success + traj (dict): dictionary that corresponds to the rollout trajectory + """ + assert isinstance(env, EnvBase) or isinstance(env, EnvWrapper) + assert isinstance(policy, RolloutPolicy) + assert not (render and (video_writer is not None)) + + policy.start_episode() + obs = env.reset() + state_dict = env.get_state() + + # hack that is necessary for robosuite tasks for deterministic action playback + obs = env.reset_to(state_dict) + + results = {} + video_count = 0 # video frame counter + total_reward = 0. + traj = dict(actions=[], rewards=[], dones=[], states=[], initial_state_dict=state_dict) + if return_obs: + # store observations too + traj.update(dict(obs=[], next_obs=[])) + try: + for step_i in range(horizon): + + # get action from policy + act = policy(ob=obs) + + # play action + next_obs, r, done, _ = env.step(act) + + # compute reward + total_reward += r + success = env.is_success()["task"] + + # visualization + if render: + env.render(mode="human", camera_name=camera_names[0]) + if video_writer is not None: + if video_count % video_skip == 0: + video_img = [] + for cam_name in camera_names: + video_img.append(env.render(mode="rgb_array", height=512, width=512, camera_name=cam_name)) + video_img = np.concatenate(video_img, axis=1) # concatenate horizontally + video_writer.append_data(video_img) + video_count += 1 + + # collect transition + traj["actions"].append(act) + traj["rewards"].append(r) + traj["dones"].append(done) + traj["states"].append(state_dict["states"]) + if return_obs: + # Note: We need to "unprocess" the observations to prepare to write them to dataset. + # This includes operations like channel swapping and float to uint8 conversion + # for saving disk space. + traj["obs"].append(ObsUtils.unprocess_obs_dict(obs)) + traj["next_obs"].append(ObsUtils.unprocess_obs_dict(next_obs)) + + # break if done or if success + if done or success: + break + + # update for next iter + obs = deepcopy(next_obs) + state_dict = env.get_state() + + except env.rollout_exceptions as e: + print("WARNING: got rollout exception {}".format(e)) + + stats = dict(Return=total_reward, Horizon=(step_i + 1), Success_Rate=float(success)) + + if return_obs: + # convert list of dict to dict of list for obs dictionaries (for convenient writes to hdf5 dataset) + traj["obs"] = TensorUtils.list_of_flat_dict_to_dict_of_list(traj["obs"]) + traj["next_obs"] = TensorUtils.list_of_flat_dict_to_dict_of_list(traj["next_obs"]) + + # list to numpy array + for k in traj: + if k == "initial_state_dict": + continue + if isinstance(traj[k], dict): + for kp in traj[k]: + traj[k][kp] = np.array(traj[k][kp]) + else: + traj[k] = np.array(traj[k]) + + return stats, traj + + +def run_trained_agent(args): + # some arg checking + write_video = (args.video_path is not None) + assert not (args.render and write_video) # either on-screen or video but not both + if args.render: + # on-screen rendering can only support one camera + assert len(args.camera_names) == 1 + + # relative path to agent + ckpt_path = args.agent + + # device + device = TorchUtils.get_torch_device(try_to_use_cuda=True) + + # restore policy + policy, ckpt_dict = FileUtils.policy_from_checkpoint(ckpt_path=ckpt_path, device=device, verbose=True) + + # read rollout settings + rollout_num_episodes = args.n_rollouts + rollout_horizon = args.horizon + if rollout_horizon is None: + # read horizon from config + config, _ = FileUtils.config_from_checkpoint(ckpt_dict=ckpt_dict) + rollout_horizon = config.experiment.rollout.horizon + + # create environment from saved checkpoint + env, _ = FileUtils.env_from_checkpoint( + ckpt_dict=ckpt_dict, + env_name=args.env, + render=args.render, + render_offscreen=(args.video_path is not None), + verbose=True, + ) + + # maybe set seed + if args.seed is not None: + np.random.seed(args.seed) + torch.manual_seed(args.seed) + + # maybe create video writer + video_writer = None + if write_video: + video_writer = imageio.get_writer(args.video_path, fps=20) + + # maybe open hdf5 to write rollouts + write_dataset = (args.dataset_path is not None) + if write_dataset: + data_writer = h5py.File(args.dataset_path, "w") + data_grp = data_writer.create_group("data") + total_samples = 0 + + rollout_stats = [] + for i in range(rollout_num_episodes): + stats, traj = rollout( + policy=policy, + env=env, + horizon=rollout_horizon, + render=args.render, + video_writer=video_writer, + video_skip=args.video_skip, + return_obs=(write_dataset and args.dataset_obs), + camera_names=args.camera_names, + ) + rollout_stats.append(stats) + + if write_dataset: + # store transitions + ep_data_grp = data_grp.create_group("demo_{}".format(i)) + ep_data_grp.create_dataset("actions", data=np.array(traj["actions"])) + ep_data_grp.create_dataset("states", data=np.array(traj["states"])) + ep_data_grp.create_dataset("rewards", data=np.array(traj["rewards"])) + ep_data_grp.create_dataset("dones", data=np.array(traj["dones"])) + if args.dataset_obs: + for k in traj["obs"]: + ep_data_grp.create_dataset("obs/{}".format(k), data=np.array(traj["obs"][k])) + ep_data_grp.create_dataset("next_obs/{}".format(k), data=np.array(traj["next_obs"][k])) + + # episode metadata + if "model" in traj["initial_state_dict"]: + ep_data_grp.attrs["model_file"] = traj["initial_state_dict"]["model"] # model xml for this episode + ep_data_grp.attrs["num_samples"] = traj["actions"].shape[0] # number of transitions in this episode + total_samples += traj["actions"].shape[0] + + rollout_stats = TensorUtils.list_of_flat_dict_to_dict_of_list(rollout_stats) + avg_rollout_stats = { k : np.mean(rollout_stats[k]) for k in rollout_stats } + avg_rollout_stats["Num_Success"] = np.sum(rollout_stats["Success_Rate"]) + print("Average Rollout Stats") + print(json.dumps(avg_rollout_stats, indent=4)) + + if write_video: + video_writer.close() + + if write_dataset: + # global metadata + data_grp.attrs["total"] = total_samples + data_grp.attrs["env_args"] = json.dumps(env.serialize(), indent=4) # environment info + data_writer.close() + print("Wrote dataset trajectories to {}".format(args.dataset_path)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + # Path to trained model + parser.add_argument( + "--agent", + type=str, + required=True, + help="path to saved checkpoint pth file", + ) + + # number of rollouts + parser.add_argument( + "--n_rollouts", + type=int, + default=27, + help="number of rollouts", + ) + + # maximum horizon of rollout, to override the one stored in the model checkpoint + parser.add_argument( + "--horizon", + type=int, + default=None, + help="(optional) override maximum horizon of rollout from the one in the checkpoint", + ) + + # Env Name (to override the one stored in model checkpoint) + parser.add_argument( + "--env", + type=str, + default=None, + help="(optional) override name of env from the one in the checkpoint, and use\ + it for rollouts", + ) + + # Whether to render rollouts to screen + parser.add_argument( + "--render", + action='store_true', + help="on-screen rendering", + ) + + # Dump a video of the rollouts to the specified path + parser.add_argument( + "--video_path", + type=str, + default=None, + help="(optional) render rollouts to this video file path", + ) + + # How often to write video frames during the rollout + parser.add_argument( + "--video_skip", + type=int, + default=5, + help="render frames to video every n steps", + ) + + # camera names to render + parser.add_argument( + "--camera_names", + type=str, + nargs='+', + default=["agentview"], + help="(optional) camera name(s) to use for rendering on-screen or to video", + ) + + # If provided, an hdf5 file will be written with the rollout data + parser.add_argument( + "--dataset_path", + type=str, + default=None, + help="(optional) if provided, an hdf5 file will be written at this path with the rollout data", + ) + + # If True and @dataset_path is supplied, will write possibly high-dimensional observations to dataset. + parser.add_argument( + "--dataset_obs", + action='store_true', + help="include possibly high-dimensional observations in output dataset hdf5 file (by default,\ + observations are excluded and only simulator states are saved)", + ) + + # for seeding before starting rollouts + parser.add_argument( + "--seed", + type=int, + default=None, + help="(optional) set seed for rollouts", + ) + + args = parser.parse_args() + run_trained_agent(args) + diff --git a/aloha-devel/robomimic/utils/__pycache__/action_utils.cpython-38.pyc b/aloha-devel/robomimic/utils/__pycache__/action_utils.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67628702d39a558d2d2c5f236f66b9eca11a4fff Binary files /dev/null and b/aloha-devel/robomimic/utils/__pycache__/action_utils.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/utils/__pycache__/log_utils.cpython-38.pyc b/aloha-devel/robomimic/utils/__pycache__/log_utils.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..710d3aaf293c3512c472a38cc04f9a3ccdd849fd Binary files /dev/null and b/aloha-devel/robomimic/utils/__pycache__/log_utils.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/utils/__pycache__/loss_utils.cpython-38.pyc b/aloha-devel/robomimic/utils/__pycache__/loss_utils.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..491c6959955b4638ee7a50af321e929d0d0fdc5a Binary files /dev/null and b/aloha-devel/robomimic/utils/__pycache__/loss_utils.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/utils/__pycache__/obs_utils.cpython-38.pyc b/aloha-devel/robomimic/utils/__pycache__/obs_utils.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..29a9f6369ded7409cb2181af04adc947ea5d2e15 Binary files /dev/null and b/aloha-devel/robomimic/utils/__pycache__/obs_utils.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/utils/__pycache__/python_utils.cpython-38.pyc b/aloha-devel/robomimic/utils/__pycache__/python_utils.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50d40e9465b93c76e5b2dbb3a24c93319149f7fe Binary files /dev/null and b/aloha-devel/robomimic/utils/__pycache__/python_utils.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/utils/__pycache__/torch_utils.cpython-38.pyc b/aloha-devel/robomimic/utils/__pycache__/torch_utils.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44da6d4f6a80df0af46ec4f5b35153689947b30d Binary files /dev/null and b/aloha-devel/robomimic/utils/__pycache__/torch_utils.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/utils/__pycache__/vis_utils.cpython-38.pyc b/aloha-devel/robomimic/utils/__pycache__/vis_utils.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f14c039fc38bbf4c8cf561bffd9d1b0f9390334a Binary files /dev/null and b/aloha-devel/robomimic/utils/__pycache__/vis_utils.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/utils/env_utils.py b/aloha-devel/robomimic/utils/env_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3cbb46ad0eeb55e11778dd18bd35dcf91d171ca0 --- /dev/null +++ b/aloha-devel/robomimic/utils/env_utils.py @@ -0,0 +1,297 @@ +""" +This file contains several utility functions for working with environment +wrappers provided by the repository, and with environment metadata saved +in dataset files. +""" +from copy import deepcopy +import robomimic.envs.env_base as EB +from robomimic.utils.log_utils import log_warning + + +def get_env_class(env_meta=None, env_type=None, env=None): + """ + Return env class from either env_meta, env_type, or env. + Note the use of lazy imports - this ensures that modules are only + imported when the corresponding env type is requested. This can + be useful in practice. For example, a training run that only + requires access to gym environments should not need to import + robosuite. + + Args: + env_meta (dict): environment metadata, which should be loaded from demonstration + hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see + @FileUtils.env_from_checkpoint). Contains 3 keys: + + :`'env_name'`: name of environment + :`'type'`: type of environment, should be a value in EB.EnvType + :`'env_kwargs'`: dictionary of keyword arguments to pass to environment constructor + + env_type (int): the type of environment, which determines the env class that will + be instantiated. Should be a value in EB.EnvType. + + env (instance of EB.EnvBase): environment instance + """ + env_type = get_env_type(env_meta=env_meta, env_type=env_type, env=env) + if env_type == EB.EnvType.ROBOSUITE_TYPE: + from robomimic.envs.env_robosuite import EnvRobosuite + return EnvRobosuite + elif env_type == EB.EnvType.GYM_TYPE: + from robomimic.envs.env_gym import EnvGym + return EnvGym + elif env_type == EB.EnvType.IG_MOMART_TYPE: + from robomimic.envs.env_ig_momart import EnvGibsonMOMART + return EnvGibsonMOMART + raise Exception("code should never reach this point") + + +def get_env_type(env_meta=None, env_type=None, env=None): + """ + Helper function to get env_type from a variety of inputs. + + Args: + env_meta (dict): environment metadata, which should be loaded from demonstration + hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see + @FileUtils.env_from_checkpoint). Contains 3 keys: + + :`'env_name'`: name of environment + :`'type'`: type of environment, should be a value in EB.EnvType + :`'env_kwargs'`: dictionary of keyword arguments to pass to environment constructor + + env_type (int): the type of environment, which determines the env class that will + be instantiated. Should be a value in EB.EnvType. + + env (instance of EB.EnvBase): environment instance + """ + checks = [(env_meta is not None), (env_type is not None), (env is not None)] + assert sum(checks) == 1, "should provide only one of env_meta, env_type, env" + if env_meta is not None: + env_type = env_meta["type"] + elif env is not None: + env_type = env.type + return env_type + + +def check_env_type(type_to_check, env_meta=None, env_type=None, env=None): + """ + Checks whether the passed env_meta, env_type, or env is of type @type_to_check. + Type corresponds to EB.EnvType. + + Args: + type_to_check (int): type to check equality against + + env_meta (dict): environment metadata, which should be loaded from demonstration + hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see + @FileUtils.env_from_checkpoint). Contains 3 keys: + + :`'env_name'`: name of environment + :`'type'`: type of environment, should be a value in EB.EnvType + :`'env_kwargs'`: dictionary of keyword arguments to pass to environment constructor + + env_type (int): the type of environment, which determines the env class that will + be instantiated. Should be a value in EB.EnvType. + + env (instance of EB.EnvBase): environment instance + """ + env_type = get_env_type(env_meta=env_meta, env_type=env_type, env=env) + return (env_type == type_to_check) + + +def check_env_version(env, env_meta): + """ + Checks whether the passed env and env_meta dictionary having matching environment versions. + Logs warning if cannot find version or versions do not match. + + Args: + env (instance of EB.EnvBase): environment instance + + env_meta (dict): environment metadata, which should be loaded from demonstration + hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see + @FileUtils.env_from_checkpoint). Contains following key: + + :`'env_version'`: environment version, type str + """ + env_system_version = env.version + env_meta_version = env_meta.get("env_version", None) + + if env_meta_version is None: + log_warning( + "No environment version found in dataset!"\ + "\nCannot verify if dataset and installed environment versions match"\ + ) + elif env_system_version != env_meta_version: + log_warning( + "Dataset and installed environment version mismatch!"\ + "\nDataset environment version: {meta}"\ + "\nInstalled environment version: {sys}".format( + sys=env_system_version, + meta=env_meta_version, + ) + ) + + +def is_robosuite_env(env_meta=None, env_type=None, env=None): + """ + Determines whether the environment is a robosuite environment. Accepts + either env_meta, env_type, or env. + """ + return check_env_type(type_to_check=EB.EnvType.ROBOSUITE_TYPE, env_meta=env_meta, env_type=env_type, env=env) + + +def create_env( + env_type, + env_name, + render=False, + render_offscreen=False, + use_image_obs=False, + lang=None, + **kwargs, +): + """ + Create environment. + + Args: + env_type (int): the type of environment, which determines the env class that will + be instantiated. Should be a value in EB.EnvType. + + env_name (str): name of environment + + render (bool): if True, environment supports on-screen rendering + + render_offscreen (bool): if True, environment supports off-screen rendering. This + is forced to be True if @use_image_obs is True. + + use_image_obs (bool): if True, environment is expected to render rgb image observations + on every env.step call. Set this to False for efficiency reasons, if image + observations are not required. + + lang: TODO documentation + """ + + # note: pass @postprocess_visual_obs True, to make sure images are processed for network inputs + env_class = get_env_class(env_type=env_type) + env = env_class( + env_name=env_name, + render=render, + render_offscreen=render_offscreen, + use_image_obs=use_image_obs, + postprocess_visual_obs=True, + lang=lang, + **kwargs, + ) + print("Created environment with name {}".format(env_name)) + print("Action size is {}".format(env.action_dimension)) + return env + + +def create_env_from_metadata( + env_meta, + env_name=None, + render=False, + render_offscreen=False, + use_image_obs=False, +): + """ + Create environment. + + Args: + env_meta (dict): environment metadata, which should be loaded from demonstration + hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see + @FileUtils.env_from_checkpoint). Contains 3 keys: + + :`'env_name'`: name of environment + :`'type'`: type of environment, should be a value in EB.EnvType + :`'env_kwargs'`: dictionary of keyword arguments to pass to environment constructor + + env_name (str): name of environment. Only needs to be provided if making a different + environment from the one in @env_meta. + + render (bool): if True, environment supports on-screen rendering + + render_offscreen (bool): if True, environment supports off-screen rendering. This + is forced to be True if @use_image_obs is True. + + use_image_obs (bool): if True, environment is expected to render rgb image observations + on every env.step call. Set this to False for efficiency reasons, if image + observations are not required. + """ + if env_name is None: + env_name = env_meta["env_name"] + env_type = get_env_type(env_meta=env_meta) + env_kwargs = env_meta["env_kwargs"] + env_kwargs["env_name"] = env_name + lang = env_meta.get("lang", None) + + env = create_env( + env_type=env_type, + render=render, + render_offscreen=render_offscreen, + use_image_obs=use_image_obs, + lang=lang, + **env_kwargs, + ) + check_env_version(env, env_meta) + return env + + +def create_env_for_data_processing( + env_meta, + camera_names, + camera_height, + camera_width, + reward_shaping, +): + """ + Creates environment for processing dataset observations and rewards. + + Args: + env_meta (dict): environment metadata, which should be loaded from demonstration + hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see + @FileUtils.env_from_checkpoint). Contains 3 keys: + + :`'env_name'`: name of environment + :`'type'`: type of environment, should be a value in EB.EnvType + :`'env_kwargs'`: dictionary of keyword arguments to pass to environment constructor + + camera_names (list of st): list of camera names that correspond to image observations + + camera_height (int): camera height for all cameras + + camera_width (int): camera width for all cameras + + reward_shaping (bool): if True, use shaped environment rewards, else use sparse task completion rewards + """ + env_name = env_meta["env_name"] + env_type = get_env_type(env_meta=env_meta) + env_kwargs = env_meta["env_kwargs"] + env_class = get_env_class(env_type=env_type) + + # remove possibly redundant values in kwargs + env_kwargs = deepcopy(env_kwargs) + env_kwargs.pop("env_name", None) + env_kwargs.pop("camera_names", None) + env_kwargs.pop("camera_height", None) + env_kwargs.pop("camera_width", None) + env_kwargs.pop("reward_shaping", None) + + env = env_class.create_for_data_processing( + env_name=env_name, + camera_names=camera_names, + camera_height=camera_height, + camera_width=camera_width, + reward_shaping=reward_shaping, + **env_kwargs, + ) + check_env_version(env, env_meta) + return env + + +def wrap_env_from_config(env, config): + """ + Wraps environment using the provided Config object to determine which wrappers + to use (if any). + """ + if config.train.frame_stack > 1: + from robomimic.envs.wrappers import FrameStackWrapper + env = FrameStackWrapper(env, num_frames=config.train.frame_stack) + + return env diff --git a/aloha-devel/robomimic/utils/file_utils.py b/aloha-devel/robomimic/utils/file_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d140b2a054012cb4124e7fdcde3121d916f70a48 --- /dev/null +++ b/aloha-devel/robomimic/utils/file_utils.py @@ -0,0 +1,573 @@ +""" +A collection of utility functions for working with files, such as reading metadata from +demonstration datasets, loading model checkpoints, or downloading dataset files. +""" +import os +import h5py +import json +import time +import urllib.request +import numpy as np +from collections import OrderedDict +from tqdm import tqdm + +import torch + +import robomimic.utils.obs_utils as ObsUtils +import robomimic.utils.env_utils as EnvUtils +import robomimic.utils.torch_utils as TorchUtils +from robomimic.config import config_factory +from robomimic.algo import algo_factory +from robomimic.algo import RolloutPolicy + + +def create_hdf5_filter_key(hdf5_path, demo_keys, key_name): + """ + Creates a new hdf5 filter key in hdf5 file @hdf5_path with + name @key_name that corresponds to the demonstrations + @demo_keys. Filter keys are generally useful to create + named subsets of the demonstrations in an hdf5, making it + easy to train, test, or report statistics on a subset of + the trajectories in a file. + + Returns the list of episode lengths that correspond to the filtering. + + Args: + hdf5_path (str): path to hdf5 file + demo_keys ([str]): list of demonstration keys which should + correspond to this filter key. For example, ["demo_0", + "demo_1"]. + key_name (str): name of filter key to create + + Returns: + ep_lengths ([int]): list of episode lengths that corresponds to + each demonstration in the new filter key + """ + f = h5py.File(hdf5_path, "a") + demos = sorted(list(f["data"].keys())) + + # collect episode lengths for the keys of interest + ep_lengths = [] + for ep in demos: + ep_data_grp = f["data/{}".format(ep)] + if ep in demo_keys: + ep_lengths.append(ep_data_grp.attrs["num_samples"]) + + # store list of filtered keys under mask group + k = "mask/{}".format(key_name) + if k in f: + del f[k] + f[k] = np.array(demo_keys, dtype='S') + + f.close() + return ep_lengths + + +def get_demos_for_filter_key(hdf5_path, filter_key): + """ + Gets demo keys that correspond to a particular filter key. + + Args: + hdf5_path (str): path to hdf5 file + filter_key (str): name of filter key + + Returns: + demo_keys ([str]): list of demonstration keys that + correspond to this filter key. For example, ["demo_0", + "demo_1"]. + """ + f = h5py.File(hdf5_path, "r") + demo_keys = [elem.decode("utf-8") for elem in np.array(f["mask/{}".format(filter_key)][:])] + f.close() + return demo_keys + + +def get_env_metadata_from_dataset(dataset_path, ds_format="robomimic"): + """ + Retrieves env metadata from dataset. + + Args: + dataset_path (str): path to dataset + + Returns: + env_meta (dict): environment metadata. Contains 3 keys: + + :`'env_name'`: name of environment + :`'type'`: type of environment, should be a value in EB.EnvType + :`'env_kwargs'`: dictionary of keyword arguments to pass to environment constructor + """ + dataset_path = os.path.expanduser(dataset_path) + f = h5py.File(dataset_path, "r") + if ds_format == "robomimic": + env_meta = json.loads(f["data"].attrs["env_args"]) + elif ds_format == "r2d2": + env_meta = dict(f.attrs) + else: + raise ValueError + f.close() + return env_meta + + +def get_shape_metadata_from_dataset(dataset_path, action_keys, all_obs_keys=None, ds_format="robomimic", verbose=False): + """ + Retrieves shape metadata from dataset. + + Args: + dataset_path (str): path to dataset + action_keys (list): list of all action key strings + all_obs_keys (list): list of all modalities used by the model. If not provided, all modalities + present in the file are used. + verbose (bool): if True, include print statements + + Returns: + shape_meta (dict): shape metadata. Contains the following keys: + + :`'ac_dim'`: action space dimension + :`'all_shapes'`: dictionary that maps observation key string to shape + :`'all_obs_keys'`: list of all observation modalities used + :`'use_images'`: bool, whether or not image modalities are present + """ + + shape_meta = {} + + # read demo file for some metadata + dataset_path = os.path.expanduser(dataset_path) + f = h5py.File(dataset_path, "r") + + if ds_format == "robomimic": + demo_id = list(f["data"].keys())[0] + demo = f["data/{}".format(demo_id)] + + for key in action_keys: + assert len(demo[key].shape) == 2 # shape should be (B, D) + action_dim = sum([demo[key].shape[1] for key in action_keys]) + shape_meta["ac_dim"] = action_dim + + # observation dimensions + all_shapes = OrderedDict() + + if all_obs_keys is None: + # use all modalities present in the file + all_obs_keys = [k for k in demo["obs"]] + + for k in sorted(all_obs_keys): + initial_shape = demo["obs/{}".format(k)].shape[1:] + if verbose: + print("obs key {} with shape {}".format(k, initial_shape)) + # Store processed shape for each obs key + all_shapes[k] = ObsUtils.get_processed_shape( + obs_modality=ObsUtils.OBS_KEYS_TO_MODALITIES[k], + input_shape=initial_shape, + ) + elif ds_format == "r2d2": + for key in action_keys: + assert len(f[key].shape) == 2 # shape should be (B, D) + action_dim = sum([f[key].shape[1] for key in action_keys]) + shape_meta["ac_dim"] = action_dim + + # observation dimensions + all_shapes = OrderedDict() + + # hack all relevant obs shapes for now + for k in [ + "robot_state/cartesian_position", + "robot_state/gripper_position", + "robot_state/joint_positions", + "camera/image/hand_camera_left_image", + "camera/image/hand_camera_right_image", + "camera/image/varied_camera_1_left_image", + "camera/image/varied_camera_1_right_image", + "camera/image/varied_camera_2_left_image", + "camera/image/varied_camera_2_right_image", + "camera/extrinsics/hand_camera_left", + # "camera/extrinsics/hand_camera_left_gripper_offset", + "camera/extrinsics/hand_camera_right", + # "camera/extrinsics/hand_camera_right_gripper_offset", + "camera/extrinsics/varied_camera_1_left", + "camera/extrinsics/varied_camera_1_right", + "camera/extrinsics/varied_camera_2_left", + "camera/extrinsics/varied_camera_2_right", + ]: + initial_shape = f["observation/{}".format(k)].shape[1:] + if len(initial_shape) == 0: + initial_shape = (1,) + + all_shapes[k] = ObsUtils.get_processed_shape( + obs_modality=ObsUtils.OBS_KEYS_TO_MODALITIES[k], + input_shape=initial_shape, + ) + else: + raise ValueError + + f.close() + + shape_meta['all_shapes'] = all_shapes + shape_meta['all_obs_keys'] = all_obs_keys + shape_meta['use_images'] = ObsUtils.has_modality("rgb", all_obs_keys) + + return shape_meta + + +def load_dict_from_checkpoint(ckpt_path): + """ + Load checkpoint dictionary from a checkpoint file. + + Args: + ckpt_path (str): Path to checkpoint file. + + Returns: + ckpt_dict (dict): Loaded checkpoint dictionary. + """ + ckpt_path = os.path.expanduser(ckpt_path) + if not torch.cuda.is_available(): + ckpt_dict = torch.load(ckpt_path, map_location=lambda storage, loc: storage) + else: + ckpt_dict = torch.load(ckpt_path) + return ckpt_dict + + +def maybe_dict_from_checkpoint(ckpt_path=None, ckpt_dict=None): + """ + Utility function for the common use case where either an ckpt path + or a ckpt_dict is provided. This is a no-op if ckpt_dict is not + None, otherwise it loads the model dict from the ckpt path. + + Args: + ckpt_path (str): Path to checkpoint file. Only needed if not providing @ckpt_dict. + + ckpt_dict(dict): Loaded model checkpoint dictionary. Only needed if not providing @ckpt_path. + + Returns: + ckpt_dict (dict): Loaded checkpoint dictionary. + """ + assert (ckpt_path is not None) or (ckpt_dict is not None) + if ckpt_dict is None: + ckpt_dict = load_dict_from_checkpoint(ckpt_path) + return ckpt_dict + + +def algo_name_from_checkpoint(ckpt_path=None, ckpt_dict=None): + """ + Return algorithm name that was used to train a checkpoint or + loaded model dictionary. + + Args: + ckpt_path (str): Path to checkpoint file. Only needed if not providing @ckpt_dict. + + ckpt_dict(dict): Loaded model checkpoint dictionary. Only needed if not providing @ckpt_path. + + Returns: + algo_name (str): algorithm name + + ckpt_dict (dict): loaded checkpoint dictionary (convenient to avoid + re-loading checkpoint from disk multiple times) + """ + ckpt_dict = maybe_dict_from_checkpoint(ckpt_path=ckpt_path, ckpt_dict=ckpt_dict) + algo_name = ckpt_dict["algo_name"] + return algo_name, ckpt_dict + + +def update_config(cfg): + """ + Updates the config for backwards-compatibility if it uses outdated configurations. + + See https://github.com/ARISE-Initiative/robomimic/releases/tag/v0.2.0 for more info. + + Args: + cfg (dict): Raw dictionary of config values + """ + # Check if image modality is defined -- this means we're using an outdated config + # Note: There may be a nested hierarchy, so we possibly check all the nested obs cfgs which can include + # e.g. a planner and actor for HBC + + def find_obs_dicts_recursively(dic): + dics = [] + if "modalities" in dic: + dics.append(dic) + else: + for child_dic in dic.values(): + dics += find_obs_dicts_recursively(child_dic) + return dics + + obs_cfgs = find_obs_dicts_recursively(cfg["observation"]) + for obs_cfg in obs_cfgs: + modalities = obs_cfg["modalities"] + + found_img = False + for modality_group in ("obs", "subgoal", "goal"): + if modality_group in modalities: + img_modality = modalities[modality_group].pop("image", None) + if img_modality is not None: + found_img = True + modalities[modality_group]["rgb"] = img_modality + + if found_img: + # Also need to map encoder kwargs correctly + old_encoder_cfg = obs_cfg.pop("encoder") + + # Create new encoder entry for RGB + rgb_encoder_cfg = { + "core_class": "VisualCore", + "core_kwargs": { + "backbone_kwargs": dict(), + "pool_kwargs": dict(), + }, + "obs_randomizer_class": None, + "obs_randomizer_kwargs": dict(), + } + + if "visual_feature_dimension" in old_encoder_cfg: + rgb_encoder_cfg["core_kwargs"]["feature_dimension"] = old_encoder_cfg["visual_feature_dimension"] + + if "visual_core" in old_encoder_cfg: + rgb_encoder_cfg["core_kwargs"]["backbone_class"] = old_encoder_cfg["visual_core"] + + for kwarg in ("pretrained", "input_coord_conv"): + if "visual_core_kwargs" in old_encoder_cfg and kwarg in old_encoder_cfg["visual_core_kwargs"]: + rgb_encoder_cfg["core_kwargs"]["backbone_kwargs"][kwarg] = old_encoder_cfg["visual_core_kwargs"][kwarg] + + # Optionally add pooling info too + if old_encoder_cfg.get("use_spatial_softmax", True): + rgb_encoder_cfg["core_kwargs"]["pool_class"] = "SpatialSoftmax" + + for kwarg in ("num_kp", "learnable_temperature", "temperature", "noise_std"): + if "spatial_softmax_kwargs" in old_encoder_cfg and kwarg in old_encoder_cfg["spatial_softmax_kwargs"]: + rgb_encoder_cfg["core_kwargs"]["pool_kwargs"][kwarg] = old_encoder_cfg["spatial_softmax_kwargs"][kwarg] + + # Update obs randomizer as well + for kwarg in ("obs_randomizer_class", "obs_randomizer_kwargs"): + if kwarg in old_encoder_cfg: + rgb_encoder_cfg[kwarg] = old_encoder_cfg[kwarg] + + # Store rgb config + obs_cfg["encoder"] = {"rgb": rgb_encoder_cfg} + + # Also add defaults for low dim + obs_cfg["encoder"]["low_dim"] = { + "core_class": None, + "core_kwargs": { + "backbone_kwargs": dict(), + "pool_kwargs": dict(), + }, + "obs_randomizer_class": None, + "obs_randomizer_kwargs": dict(), + } + + +def config_from_checkpoint(algo_name=None, ckpt_path=None, ckpt_dict=None, verbose=False): + """ + Helper function to restore config from a checkpoint file or loaded model dictionary. + + Args: + algo_name (str): Algorithm name. + + ckpt_path (str): Path to checkpoint file. Only needed if not providing @ckpt_dict. + + ckpt_dict(dict): Loaded model checkpoint dictionary. Only needed if not providing @ckpt_path. + + verbose (bool): if True, include print statements + + Returns: + config (dict): Raw loaded configuration, without properties replaced. + + ckpt_dict (dict): loaded checkpoint dictionary (convenient to avoid + re-loading checkpoint from disk multiple times) + """ + ckpt_dict = maybe_dict_from_checkpoint(ckpt_path=ckpt_path, ckpt_dict=ckpt_dict) + if algo_name is None: + algo_name, _ = algo_name_from_checkpoint(ckpt_dict=ckpt_dict) + + # restore config from loaded model dictionary + config_dict = json.loads(ckpt_dict['config']) + update_config(cfg=config_dict) + + if verbose: + print("============= Loaded Config =============") + print(json.dumps(config_dict, indent=4)) + + config = config_factory(algo_name, dic=config_dict) + + # lock config to prevent further modifications and ensure missing keys raise errors + config.lock() + + return config, ckpt_dict + + +def policy_from_checkpoint(device=None, ckpt_path=None, ckpt_dict=None, verbose=False): + """ + This function restores a trained policy from a checkpoint file or + loaded model dictionary. + + Args: + device (torch.device): if provided, put model on this device + + ckpt_path (str): Path to checkpoint file. Only needed if not providing @ckpt_dict. + + ckpt_dict(dict): Loaded model checkpoint dictionary. Only needed if not providing @ckpt_path. + + verbose (bool): if True, include print statements + + Returns: + model (RolloutPolicy): instance of Algo that has the saved weights from + the checkpoint file, and also acts as a policy that can easily + interact with an environment in a training loop + + ckpt_dict (dict): loaded checkpoint dictionary (convenient to avoid + re-loading checkpoint from disk multiple times) + """ + ckpt_dict = maybe_dict_from_checkpoint(ckpt_path=ckpt_path, ckpt_dict=ckpt_dict) + + # algo name and config from model dict + algo_name, _ = algo_name_from_checkpoint(ckpt_dict=ckpt_dict) + config, _ = config_from_checkpoint(algo_name=algo_name, ckpt_dict=ckpt_dict, verbose=verbose) + + # read config to set up metadata for observation modalities (e.g. detecting rgb observations) + ObsUtils.initialize_obs_utils_with_config(config) + + # shape meta from model dict to get info needed to create model + shape_meta = ckpt_dict["shape_metadata"] + + # maybe restore observation normalization stats + obs_normalization_stats = ckpt_dict.get("obs_normalization_stats", None) + if obs_normalization_stats is not None: + assert config.train.hdf5_normalize_obs + for m in obs_normalization_stats: + for k in obs_normalization_stats[m]: + obs_normalization_stats[m][k] = np.array(obs_normalization_stats[m][k]) + + # maybe restore action normalization stats + action_normalization_stats = ckpt_dict.get("action_normalization_stats", None) + if action_normalization_stats is not None: + for m in action_normalization_stats: + for k in action_normalization_stats[m]: + action_normalization_stats[m][k] = np.array(action_normalization_stats[m][k]) + + if device is None: + # get torch device + device = TorchUtils.get_torch_device(try_to_use_cuda=config.train.cuda) + + # create model and load weights + model = algo_factory( + algo_name, + config, + obs_key_shapes=shape_meta["all_shapes"], + ac_dim=shape_meta["ac_dim"], + device=device, + ) + model.deserialize(ckpt_dict["model"]) + model.set_eval() + model = RolloutPolicy( + model, + obs_normalization_stats=obs_normalization_stats, + action_normalization_stats=action_normalization_stats + ) + if verbose: + print("============= Loaded Policy =============") + print(model) + return model, ckpt_dict + + +def env_from_checkpoint(ckpt_path=None, ckpt_dict=None, env_name=None, render=False, render_offscreen=False, verbose=False): + """ + Creates an environment using the metadata saved in a checkpoint. + + Args: + ckpt_path (str): Path to checkpoint file. Only needed if not providing @ckpt_dict. + + ckpt_dict(dict): Loaded model checkpoint dictionary. Only needed if not providing @ckpt_path. + + env_name (str): if provided, override environment name saved in checkpoint + + render (bool): if True, environment supports on-screen rendering + + render_offscreen (bool): if True, environment supports off-screen rendering. This + is forced to be True if saved model uses image observations. + + Returns: + env (EnvBase instance): environment created using checkpoint + + ckpt_dict (dict): loaded checkpoint dictionary (convenient to avoid + re-loading checkpoint from disk multiple times) + """ + ckpt_dict = maybe_dict_from_checkpoint(ckpt_path=ckpt_path, ckpt_dict=ckpt_dict) + + # metadata from model dict to get info needed to create environment + env_meta = ckpt_dict["env_metadata"] + shape_meta = ckpt_dict["shape_metadata"] + + # create env from saved metadata + env = EnvUtils.create_env_from_metadata( + env_meta=env_meta, + render=render, + render_offscreen=render_offscreen, + use_image_obs=shape_meta["use_images"], + ) + config, _ = config_from_checkpoint(algo_name=ckpt_dict["algo_name"], ckpt_dict=ckpt_dict, verbose=False) + env = EnvUtils.wrap_env_from_config(env, config=config) # apply environment warpper, if applicable + if verbose: + print("============= Loaded Environment =============") + print(env) + return env, ckpt_dict + + +class DownloadProgressBar(tqdm): + def update_to(self, b=1, bsize=1, tsize=None): + if tsize is not None: + self.total = tsize + self.update(b * bsize - self.n) + + +def url_is_alive(url): + """ + Checks that a given URL is reachable. + From https://gist.github.com/dehowell/884204. + + Args: + url (str): url string + + Returns: + is_alive (bool): True if url is reachable, False otherwise + """ + request = urllib.request.Request(url) + request.get_method = lambda: 'HEAD' + + try: + urllib.request.urlopen(request) + return True + except urllib.request.HTTPError: + return False + + +def download_url(url, download_dir, check_overwrite=True): + """ + First checks that @url is reachable, then downloads the file + at that url into the directory specified by @download_dir. + Prints a progress bar during the download using tqdm. + + Modified from https://github.com/tqdm/tqdm#hooks-and-callbacks, and + https://stackoverflow.com/a/53877507. + + Args: + url (str): url string + download_dir (str): path to directory where file should be downloaded + check_overwrite (bool): if True, will sanity check the download fpath to make sure a file of that name + doesn't already exist there + """ + + # check if url is reachable. We need the sleep to make sure server doesn't reject subsequent requests + assert url_is_alive(url), "@download_url got unreachable url: {}".format(url) + time.sleep(0.5) + + # infer filename from url link + fname = url.split("/")[-1] + file_to_write = os.path.join(download_dir, fname) + + # If we're checking overwrite and the path already exists, + # we ask the user to verify that they want to overwrite the file + if check_overwrite and os.path.exists(file_to_write): + user_response = input(f"Warning: file {file_to_write} already exists. Overwrite? y/n\n") + assert user_response.lower() in {"yes", "y"}, f"Did not receive confirmation. Aborting download." + + with DownloadProgressBar(unit='B', unit_scale=True, + miniters=1, desc=fname) as t: + urllib.request.urlretrieve(url, filename=file_to_write, reporthook=t.update_to) diff --git a/aloha-devel/robomimic/utils/log_utils.py b/aloha-devel/robomimic/utils/log_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9edd6e5ce3bc0688d1386bf0305b065ab405e959 --- /dev/null +++ b/aloha-devel/robomimic/utils/log_utils.py @@ -0,0 +1,230 @@ +""" +This file contains utility classes and functions for logging to stdout, stderr, +and to tensorboard. +""" +import os +import sys +import numpy as np +from datetime import datetime +from contextlib import contextmanager +import textwrap +import time +from tqdm import tqdm +from termcolor import colored + +import robomimic + +# global list of warning messages can be populated with @log_warning and flushed with @flush_warnings +WARNINGS_BUFFER = [] + + +class PrintLogger(object): + """ + This class redirects print statements to both console and a file. + """ + def __init__(self, log_file): + self.terminal = sys.stdout + print('STDOUT will be forked to %s' % log_file) + self.log_file = open(log_file, "a") + + def write(self, message): + self.terminal.write(message) + self.log_file.write(message) + self.log_file.flush() + + def flush(self): + # this flush method is needed for python 3 compatibility. + # this handles the flush command by doing nothing. + # you might want to specify some extra behavior here. + pass + + +class DataLogger(object): + """ + Logging class to log metrics to tensorboard and/or retrieve running statistics about logged data. + """ + def __init__(self, log_dir, config, log_tb=True, log_wandb=False): + """ + Args: + log_dir (str): base path to store logs + log_tb (bool): whether to use tensorboard logging + """ + self._tb_logger = None + self._wandb_logger = None + self._data = dict() # store all the scalar data logged so far + + if log_tb: + from tensorboardX import SummaryWriter + self._tb_logger = SummaryWriter(os.path.join(log_dir, 'tb')) + + if log_wandb: + import wandb + import robomimic.macros as Macros + + # set up wandb api key if specified in macros + if Macros.WANDB_API_KEY is not None: + os.environ["WANDB_API_KEY"] = Macros.WANDB_API_KEY + + assert Macros.WANDB_ENTITY is not None, "WANDB_ENTITY macro is set to None." \ + "\nSet this macro in {base_path}/macros_private.py" \ + "\nIf this file does not exist, first run python {base_path}/scripts/setup_macros.py".format(base_path=robomimic.__path__[0]) + + # attempt to set up wandb 10 times. If unsuccessful after these trials, don't use wandb + num_attempts = 10 + for attempt in range(num_attempts): + try: + # set up wandb + self._wandb_logger = wandb + + self._wandb_logger.init( + entity=Macros.WANDB_ENTITY, + project=config.experiment.logging.wandb_proj_name, + name=config.experiment.name, + dir=log_dir, + mode=("offline" if attempt == num_attempts - 1 else "online"), + ) + + # set up info for identifying experiment + wandb_config = {k: v for (k, v) in config.meta.items() if k not in ["hp_keys", "hp_values"]} + for (k, v) in zip(config.meta["hp_keys"], config.meta["hp_values"]): + wandb_config[k] = v + if "algo" not in wandb_config: + wandb_config["algo"] = config.algo_name + self._wandb_logger.config.update(wandb_config) + + break + except Exception as e: + log_warning("wandb initialization error (attempt #{}): {}".format(attempt + 1, e)) + self._wandb_logger = None + time.sleep(30) + + def record(self, k, v, epoch, data_type='scalar', log_stats=False): + """ + Record data with logger. + Args: + k (str): key string + v (float or image): value to store + epoch: current epoch number + data_type (str): the type of data. either 'scalar' or 'image' + log_stats (bool): whether to store the mean/max/min/std for all data logged so far with key k + """ + + assert data_type in ['scalar', 'image'] + + if data_type == 'scalar': + # maybe update internal cache if logging stats for this key + if log_stats or k in self._data: # any key that we're logging or previously logged + if k not in self._data: + self._data[k] = [] + self._data[k].append(v) + + # maybe log to tensorboard + if self._tb_logger is not None: + if data_type == 'scalar': + self._tb_logger.add_scalar(k, v, epoch) + if log_stats: + stats = self.get_stats(k) + for (stat_k, stat_v) in stats.items(): + stat_k_name = '{}-{}'.format(k, stat_k) + self._tb_logger.add_scalar(stat_k_name, stat_v, epoch) + elif data_type == 'image': + if len(v.shape) == 3: + v = v[None, ...] + self._tb_logger.add_images(k, img_tensor=v, global_step=epoch, dataformats="NHWC") + + if self._wandb_logger is not None: + try: + if data_type == 'scalar': + self._wandb_logger.log({k: v}, step=epoch) + if log_stats: + stats = self.get_stats(k) + for (stat_k, stat_v) in stats.items(): + self._wandb_logger.log({stat_k: stat_v}, step=epoch) + elif data_type == 'image': + import wandb + self._wandb_logger.log({k: wandb.Image(v)}, step=epoch) + except Exception as e: + log_warning("wandb logging: {}".format(e)) + + def get_stats(self, k): + """ + Computes running statistics for a particular key. + Args: + k (str): key string + Returns: + stats (dict): dictionary of statistics + """ + stats = dict() + stats['mean'] = np.mean(self._data[k]) + stats['std'] = np.std(self._data[k]) + stats['min'] = np.min(self._data[k]) + stats['max'] = np.max(self._data[k]) + return stats + + def close(self): + """ + Run before terminating to make sure all logs are flushed + """ + if self._tb_logger is not None: + self._tb_logger.close() + + if self._wandb_logger is not None: + self._wandb_logger.finish() + + +class custom_tqdm(tqdm): + """ + Small extension to tqdm to make a few changes from default behavior. + By default tqdm writes to stderr. Instead, we change it to write + to stdout. + """ + def __init__(self, *args, **kwargs): + assert "file" not in kwargs + super(custom_tqdm, self).__init__(*args, file=sys.stdout, **kwargs) + + +@contextmanager +def silence_stdout(): + """ + This contextmanager will redirect stdout so that nothing is printed + to the terminal. Taken from the link below: + + https://stackoverflow.com/questions/6735917/redirecting-stdout-to-nothing-in-python + """ + old_target = sys.stdout + try: + with open(os.devnull, "w") as new_target: + sys.stdout = new_target + yield new_target + finally: + sys.stdout = old_target + + +def log_warning(message, color="yellow", print_now=True): + """ + This function logs a warning message by recording it in a global warning buffer. + The global registry will be maintained until @flush_warnings is called, at + which point the warnings will get printed to the terminal. + + Args: + message (str): warning message to display + color (str): color of message - defaults to "yellow" + print_now (bool): if True (default), will print to terminal immediately, in + addition to adding it to the global warning buffer + """ + global WARNINGS_BUFFER + buffer_message = colored("ROBOMIMIC WARNING(\n{}\n)".format(textwrap.indent(message, " ")), color) + WARNINGS_BUFFER.append(buffer_message) + if print_now: + print(buffer_message) + + +def flush_warnings(): + """ + This function flushes all warnings from the global warning buffer to the terminal and + clears the global registry. + """ + global WARNINGS_BUFFER + for msg in WARNINGS_BUFFER: + print(msg) + WARNINGS_BUFFER = [] diff --git a/aloha-devel/robomimic/utils/script_utils.py b/aloha-devel/robomimic/utils/script_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e42eaf627715b7d8ce30aaf43dfbfd992678beab --- /dev/null +++ b/aloha-devel/robomimic/utils/script_utils.py @@ -0,0 +1,15 @@ +""" +Collection of miscellaneous utility tools +""" + +def deep_update(d, u): + """ + Copied from https://stackoverflow.com/a/3233356 + """ + import collections + for k, v in u.items(): + if isinstance(v, collections.abc.Mapping): + d[k] = deep_update(d.get(k, {}), v) + else: + d[k] = v + return d \ No newline at end of file