机器人模仿学习之动作分块ACT算法的代码剖析、部署训练

前言

本文最早是属于《斯坦福Mobile ALOHA背后的关键技术：动作分块ACT算法的原理解析》的第二、第三部分，涉及到动作分块ACT的代码剖析与部署训练

但因为想把ACT的代码逐行剖析的更细致些，加之为避免上一篇文章太过于长，故把动作分块ACT的代码剖析与部署实践这块独立出来成本文

第一部分 动作分块算法ACT的代码剖析

关于ACT的代码，我们可以重点研究下这个仓库：GitHub - tonyzhaozh/act，我司同事杜老师也于24年1.10日跑通了这份代码(如何跑通的教程见下文第二部分)

• imitate_episodes.py，训练和评估 ACT
• policy.py，An adaptor for ACT policy
• detr，ACT 的模型定义 修改自 DETR
• sim_env.py，具有 joint space control的 Mujoco + DM_Control 环境
• ee_sim_env.py，具有EE space control的 Mujoco + DM_Control 环境
• scripted_policy.py，模拟环境的脚本化策略
• constants.py，跨文件共享的常量
• utils.py，数据加载和辅助函数等实用程序
• visualize_episodes.py，保存 .hdf5 数据集中的视频

1.1 ACT的训练与评估imitate_episodes.py

1.1.1 主程序

1. 从命令行参数中获取模型训练和评估的相关配置

   def main(args):
       set_seed(1)  # 设置随机种子以保证结果可重现
       # 解析命令行参数
       is_eval = args["eval"]  # 是否为评估模式的布尔标志
       ckpt_dir = args["ckpt_dir"]  # 保存/加载检查点的目录
       policy_class = args["policy_class"]  # 使用的策略类
       onscreen_render = args["onscreen_render"]  # 是否进行屏幕渲染的标志
       task_name = args["task_name"]  # 任务名称
       batch_size_train = args["batch_size"]  # 训练批大小
       batch_size_val = args["batch_size"]  # 验证批大小
       num_epochs = args["num_epochs"]  # 训练的总周期数
       use_waypoint = args["use_waypoint"]  # 是否使用航点
       constant_waypoint = args["constant_waypoint"]  # 持续航点的设置
   
       # 根据是否使用航点打印相应信息
       if use_waypoint:
           print("Using waypoint")  # 使用航点
       if constant_waypoint is not None:
           print(f"Constant waypoint: {constant_waypoint}")  # 持续航点

2. 根据任务名称和配置获取任务参数，例如数据集目录、任务类型等

       # 获取任务参数
       is_sim = True  # 硬编码为True以避免从aloha中查找常量
       # 如果是模拟任务，从constants导入SIM_TASK_CONFIGS
       if is_sim:
           from constants import SIM_TASK_CONFIGS
           task_config = SIM_TASK_CONFIGS[task_name]
       else:
           from aloha_scripts.constants import TASK_CONFIGS
           task_config = TASK_CONFIGS[task_name]
   
       # 从任务配置中获取相关参数
       dataset_dir = task_config["dataset_dir"]
       num_episodes = task_config["num_episodes"]
       episode_len = task_config["episode_len"]
       camera_names = task_config["camera_names"]

3. 定义模型的架构和超参数，包括学习率、网络结构、层数等

      # 固定参数
       state_dim = 14  # 状态维度
       lr_backbone = 1e-5  # 主干网络的学习率
       backbone = "resnet18"  # 使用的主干网络类型

4. 根据策略类别设置策略配置

       # 根据策略类别设置策略配置
       if policy_class == "ACT":
           # ACT策略的特定参数
           enc_layers = 4
           dec_layers = 7
           nheads = 8
           policy_config = {
               "lr": args["lr"],
               "num_queries": args["chunk_size"],
               "kl_weight": args["kl_weight"],
               "hidden_dim": args["hidden_dim"],
               "dim_feedforward": args["dim_feedforward"],
               "lr_backbone": lr_backbone,
               "backbone": backbone,
               "enc_layers": enc_layers,
               "dec_layers": dec_layers,
               "nheads": nheads,
               "camera_names": camera_names,
           }
       elif policy_class == "CNNMLP":
           # CNNMLP策略的特定参数
           policy_config = {
               "lr": args["lr"],
               "lr_backbone": lr_backbone,
               "backbone": backbone,
               "num_queries": 1,
               "camera_names": camera_names,
           }
       else:
           raise NotImplementedError

