机器人抓取与操作学习(六):在机器人控制任务场景下训练和评估的深度学习模型
·
下面的程序整体是一个用于训练和评估深度学习模型的框架,主要用于机器人控制任务。
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 einops import rearrange
from constants import DT
from constants import PUPPET_GRIPPER_JOINT_OPEN
from utils import load_data # 数据加载函数
from utils import sample_box_pose, sample_insertion_pose # 机器人姿态采样函数
from utils import compute_dict_mean, set_seed, detach_dict # 辅助函数
from policy import ACTPolicy, CNNMLPPolicy # 策略模型
from visualize_episodes import save_videos # 视频保存函数
from sim_env import BOX_POSE # 模拟环境中的盒子位置
import IPython
e = IPython.embed # 用于在特定位置调试代码
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']
# 获取任务参数
is_sim = task_name[:4] == 'sim_' # 判断任务是否为模拟任务
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']
# 固定参数
state_dim = 7 # 状态维度
lr_backbone = 1e-5 # backbone学习率
backbone = 'resnet18' # 使用的网络架构
if policy_class == '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':
policy_config = {
'lr': args['lr'],
'lr_backbone': lr_backbone,
'backbone': backbone,
'num_queries': 1,
'camera_names': camera_names,
}
else:
raise NotImplementedError # 如果策略类未实现,则抛出错误
# 配置字典
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 # 如果是模拟任务,则设置为False
}
# 如果是评估模式
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)
# 保存数据集统计信息
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) # 保存统计信息
# 训练并返回最佳检查点信息
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}')
# 创建策略
def make_policy(policy_class, policy_config):
if policy_class == 'ACT':
policy = ACTPolicy(policy_config) # 创建ACT策略
elif policy_class == 'CNNMLP':
policy = CNNMLPPolicy(policy_config) # 创建CNNMLP策略
else:
raise NotImplementedError # 抛出错误
return policy
# 创建优化器
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
# 获取图像数据
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) # 转换为张量并归一化
return curr_image
# 评估行为克隆
def eval_bc(config, ckpt_name, save_episode=True):
set_seed(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'] # 后处理函数
# 加载环境
if real_robot:
from aloha_scripts.robot_utils import move_grippers # 引入机器人控制函数
from aloha_scripts.real_env import make_real_env # 引入真实环境创建函数
env = make_real_env(init_node=True) # 创建真实环境
env_max_reward = 0
else:
from sim_env import 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) # 最大时间步长
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() # 用于模拟重置
elif 'sim_pick_n_place_cube' in task_name:
BOX_POSE[0] = sample_box_pose()
elif 'sim_insertion' in task_name:
BOX_POSE[0] = np.concatenate(sample_insertion_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) # 等待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
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() # 转换为numpy数组
action = post_process(raw_action) # 后处理动作
target_qpos = action # 目标qpos
### 步进环境
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')) # 保存视频
success_rate = np.mean(np.array(highest_rewards) == env_max_reward) # 计算成功率
avg_return = np.mean(episode_returns) # 计算平均返回值
summary_str = f'\nSuccess rate: {success_rate}\nAverage return: {avg_return}\n\n'
for r in range(env_max_reward + 1):
more_or_equal_r = (np.array(highest_rewards) >= r).sum() # 统计达到奖励r的回合数
more_or_equal_r_rate = more_or_equal_r / num_rollouts # 计算达到奖励r的比例
summary_str += f'Reward >= {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 # 返回成功率和平均返回值
# 前向传播
def forward_pass(data, policy):
image_data, qpos_data, action_data, is_pad = data
image_data, qpos_data, action_data, is_pad = image_data.cuda(), qpos_data.cuda(), action_data.cuda(), is_pad.cuda() # 将数据转移到GPU
return policy(qpos_data, image_data, action_data, is_pad) # 查询策略输出
# 训练行为克隆
def train_bc(train_dataloader, val_dataloader, config):
num_epochs = config['num_epochs']
ckpt_dir = config['ckpt_dir']
seed = config['seed']
policy_class = config['policy_class']
policy_config = config['policy_config']
set_seed(seed) # 设置随机种子
policy = make_policy(policy_class, policy_config) # 创建策略
policy.cuda() # 转移到GPU
optimizer = make_optimizer(policy_class, policy) # 创建优化器
train_history = [] # 存储训练历史
validation_history = [] # 存储验证历史
min_val_loss = np.inf # 初始化最小验证损失
best_ckpt_info = None # 初始化最佳检查点信息
for epoch in tqdm(range(num_epochs)): # 迭代每个周期
print(f'\nEpoch {epoch}')
# 验证
with torch.inference_mode(): # 关闭梯度计算
policy.eval() # 设置为评估模式
epoch_dicts = []
for batch_idx, data in enumerate(val_dataloader):
forward_dict = forward_pass(data, policy) # 前向传播
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.state_dict()))
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)
# 训练
policy.train() # 设置为训练模式
optimizer.zero_grad() # 清空梯度
for batch_idx, data in enumerate(train_dataloader):
forward_dict = forward_pass(data, policy) # 前向传播
# 反向传播
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 % 100 == 0: # 每100个周期保存一次检查点
ckpt_path = os.path.join(ckpt_dir, f'policy_epoch_{epoch}_seed_{seed}.ckpt')
torch.save(policy.state_dict(), ckpt_path) # 保存模型参数
plot_history(train_history, validation_history, epoch, ckpt_dir, seed) # 画出训练曲线
ckpt_path = os.path.join(ckpt_dir, f'policy_last.ckpt')
torch.save(policy.state_dict(), 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}')
# 保存训练曲线
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):
# 保存训练曲线
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]) # 可选:设置y轴范围
plt.tight_layout() # 自动调整布局
plt.legend()
plt.title(key) # 设置标题
plt.savefig(plot_path) # 保存图像
print(f'Saved plots to {ckpt_dir}') # 打印保存路径
if __name__ == '__main__':
# 命令行参数解析
parser = argparse.ArgumentParser()
parser.add_argument('--eval', action='store_true') # 是否评估模式
parser.add_argument('--onscreen_render', action='store_true') # 是否在屏幕上渲染
parser.add_argument('--ckpt_dir', action='store', type=str, help='ckpt_dir', required=True) # 检查点目录
parser.add_argument('--policy_class', action='store', type=str, help='policy_class, capitalize', required=True) # 策略类
parser.add_argument('--task_name', action='store', type=str, help='task_name', required=True) # 任务名称
parser.add_argument('--batch_size', action='store', type=int, help='batch_size', required=True) # 批量大小
parser.add_argument('--seed', action='store', type=int, help='seed', required=True) # 随机种子
parser.add_argument('--num_epochs', action='store', type=int, help='num_epochs', required=True) # 训练周期数
parser.add_argument('--lr', action='store', type=float, help='lr', required=True) # 学习率
# 针对ACT的额外参数
parser.add_argument('--kl_weight', action='store', type=int, help='KL Weight', required=False) # KL权重
parser.add_argument('--chunk_size', action='store', type=int, help='chunk_size', required=False) # 块大小
parser.add_argument('--hidden_dim', action='store', type=int, help='hidden_dim', required=False) # 隐藏层维度
parser.add_argument('--dim_feedforward', action='store', type=int, help='dim_feedforward', required=False) # 前馈层维度
parser.add_argument('--temporal_agg', action='store_true') # 是否使用时间聚合
main(vars(parser.parse_args())) # 执行主函数
程序结构和功能说明
-
导入必要的库:
- 导入了深度学习、数据处理、可视化和其他辅助库。
-
主函数
main(args):- 处理命令行参数。
- 根据任务名称加载相应的配置。
- 如果是评估模式,则加载并评估模型;否则,加载数据并进行训练。
- 训练后保存最佳模型及其参数。
-
策略创建和优化器配置:
make_policy函数:根据传入的策略类型创建对应的策略对象。make_optimizer函数:根据传入的策略类型创建优化器。
-
数据获取和处理:
get_image函数:处理状态信息,提取和转换图像数据。eval_bc函数:实现行为克隆的评估,包括环境重置、策略查询和回合奖励计算。
-
训练过程:
train_bc函数:实现训练过程,包括训练和验证循环,损失计算和优化步骤。forward_pass函数:执行单次前向传播。
-
绘制和保存训练曲线:
plot_history函数:绘制训练和验证曲线,并保存图像文件。
-
命令行解析部分:
- 通过 argparse 库解析命令行参数,方便用户配置训练过程的相关参数。
运行方式
程序通过命令行输入参数来运行,用户需要指定检查点目录、策略类型、任务名称、批量大小、随机种子、训练周期数和学习率等参数。根据参数,程序将选择相应的操作(训练或评估)并执行。
训练:
python3 imitate_episodes.py \
--task_name sim_pick_n_place_cube_scripted \
--ckpt_dir ckpt_dir --policy_class ACT \
--kl_weight 10 --chunk_size 100 \
--hidden_dim 512 --batch_size 8 \
--dim_feedforward 3200 --num_epochs 2000 \
--lr 1e-5 --seed 0 --temporal_agg
评估训练结果:
python3 imitate_episodes.py \
--task_name sim_pick_n_place_cube_scripted \
--ckpt_dir ckpt_dir --policy_class ACT \
--kl_weight 10 --chunk_size 100 \
--hidden_dim 512 --batch_size 8 \
--dim_feedforward 3200 --num_epochs 2000 \
--lr 1e-5 --seed 0 --temporal_agg --eval
程序的主要功能概括
该程序主要用于训练和评估用于机器人控制任务的深度学习模型,具体功能包括:
-
配置和初始化:解析命令行参数,设置随机种子,加载任务相关的配置和数据集。
-
模型创建:根据指定的策略类型(如ACT或CNNMLP)创建对应的策略模型,并配置相应的优化器。
-
数据加载:从指定路径加载训练和验证数据集,以便用于模型训练和评估。
-
训练过程:包含训练循环和验证过程:
- 在训练过程中,模型不断更新参数以最小化损失。
- 在每个训练周期后,验证模型性能并记录损失。
-
评估功能:支持评估模式,加载最佳模型,对其进行评估,并计算成功率和平均返回值。
-
结果存储:将训练和验证过程中生成的数据(如损失、奖励、图像等)保存在指定目录中,并生成训练曲线图。
-
可视化:支持在屏幕上实时渲染环境,对于机器人控制任务进行可视化展示。
训练的输入和输出
输入
-
数据输入:
- 训练数据:包含状态(qpos)、图像(image)、动作(action)和填充信息(is_pad)的批量数据,由
train_dataloader提供。 - 验证数据:同样结构的数据,由
val_dataloader提供,用于在每个训练周期后评估模型性能。
- 训练数据:包含状态(qpos)、图像(image)、动作(action)和填充信息(is_pad)的批量数据,由
-
超参数:
- 学习率(lr)
- 批量大小(batch_size)
- 训练周期数(num_epochs)
- 策略类(policy_class)
- 任务名称(task_name)
- 随机种子(seed)
- 其他可选参数,例如KL权重、块大小、隐藏层维度等。
输出
-
训练输出:
- 训练损失:每个训练周期的损失值,表示模型在训练集上的表现。
- 验证损失:每个训练周期后计算的验证损失,用于判断模型的泛化能力。
- 最佳检查点:保存具有最低验证损失的模型参数,便于后续加载和评估。
-
评估输出(在评估模式下):
- 成功率:模型在指定回合上成功完成任务的比例。
- 平均返回值:模型在评估过程中获得的平均奖励。
- 详细的奖励分布:对于每个设定的奖励值,统计达到此奖励的回合数和比例。
-
可视化输出:
- 保存的训练和验证曲线图,提供损失随训练进展的变化情况。
-
视频记录(可选):评估过程中生成的环境渲染视频,便于可视化模型的决策过程。
通过这些输入和输出,程序能够有效地训练机器人控制模型,并评估其在真实或模拟环境中的表现。
DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐

所有评论(0)