Add files using upload-large-folder tool

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}
+    }

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(),
+)

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/

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

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() 

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 

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

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
+

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

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

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

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)

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

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)

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__(), '  ')

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)

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)

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)

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

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
+         ])))

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

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

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
+

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
+        ])))

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
+        ])))

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"]

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": []
+    }
+}

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": {}
+            }
+        }
+    }
+}

aloha-devel/robomimic/models/__init__.py

@@ -0,0 +1 @@
+from .obs_core import EncoderCore, Randomizer

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)

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

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"]

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

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

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)