5. 配置训练参数

       # 配置训练参数
       config = {
           "num_epochs": num_epochs,
           "ckpt_dir": ckpt_dir,
           "episode_len": episode_len,
           "state_dim": state_dim,
           "lr": args["lr"],
           "policy_class": policy_class,
           "onscreen_render": onscreen_render,
           "policy_config": policy_config,
           "task_name": task_name,
           "seed": args["seed"],
           "temporal_agg": args["temporal_agg"],
           "camera_names": camera_names,
           "real_robot": not is_sim,
       }

6. 如果设置为评估模式，加载保存的模型权重并在验证集上评估模型性能，计算成功率和平均回报

       # 如果为评估模式，执行评估流程
       if is_eval:
           ckpt_names = [f"policy_best.ckpt"]
           results = []
           for ckpt_name in ckpt_names:
               success_rate, avg_return = eval_bc(config, ckpt_name, save_episode=True)
               results.append([ckpt_name, success_rate, avg_return])
   
           for ckpt_name, success_rate, avg_return in results:
               print(f"{ckpt_name}: {success_rate=} {avg_return=}")
           print()
           exit()
   
       # 加载数据
       train_dataloader, val_dataloader, stats, _ = load_data(
           dataset_dir,
           num_episodes,
           camera_names,
           batch_size_train,
           batch_size_val,
           use_waypoint,
           constant_waypoint,
       )
   
       # 保存数据集统计信息
       if not os.path.isdir(ckpt_dir):
           os.makedirs(ckpt_dir)
       stats_path = os.path.join(ckpt_dir, f"dataset_stats.pkl")
       with open(stats_path, "wb") as f:
           pickle.dump(stats, f)
   

7. 最后，将结果打印出来

       # 训练并获取最佳检查点信息
       best_ckpt_info = train_bc(train_dataloader, val_dataloader, config)
       best_epoch, min_val_loss, best_state_dict = best_ckpt_info
   
       # 保存最佳检查点
       ckpt_path = os.path.join(ckpt_dir, f"policy_best.ckpt")
       torch.save(best_state_dict, ckpt_path)
       print(f"Best ckpt, val loss {min_val_loss:.6f} @ epoch{best_epoch}")

1.1.2 make_policy、make_optimizer、get_image

根据指定的policy_class(策略类别，目前支持两种类型："ACT"和"CNNMLP")，和policy_config(策略配置)创建一个策略模型对象

def make_policy(policy_class, policy_config):
    if policy_class == 'ACT':
        policy = ACTPolicy(policy_config)  # 如果策略类是 ACT，创建 ACTPolicy
    elif policy_class == 'CNNMLP':
        policy = CNNMLPPolicy(policy_config)  # 如果策略类是 CNNMLP，创建 CNNMLPPolicy
    else:
        raise NotImplementedError  # 如果不是以上两种类型，则抛出未实现错误
    return policy  # 返回创建的策略对象

make_optimizer用于创建策略模型的优化器(optimizer)，并返回创建的优化器对象。优化器的作用是根据策略模型的损失函数来更新模型的参数，以使损失函数尽量减小

def make_optimizer(policy_class, policy):
    if policy_class == 'ACT':
        optimizer = policy.configure_optimizers()  # 如果策略类是 ACT，配置优化器
    elif policy_class == 'CNNMLP':
        optimizer = policy.configure_optimizers()  # 如果策略类是 CNNMLP，配置优化器
    else:
        raise NotImplementedError  # 如果不是以上两种类型，则抛出未实现错误
    return optimizer  # 返回配置的优化器

get_image的作用是获取一个时间步(ts)的图像数据。函数接受两个参数：ts和camera_names

def get_image(ts, camera_names):
    curr_images = []
    for cam_name in camera_names:
        curr_image = rearrange(ts.observation['images'][cam_name], 'h w c -> c h w')  # 重排图像数组
        curr_images.append(curr_image)  # 将处理后的图像添加到列表中
    curr_image = np.stack(curr_images, axis=0)  # 将图像列表堆叠成数组
    curr_image = torch.from_numpy(curr_image / 255.0).float().cuda().unsqueeze(0)  # 将数组转换为 PyTorch 张量
    return curr_image  # 返回处理后的图像张量

1. ts是一个时间步的数据，包含了多个相机（摄像头）拍摄的图像
ts.observation["images"]包含了各个相机拍摄的图像数据，而camera_names是一个列表，包含了要获取的相机的名称
2. 函数通过循环遍历camera_names中的相机名称，从ts.observation["images"]中获取对应相机的图像数据
这些图像数据首先通过rearrange函数重新排列维度，将"height-width-channels"的顺序变为"channels-height-width"，以适应PyTorch的数据格式
3. 获取的图像数据被放入curr_images列表中
4. 接下来，函数将curr_images列表中的所有图像数据堆叠成一个张量（tensor），np.stack(curr_images, axis=0)这一行代码实现了这个操作
5. 接着，图像数据被归一化到[0, 1]的范围，然后转换为PyTorch的float类型，并移到GPU上（如果可用）。最后，图像数据被增加了一个额外的维度（unsqueeze(0)），以适应模型的输入要求

最终，函数返回包含时间步图像数据的PyTorch张量。这个图像数据可以被用于输入到神经网络模型中进行处理

1.1.3 eval_bc：评估一个行为克隆(behavior cloning)模型

1. 的

   def eval_bc(config, ckpt_name, save_episode=True):
       set_seed(1000)  # 设置随机种子为 1000
       # 从配置中获取参数
       ckpt_dir = config['ckpt_dir']
       state_dim = config['state_dim']
       real_robot = config['real_robot']
       policy_class = config['policy_class']
       onscreen_render = config['onscreen_render']
       policy_config = config['policy_config']
       camera_names = config['camera_names']
       max_timesteps = config['episode_len']
       task_name = config['task_name']
       temporal_agg = config['temporal_agg']
       onscreen_cam = 'angle'
   
       # 加载策略和统计信息
       ckpt_path = os.path.join(ckpt_dir, ckpt_name)
       policy = make_policy(policy_class, policy_config)
       loading_status = policy.load_state_dict(torch.load(ckpt_path))
       print(loading_status)
       policy.cuda()
       policy.eval()
       print(f'Loaded: {ckpt_path}')
       stats_path = os.path.join(ckpt_dir, f'dataset_stats.pkl')
       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['action_std'] + stats['action_mean']

2. 的

       # 加载环境
       if real_robot:
           from aloha_scripts.robot_utils import move_grippers  # 从 aloha_scripts.robot_utils 导入 move_grippers
           from aloha_scripts.real_env import make_real_env  # 从 aloha_scripts.real_env 导入 make_real_env
           env = make_real_env(init_node=True)  # 创建真实机器人环境
           env_max_reward = 0
       else:
           from sim_env import make_sim_env  # 从 sim_env 导入 make_sim_env
           env = make_sim_env(task_name)  # 创建模拟环境
           env_max_reward = env.task.max_reward
   
       # 设置查询频率和时间聚合参数
       query_frequency = policy_config['num_queries']
       if temporal_agg:
           query_frequency = 1
           num_queries = policy_config['num_queries']
   
       # 设置最大时间步数
       max_timesteps = int(max_timesteps * 1)  # 可以根据实际任务调整最大时间步数

3. 设置评估的循环次数(num_rollouts)，每次循环都会进行一次评估
   在每次循环中，初始化环境，执行模型生成的动作并观测环境的响应
   将每个时间步的观测数据(包括图像、关节位置等)存储在相应的列表中

       # 设置回放次数和初始化结果列表
       num_rollouts = 50
       episode_returns = []
       highest_rewards = []
   
       # 回放循环
       for rollout_id in range(num_rollouts):
           rollout_id += 0
           # 设置任务
           if 'sim_transfer_cube' in task_name:
               BOX_POSE[0] = sample_box_pose()  # 在模拟重置中使用的 BOX_POSE
           elif 'sim_insertion' in task_name:
               BOX_POSE[0] = np.concatenate(sample_insertion_pose())  # 在模拟重置中使用的 BOX_POSE
   
           ts = env.reset()  # 重置环境
   
           # 处理屏幕渲染
           if onscreen_render:
               ax = plt.subplot()
               plt_img = ax.imshow(env._physics.render(height=480, width=640, camera_id=onscreen_cam))
               plt.ion()
   
           # 评估循环
           if temporal_agg:
               all_time_actions = torch.zeros([max_timesteps, max_timesteps+num_queries, state_dim]).cuda()
   
           qpos_history = torch.zeros((1, max_timesteps, state_dim)).cuda()
           image_list = []  # 用于可视化的图像列表
           qpos_list = []
           target_qpos_list = []
           rewards = []
   
           # 在不计算梯度的模式下执行
           with torch.inference_mode():
               for t in range(max_timesteps):
                   # 更新屏幕渲染和等待时间
                   if onscreen_render:
                       image = env._physics.render(height=480, width=640, camera_id=onscreen_cam)
                       plt_img.set_data(image)
                       plt.pause(DT)
   
                   # 处理上一时间步的观测值以获取 qpos 和图像列表
                   obs = ts.observation
                   if 'images' in obs:
                       image_list.append(obs['images'])
                   else:
                       image_list.append({'main': obs['image']})
                   qpos_numpy = np.array(obs['qpos'])
                   qpos = pre_process(qpos_numpy)
                   qpos = torch.from_numpy(qpos).float().cuda().unsqueeze(0)
                   qpos_history[:, t] = qpos
                   curr_image = get_image(ts, camera_names)
   
                   # 查询策略
                   if config['policy_class'] == "ACT":
                       if t % query_frequency == 0:
                           all_actions = policy(qpos, curr_image)
                       if temporal_agg:
                           all_time_actions[[t], t:t+num_queries] = all_actions
                           actions_for_curr_step = all_time_actions[:, t]
                           actions_populated = torch.all(actions_for_curr_step != 0, axis=1)
                           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 = torch.from_numpy(exp_weights).cuda().unsqueeze(dim=1)
                           raw_action = (actions_for_curr_step * exp_weights).sum(dim=0, keepdim=True)
                       else:
                           raw_action = all_actions[:, t % query_frequency]
                   elif config['policy_class'] == "CNNMLP":
                       raw_action = policy(qpos, curr_image)
                   else:
                       raise NotImplementedError
   
                   # 后处理动作
                   raw_action = raw_action.squeeze(0).cpu().numpy()
                   action = post_process(raw_action)
                   target_qpos = action
   
                   # 步进环境
                   ts = env.step(target_qpos)
   
                   # 用于可视化的列表
                   qpos_list.append(qpos_numpy)
                   target_qpos_list.append(target_qpos)
                   rewards.append(ts.reward)
   
               plt.close()  # 关闭绘图窗口
           if real_robot:
               move_grippers([env.puppet_bot_left, env.puppet_bot_right], [PUPPET_GRIPPER_JOINT_OPEN] * 2, move_time=0.5)  # 打开夹持器
               pass

   计算每次评估的总回报，以及每次评估的最高回报，并记录成功率

           # 计算回报和奖励
           rewards = np.array(rewards)
           episode_return = np.sum(rewards[rewards != None])
           episode_returns.append(episode_return)
           episode_highest_reward = np.max(rewards)
           highest_rewards.append(episode_highest_reward)
           print(f'Rollout {rollout_id}\n{episode_return=}, {episode_highest_reward=}, {env_max_reward=}, Success: {episode_highest_reward == env_max_reward}')

   如果指定了保存评估过程中的图像数据，将每次评估的图像数据保存为视频

           # 保存视频
           if save_episode:
               save_videos(image_list, DT, video_path=os.path.join(ckpt_dir, f'video{rollout_id}.mp4'))

4. 输出评估结果，包括成功率、平均回报以及回报分布
   将评估结果保存到文本文件中

   
       # 计算成功率和平均回报
       # 计算成功率，即最高奖励的次数与环境最大奖励相等的比率
       success_rate = np.mean(np.array(highest_rewards) == env_max_reward)
   
       # 计算平均回报
       avg_return = np.mean(episode_returns)
   
       # 创建一个包含成功率和平均回报的摘要字符串
       summary_str = f'\n成功率: {success_rate}\n平均回报: {avg_return}\n\n'
   
       # 遍历奖励范围，计算每个奖励范围内的成功率
       for r in range(env_max_reward + 1):
           # 统计最高奖励大于等于 r 的次数
           more_or_equal_r = (np.array(highest_rewards) >= r).sum()
       
           # 计算成功率
           more_or_equal_r_rate = more_or_equal_r / num_rollouts
       
           # 将结果添加到摘要字符串中
           summary_str += f'奖励 >= {r}: {more_or_equal_r}/{num_rollouts} = {more_or_equal_r_rate*100}%\n'
   
       # 打印摘要字符串
       print(summary_str)
   
       # 将成功率保存到文本文件
       result_file_name = 'result_' + ckpt_name.split('.')[0] + '.txt'
   with open(os.path.join(ckpt_dir, result_file_name), 'w') as f:
           f.write(summary_str)  # 写入摘要字符串
           f.write(repr(episode_returns))  # 写入回报数据
           f.write('\n\n')
           f.write(repr(highest_rewards))  # 写入最高奖励数据
   
       # 返回成功率和平均回报
       return success_rate, avg_return

// 待更

第二部分 Mobile Aloha或Aloha软件层面代码的跑通与部署

// 待更

参考文献与推荐阅读

1. Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware
2. Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware(阅读笔记)
3. Aloha 机械臂的学习记录2——AWE：AWE + ACT
4. ..