前言
本文一开始是属于此文 《HumanPlus——斯坦福ALOHA团队开源的人形机器人：融合影子学习技术、RL、模仿学习》的第四部分，但为避免原文篇幅过长，故把该部分抽取出来独立成文
促成本文主要有以下两个因素
- 大模型时代，技术更迭速度超过以往任何时期，而个人认为机器人(具身智能)将是未来几年最大的趋势
包括我司机器人线下营曾一天连报5人(开营后，将邀请一波人加入我司机器人开发队伍)，期待与更多有缘人共同开发机器人以服务更多工厂 - 加之后来，24年Q3时，我司除了开发一些特定场景的工业协作机器人业务外，我们还在帮几家公司复现umi/dexcap
包括有公司找到我们，希望帮他们复现humanplus，那顺带把其源码分析 分享下
而过程中解读斯坦福人形机器人humanplus的代码时，还是充满乐趣的，比如又遇到了熟悉的ppo，想到
- 去年上半年啃了半年的ChatGPT原理没有白啃
- 去年下半年带队做大模型应用，直接促成我司「七月在线」从教育公司往科技公司的转型
- 今年上半年则是具身智能
也算是可谓三者合一、步步为赢了
话休絮烦，对于humanplus的整个代码框架，总计包含以下五个部分
- Humanoid Shadowing Transformer (HST)，此所谓low-level，属于机器人小脑
这个部分的代码是基于仿真的强化学习实现，使用了legged_gym和rsl_rl
需要先安装IsaacGym v4，并将isaacgym文件夹放置在HST文件夹中
提供了训练HST的示例命令，以及如何使用训练好的策略 - Humanoid Imitation Transformer (HIT)，此所谓high-level，属于机器人大脑
这部分代码用于现实世界中的模仿学习，基于ACT repo和Mobile ALOHA repo「关于什么时ACT、详见此文：ACT的原理解析：斯坦福炒虾机器人Moblie Aloha的动作分块算法ACT，至于什么moblie aloha，详见此文《以Mobile ALOHA为代表的模仿学习的爆发：从Dobb·E、Gello到斯坦福ALOHA、UMI、DexCap、伯克利FMB》的第三部分」
提供了安装指南，包括创建conda环境和安装所需的Python库 - Pose Estimation
身体姿态估计使用WHAM，手部姿态估计使用HaMeR - Hardware Codebase
硬件代码基于unitree_ros2，适用于与真实机器人的交互
提供了安装指南，包括创建conda环境、安装unitree_sdk和unitree_ros2 - Example Usages
提供了训练HST和HIT的具体命令示例。
对于硬件代码，提供了如何放置训练好的策略文件以及如何运行硬件脚本的示例
至于上面前两个部分的关系，我已经在此文《HumanPlus——斯坦福ALOHA团队开源的人形机器人：融合影子学习技术、RL、模仿学习》中写的很清晰了
为更形象的说明，再次贴下这个模型架构图
- 下图左侧是用于低级控制的仅解码器Transformer(Humanoid Shadowing Transformer)，是所谓的机器人小脑low-level
- 下图右下角是一个用于模仿学习的仅解码器Transformer(Humanoid Imitation Transformer)，是所谓的机器人大脑high-level
对于上图右下角的仅解码器的人形模仿Transformer(此HIT大脑，其Transformer层数:12、隐藏层维度:1024、注意力头数:16、学习率:1e-4、训练样本:40个演示)，其用于技能策略，具体而言
第一部分 机器人小脑之low-level控制策略Humanoid Shadowing Transformer(HST)
为了对整个humanplus的代码做更好的解读，先把整体的代码结构梳理一下(如下4张图所示，总计4个部分，前3个部分都是HST相关，第4个部分则是HIT相关)
1.1 HST/rsl_rl/rsl_rl：涉及PPO、actor_critic
在humanplus/HST/rsl_rl/rsl_rl文件夹里有以下分文件夹
1.1.1 rsl_rl/rsl_rl/algorithms/ppo.py
其中，algorithms文件夹里有两个文件：__init__.py、ppo.py
接下来，咱们重点看下HST/rsl_rl/rsl_rl/algorithms/ppo.py
首先是一些类的导入、与PPO类的定义
import torch # 导入 PyTorch 库
import torch.nn as nn # 导入 PyTorch 神经网络模块
import torch.optim as optim # 导入 PyTorch 优化模块
from rsl_rl.modules import ActorCriticTransformer # 从 rsl_rl.modules 导入 ActorCriticTransformer 类
from rsl_rl.storage import RolloutStorage # 从 rsl_rl.storage 导入 RolloutStorage 类
class PPO: # 定义 PPO 类
actor_critic: ActorCriticTransformer # 定义 actor_critic 变量类型为 ActorCriticTransformer
def __init__(self, # 定义 PPO 类的初始化函数
actor_critic, # 定义 actor_critic 参数
num_learning_epochs=1, # 定义 num_learning_epochs 参数，默认值为1
num_mini_batches=1, # 定义 num_mini_batches 参数，默认值为1
clip_param=0.2, # 定义 clip_param 参数，默认值为0.2
gamma=0.998, # 定义 gamma 参数，默认值为0.998
lam=0.95, # 定义 lam 参数，默认值为0.95
value_loss_coef=1.0, # 定义 value_loss_coef 参数，默认值为1.0
entropy_coef=0.0, # 定义 entropy_coef 参数，默认值为0.0
learning_rate=1e-3, # 定义 learning_rate 参数，默认值为1e-3
max_grad_norm=1.0, # 定义 max_grad_norm 参数，默认值为1.0
use_clipped_value_loss=True, # 定义 use_clipped_value_loss 参数，默认值为 True
schedule="fixed", # 定义 schedule 参数，默认值为 "fixed"
desired_kl=0.01, # 定义 desired_kl 参数，默认值为 0.01
device='cpu', # 定义 device 参数，默认值为 'cpu'
):
self.device = device # 初始化 self.device 为传入的 device 参数
self.desired_kl = desired_kl # 初始化 self.desired_kl 为传入的 desired_kl 参数
self.schedule = schedule # 初始化 self.schedule 为传入的 schedule 参数
self.learning_rate = learning_rate # 初始化 self.learning_rate 为传入的 learning_rate 参数
# PPO components
# 最上面PPO类里一系列参数的self初始化
# PPO parameters
# 最上面PPO类里一系列参数的self初始化//..
待补充
def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape):
# 初始化存储
self.storage = RolloutStorage(num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape, self.device)
def test_mode(self): # 定义测试模式函数
self.actor_critic.test() # 设置 actor_critic 为测试模式
def train_mode(self): # 定义训练模式函数
self.actor_critic.train() # 设置 actor_critic 为训练模式
def act(self, obs, critic_obs): # 定义 act 函数
if self.actor_critic.is_recurrent: # 如果 actor_critic 是循环的
self.transition.hidden_states = self.actor_critic.get_hidden_states() # 获取隐藏状态
# Compute the actions and values
# 计算动作和值
self.transition.actions = self.actor_critic.act(obs).detach() # 获取动作
self.transition.values = self.actor_critic.evaluate(critic_obs).detach() # 获取值
self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(self.transition.actions).detach() # 获取动作的对数概率
self.transition.action_mean = self.actor_critic.action_mean.detach() # 获取动作均值
self.transition.action_sigma = self.actor_critic.action_std.detach() # 获取动作标准差
# need to record obs and critic_obs before env.step()
# 在环境步之前需要记录 obs 和 critic_obs
self.transition.observations = obs # 记录观察
self.transition.critic_observations = critic_obs # 记录评论员观察
return self.transition.actions # 返回动作//..
待补充
def process_env_step(self, rewards, dones, infos): # 定义处理环境步函数
self.transition.rewards = rewards.clone() # 克隆奖励
self.transition.dones = dones # 记录完成状态
# Bootstrapping on time outs
# 对超时进行引导
if 'time_outs' in infos: # 如果信息中有 'time_outs'
self.transition.rewards += self.gamma * torch.squeeze(self.transition.values * infos['time_outs'].unsqueeze(1).to(self.device), 1) # 更新奖励
# Record the transition
# 记录过渡
self.storage.add_transitions(self.transition) # 添加过渡到存储中
self.transition.clear() # 清除过渡数据
self.actor_critic.reset(dones) # 重置 actor_critic
def compute_returns(self, last_critic_obs): # 定义计算回报函数
last_values = self.actor_critic.evaluate(last_critic_obs).detach() # 计算最后的值
self.storage.compute_returns(last_values, self.gamma, self.lam) # 计算回报//..
接下来重点分析下，这个update函数
- 首先，它根据模型是否是循环的来选择不同的mini-batch生成器。然后，它遍历生成器产生的每个mini-batch，对每个batch进行一系列的操作
在每个mini-batch中，首先使用[`actor_critic`]模型对观察值进行行动，并获取行动的对数概率，评估价值，以及行动的均值和标准差def update(self): # 定义更新函数 mean_value_loss = 0 # 初始化平均值损失 mean_surrogate_loss = 0 # 初始化平均代理损失 if self.actor_critic.is_recurrent: # 如果 actor_critic 是循环的 generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) # 使用循环小批量生成器 else: # 否则 generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) # 使用小批量生成器 for obs_batch, critic_obs_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, \ old_mu_batch, old_sigma_batch, hid_states_batch, masks_batch in generator: # 迭代生成器中的数据 # 获取动作概率和价值 self.actor_critic.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) # 获取动作 actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch) # 获取动作的对数概率 value_batch = self.actor_critic.evaluate(critic_obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) # 获取值 mu_batch = self.actor_critic.action_mean # 获取动作均值 sigma_batch = self.actor_critic.action_std # 获取动作标准差 entropy_batch = self.actor_critic.entropy # 获取熵 - 如果设置了期望的KL散度并且调度策略为'adaptive'，则计算当前策略和旧策略之间的KL散度，并根据KL散度的大小动态调整学习率
# KL if self.desired_kl != None and self.schedule == 'adaptive': # 如果期望的 KL 不为 None 且调度为 'adaptive' with torch.inference_mode(): # 使用推理模式 kl = torch.sum( torch.log(sigma_batch / old_sigma_batch + 1.e-5) + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) / (2.0 * torch.square(sigma_batch)) - 0.5, axis=-1) # 计算 KL 散度 kl_mean = torch.mean(kl) # 计算 KL 散度的均值 # 如果 KL 散度均值大于期望的 KL 的两倍 if kl_mean > self.desired_kl * 2.0: self.learning_rate = max(1e-5, self.learning_rate / 1.5) # 减小学习率 # 如果 KL 散度均值小于期望的 KL 的一半且大于 0.0 elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: self.learning_rate = min(1e-2, self.learning_rate * 1.5) # 增大学习率 for param_group in self.optimizer.param_groups: # 对于优化器中的每个参数组 param_group['lr'] = self.learning_rate # 更新学习率
上述代码是典型的自适应KL惩罚的过程
上述公式中的是怎么取值的呢，事实上，
- 先设一个可以接受的 KL 散度的最大值
假设优化完以后，KL 散度值太大导致
，意味着
与
差距过大(即学习率/步长过大)，也就代表后面惩罚的项
惩罚效果太弱而没有发挥作用，故增大惩罚把
- 再设一个 KL 散度的最小值
如果优化完，意味着
至于详细了解请查看本博客内此文《强化学习极简入门：通俗理解MDP、DP MC TC和Q学习、策略梯度、PPO》的4.4.1 什么是近端策略优化PPO与PPO-penalty
- 接下来，计算actor损失（surrogate loss），这是PPO算法的核心部分。它首先计算新旧策略的比率，然后计算未裁剪和裁剪的代理损失，并取两者的最大值作为最终的代理损失
# Surrogate loss # 代理损失 ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) # 计算比率 surrogate = -torch.squeeze(advantages_batch) * ratio # 计算代理损失 surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(ratio, 1.0 - self.clip_param, 1.0 + self.clip_param) # 计算剪切后的代理损失 surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() # 计算代理损失的均值
其实上面的代码就是对近端策略优化裁剪PPO-clip的直接实现，其对应的公式如下(详细了解请查看本博客内此文《强化学习极简入门：通俗理解MDP、DP MC TC和Q学习、策略梯度、PPO》的第4.4.2节PPO算法的另一个变种：近端策略优化裁剪PPO-clip)
- 然后，计算价值函数的损失。如果设置了使用裁剪价值损失，那么会计算裁剪的价值损失和未裁剪的价值损失，并取两者的最大值作为最终的价值损失。否则，直接计算未裁剪的价值损失
# Value function loss # 价值函数损失 if self.use_clipped_value_loss: value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(-self.clip_param, self.clip_param) value_losses = (value_batch - returns_batch).pow(2) value_losses_clipped = (value_clipped - returns_batch).pow(2) value_loss = torch.max(value_losses, value_losses_clipped).mean() else: value_loss = (returns_batch - value_batch).pow(2).mean() loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean()
其实就是之前这篇文章《从零实现带RLHF的类ChatGPT：逐行解析微软DeepSpeed Chat的源码》中3.6中AC架构下的PPO训练：在加了β惩罚且截断后的RM之下，通过经验数据不断迭代策略且估计value讲过的
在1个ppo_batch中，critic的损失计算公式为：
裁剪新价值估计，使其不至于太偏离采集经验时的旧价值估计，使得经验回放仍能有效：
critic将拟合回报R：
可能有同学疑问上面的代码和我说的这个公式并没有一一对齐呀，为了方便大家一目了然，我们把代码逐行再分析下
- 对于这行代码
value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(-self.clip_param,self.clip_param)
转换成公式便是
V_clipped = V_old + clip(V_new - V_old, -ε, ε)
它和我上面贴的公式表达的其实是一样的
因为我上面贴的公式要表达的是：，那该不等式两边都减去个
，不就意味着
- 而接下来这三行代码
value_losses = (value_batch - returns_batch).pow(2)
value_losses_clipped = (value_clipped - returns_batch).pow(2)
value_loss = torch.max(value_losses, value_losses_clipped).mean()
则表达的就是如下公式是不一目了然了..
- 最后，将代理损失、价值损失和熵损失结合起来，形成最终的损失函数
最后的最后，进行梯度下降，更新模型的参数，并在所有mini-batch更新完成后，计算平均的价值损失和代理损失，并清空存储器# Gradient step # 梯度步 self.optimizer.zero_grad() # 清零梯度 loss.backward() # 反向传播 nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm) # 剪切梯度 self.optimizer.step() # 更新优化器 mean_value_loss += value_loss.item() # 累加平均值损失 mean_surrogate_loss += surrogate_loss.item() # 累加平均代理损失 num_updates = self.num_learning_epochs * self.num_mini_batches # 计算更新次数 mean_value_loss /= num_updates # 计算平均值损失 mean_surrogate_loss /= num_updates # 计算平均代理损失 self.storage.clear() # 清除存储 return mean_value_loss, mean_surrogate_loss # 返回平均值损失和平均代理损失
1.1.2 rsl_rl/rsl_rl/modules/actor_critic_transformer.py
actor_critic_transformer.py这段代码定义了一个基于BERT风格的Transformer模型，用于强化学习中的Actor-Critic方法
它由三个主要的类组成：Transformer_Block, Transformer、ActorCriticTransformer
- Transformer_Block 类
Transformer_Block 类是构建Transformer模型的基础块(下图来自此文的4.1.1 ViT的架构：Embedding层 + Transformer Encoder + MLP Head)
它包含了多头注意力机制 (Muitihead Attention） 和前馈神经网络 (FeedForward Neural Network）
这个类首先通过层归一代 (Layer Normalization），然后应用多头注意力机制，再次进行层归一化，并最后通过一个前馈神经网络# a BERT-style transformer block class Transformer_Block(nn.Module): def __init__(self, latent_dim, num_head, dropout_rate) -> None: super().__init__() self.num_head = num_head self.latent_dim = latent_dim self.ln_1 = nn.LayerNorm(latent_dim) self.attn = nn.MultiheadAttention(latent_dim, num_head, dropout=dropout_rate, batch_first=True) self.ln_2 = nn.LayerNorm(latent_dim) self.mlp = nn.Sequential( nn.Linear(latent_dim, 4 * latent_dim), nn.GELU(), nn.Linear(4 * latent_dim, latent_dim), nn.Dropout(dropout_rate), ) def forward(self, x): x = self.ln_1(x) x = x + self.attn(x, x, x, need_weights=False)[0] x = self.ln_2(x) x = x + self.mlp(x) return x - Transformer 类
Transformer 类构建了一个完整的 Transformer模型，它接收输入数据并通过一系列的Transformer_Block 进行处理
首先，输入数据通过一个线性层和位置嵌入 (Position Embedding） 进行处理，以增加位置信息
然后，数据通过多个 Transformer_Block 进行处理class Transformer(nn.Module): def __init__(self, input_dim, output_dim, context_len, latent_dim=128, num_head=4, num_layer=4, dropout_rate=0.1) -> None: super().__init__() self.input_dim = input_dim self.output_dim = output_dim self.context_len = context_len self.latent_dim = latent_dim self.num_head = num_head self.num_layer = num_layer self.input_layer = nn.Sequential( nn.Linear(input_dim, latent_dim), nn.Dropout(dropout_rate), ) self.weight_pos_embed = nn.Embedding(context_len, latent_dim)
最后，模型通过另一个线性层输出最終结果self.attention_blocks = nn.Sequential( *[Transformer_Block(latent_dim, num_head, dropout_rate) for _ in range(num_layer)], )self.output_layer = nn.Sequential( nn.LayerNorm(latent_dim), nn.Linear(latent_dim, output_dim), ) def forward(self, x): x = self.input_layer(x) x = x + self.weight_pos_embed(torch.arange(x.shape[1], device=x.device)) x = self.attention_blocks(x) # take the last token x = x[:, -1, :] x = self.output_layer(x) return x - ActorCriticTransformer 类
ActorcriticTransformer 类实现了Actor-Critic方法，其中Actor和Critic都使用了上述的Transformer模型
Actor负责生成动作，而Critic负责评估当前策略的价值。这个类还包括了动作噪声的处理，以及一些用于强化学习的特定方法，如 act、evaluate 和update-distribution 等
1.1.3 rsl_rl/rsl_rl/modules/actor_critic.py
这段代码定义了一个名为 ActorCritic的类，它继承自nn.Module，是一个用于强化学习的演员-评论家模型的实现
这个模型包含两个主要部分：一个用于决策的策略网络(演员）和一个用于评估动作价值的价值网络（评论家）
- 构造函数_-init--- 接收多个参数，包括观察空问的维度（分别为演员和评论家）、动作空问的维度、隐藏层的维度、激活函数类型以及初始噪声标淮差
class ActorCritic(nn.Module): is_recurrent = False def __init__(self, num_actor_obs, num_critic_obs, num_actions, actor_hidden_dims=[256, 256, 256], critic_hidden_dims=[256, 256, 256], activation='elu', init_noise_std=1.0, **kwargs): - 构造函数首先检查是否有末预期的参数传入，并打印警告信息。然后，它调用get_activation 函数来获取指定的激活函数
if kwargs: print("ActorCritic.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs.keys()])) super(ActorCritic, self).__init__() activation = get_activation(activation) - 接下来，构造函数初始化演员和评论家网络
这两个网络都是使用多层感知机MLP构建的，其中每一层都是通过 nn.Linear 创建的，并且在每个线性层之后应用了激活函数mlp_input_dim_a = num_actor_obs mlp_input_dim_c = num_critic_obs
比如这是对policy的构建
再比如这是对value的构建# Policy actor_layers = [] actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0])) actor_layers.append(activation) for l in range(len(actor_hidden_dims)): if l == len(actor_hidden_dims) - 1: actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions)) else: actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1])) actor_layers.append(activation) self.actor = nn.Sequential(*actor_layers)
演员网络的输出维度等于动作空间的维度，而评论家网络的输出维度固定为1，表示对当前状态的价值评估# Value function critic_layers = [] critic_layers.append(nn.Linear(mlp_input_dim_c, critic_hidden_dims[0])) critic_layers.append(activation) for l in range(len(critic_hidden_dims)): if l == len(critic_hidden_dims) - 1: critic_layers.append(nn.Linear(critic_hidden_dims[l], 1)) else: critic_layers.append(nn.Linear(critic_hidden_dims[l], critic_hidden_dims[l + 1])) critic_layers.append(activation) self.critic = nn.Sequential(*critic_layers) print(f"Actor MLP: {self.actor}") print(f"Critic MLP: {self.critic}") - 此外，构造函数还初始化了一个用于动作输出的噪声参数 self.std，并设置了一个分布self.distribution，该分布稍后将用于生成带有噪声的动作
类中还定义了几个方法，包括reset(重置状态，当前为空实现）、forward（抽象方法，未实现）、update_distribution(根据当前观察更新动作分布）、act （根据当前观察采取动作）、get_actions_1og-prob（计算动作的对数概率）、 act-inference（推断模式下的动作选择，不包含噪声）和evaluate（评估给定观察的价值)
最后，get_activation 函数根据传入的激活函数名称返回对应的PyTorch激活函数对象，如果传入的名称无效，则打印错误信息井返回None
1.2 HST/legged_gym/legged_gym/envs：涉及仿真环境
1.2.1 envs/base/legged_robot_config.py
这段代码定义了一个名为 LeggedRobotCfg 的类，它继承自BaseConfig，这个类用于配置一个多足机器人的仿真环境，包括环境参数、地形设置、命令配置、初始状态、控制参数、资产信息、领域随机化、奖励设置、归一化参数、噪声设置、观察者视角和仿真参数
此外，还定义了一个名为 LeggedRobotCfgPP0 的类，专门用于配置使用PPO算法进行训练的参数
首先，LeggedRobotcfg 类中的
- env 子类定义了仿真环境的基本参数，如环境数量、观察值数量、动作数量、环境问隔、是否发送超时信息、每个剧集的长度等
class LeggedRobotCfg(BaseConfig): class env: num_envs = 4096 num_observations = 235 num_privileged_obs = None # if not None a priviledge_obs_buf will be returned by step() (critic obs for assymetric training). None is returned otherwise num_actions = 12 env_spacing = 3. # not used with heightfields/trimeshes send_timeouts = True # send time out information to the algorithm episode_length_s = 20 # episode length in seconds - terrain 子类定义了地形的类型、尺寸、摩擦系数、是否进行地形课程学习等
- cormands 子类定义了机器人可以接收的命令类型和范围
- init state 子类定义了机器人的初始位置、旋转、线速度、角速度和默认关节角度
- control 子类定义了控制类型、PD控制参数、动作缩放比例和控制动作更新的频率
class control: control_type = 'P' # P: position, V: velocity, T: torques # PD Drive parameters: stiffness = {'joint_a': 10.0, 'joint_b': 15.} # [N*m/rad] damping = {'joint_a': 1.0, 'joint_b': 1.5} # [N*m*s/rad] # action scale: target angle = actionScale * action + defaultAngle action_scale = 0.5 # decimation: Number of control action updates @ sim DT per policy DT decimation = 4 - asset 子类定义了机器人模型的文件路径、名称、脚的名称、重力设置等
- domain-rand 子类定义了领域随机化的参数，如是否随机化摩擦系数、基础质量、是否推动机器人等
- rewards 子类定义了奖励的计算方式和奖励的比例
class rewards: class scales: termination = -0.0 tracking_lin_vel = 1.0 tracking_ang_vel = 0.5 lin_vel_z = -2.0 ang_vel_xy = -0.05 orientation = -0. torques = -0.00001 dof_vel = -0. dof_acc = -2.5e-7 base_height = -0. feet_air_time = 1.0 collision = -1. feet_stumble = -0.0 action_rate = -0.01 stand_still = -0. only_positive_rewards = True # if true negative total rewards are clipped at zero (avoids early termination problems) tracking_sigma = 0.25 # tracking reward = exp(-error^2/sigma) soft_dof_pos_limit = 1. # percentage of urdf limits, values above this limit are penalized soft_dof_vel_limit = 1. soft_torque_limit = 1. base_height_target = 1. max_contact_force = 100. # forces above this value are penalized - normalization 子类定义了观察值和动作的归一化参数
class normalization: class obs_scales: lin_vel = 2.0 ang_vel = 0.25 dof_pos = 1.0 dof_vel = 0.05 height_measurements = 5.0 clip_observations = 100. clip_actions = 100. - noise子类定义了是否添加噪声以及噪声的级别
- viewer子类定义了观察者视角的位置和朝向。sim 子类定义了仿真的时间步长、子步数、重力加速度等
其次，LeggedRobotCfgPP0 类专门为使用PPO算法进行训练的配置，包括
种子、运行类名称、策略配置（如隐藏层维度、激活函数）
class LeggedRobotCfgPPO(BaseConfig):
seed = 1
runner_class_name = 'OnPolicyRunner'
class policy:
init_noise_std = 1.0
actor_hidden_dims = [512, 256, 128]
critic_hidden_dims = [512, 256, 128]
activation = 'elu' # can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
# only for 'ActorCriticRecurrent':
# rnn_type = 'lstm'
# rnn_hidden_size = 512
# rnn_num_layers = 1算法配置(如价值损失系数、学习率、折扣因子）
class algorithm:
# training params
value_loss_coef = 1.0
use_clipped_value_loss = True
clip_param = 0.2
entropy_coef = 0.01
num_learning_epochs = 5
num_mini_batches = 4 # mini batch size = num_envs*nsteps / nminibatches
learning_rate = 1.e-3 #5.e-4
schedule = 'adaptive' # could be adaptive, fixed
gamma = 0.99
lam = 0.95
desired_kl = 0.01
max_grad_norm = 1.运行配置（如每个环境的步数、最大迭代次数、保存间隔）等
class runner:
policy_class_name = 'ActorCritic'
algorithm_class_name = 'PPO'
num_steps_per_env = 24 # per iteration
max_iterations = 1500 # number of policy updates
# logging
save_interval = 50 # check for potential saves every this many iterations
experiment_name = 'test'
run_name = ''
# load and resume
resume = False
load_run = -1 # -1 = last run
checkpoint = -1 # -1 = last saved model
resume_path = None # updated from load_run and chkpt整体而言，这段代码提供了一个详细的配置框架，用于设置和调整多足机器人在仿真环境中的行为、控制和训练参数，以便进行有效的机器学习和仿真实验
1.2.2 envs/base/legged_robot.py
第51-226行代码定义了一个名为 LeggedRobot 的类，它继承自 BaseTask。这个类是为了模拟和控制一个有腿的机器人在一个仿真环境中的行为
以下是代码的主要组成部分和功能：
- 初始化（--init_—方法)：这个方法用于初始化 LeggedRobot 类的实例。它接收配置参数（如仿真参数、物理弓!擎、设备类型等），并根据这些参数设置仿真环境。此外，它还初始化了一些用于训练的PyTorch 缓冲区
- 步骤(step 方法)：这个方法用于在仿真中应用动作（例如，机器人的移动或旋转），并进行一步仿真。它处理动作的裁剪、物理仿真的执行、观察值的更新等
- 物理步骤后处理 (post_physics_step 方法)：在每次物理仿真步骤之后调用，用于检查是否需要重置环境（例如，机器人是否跌倒)计算观察值和奖励，并进行一些调试可视化（如果启用）
- 检查终止条件 (check_termination 方法)：用于检查是否有环境需要重置。这可能是因为接触力超过國值或仿真达到最大步数
- 重置索引（reset_idx 方法)：当某些环境需要重置时调用。它处理机器人状态的重置、命令的重新采样、以及一些缓冲区的重置
- 计算奖励( compute_reward 方法)：根据当前状态和行为计算奖励。它可能会调用多个奖励西数，并将每个奖劢项加到总奖励中
- 计算观察值 (compute_observations 方法)：根据当前的仿真状态计算观察值，这些观察值将用于后续的决策过程
接下来的第228-412行代码是一个用于仿真机器人运动和环境交互的Python类的一部分，主要用于创建和管理一个机器人仿真环境，它使用了NVIDIA IsaacGym库来进行物理仿真和环境渲染
以下是代码的主要功能和组成部分的详细解释：
- create-sim 方法：这个方法用于创建仿真环境、地形和其他环境元素。它首先设置了重力方向，然后根据配置创建不同类型的地开(如平面、高度场、三角形网格）。此外，它还负责创建仿真环境中的其他元素
- set-camera 方法：设置相机的位置和朝向，以便在仿真中观察机器人和环境
- -process-rigid_shape-props 方法：这是一个回调方法，用于在环境创建期问存储、更改或随机化每个环境中刚体形状的属性。例如，它可以随机化每个环境的摩擦系数
- -process_dof_props 方法：另一个回调方法，用于存储、更改或随机化每个环境中的自由度(DOF)属性。这些属性包括位置、速度和扭矩限制，可以根据URDF文件中定义的值进行设置
- -process_rigid body-props 方法：用于随机化基础质量等刚体属性的方法
- -post-physics-step_cal1back 方法：这是在每个物理步骤后调用的回调，用于计算终止条件、奖励和观察结果。例如，它可以基于目标和当前朝向计算角速度命令，测量地形高度，或在一定条件下随机推动机器人。-resample-commands 方法：为某些环境随机选择新的命令（如线速度、角速度等），用于控制机器人的运动
- -compute_torques 方法：根据给定的动作计算扭矩。动作可以解释为给PD（比例-微分）控制器的位置或速度目标，或直接作为缩放后的扭矩
- -reset_dofs 和 -reset_root_states 方法：这些方法用于重置选定环境中的DOF位置和速度，以及根状态的位置和速度。这通常在仿真的每个新回合开始时进行，以初始化机器人的状态
整体而言，这段代码展示了如何使用NVIDIA lsaac Gym库创建和管理一个复杂的机器人仿真环境，包括环境的创建、机器人的控制和状态管理等方面
第414-700行代码是一个用于仿真环境中的四足机器人控制和管理的Python类的一部分。它使用了PyTorch和NVIDIA Isaac Gym库来模拟和控制机器人在不同地形上的行为
以下是代码的主要功能和组成部分的概述：
- _push_robots：这个方法通过给机器人一个随机的基础速度来模拟沖击，用于域随机化训练，增强机器人 在不同情况下的鲁棒性
- -update_terrain_curriculum：实现了一种基于游戏的课程学习方法，根据机器人 在环境中的行走距离来调整它们所处的地形难度，以此来逐步提高机器人在复杂地形上的表现
- update_command_curriculum：根据机器人追踪速度的表现来调整指令的范围，是另一种课程学习策略，旨在逐步提高机器人对更复杂指令的响应能力
- -get_noise_scale_vec：设置一个向量，用于缩放添加到观测中的噪声。这是为了模拟真实世界中的不确定性，提高机器人的鲁棒性
- -init_buffers:初始化用于存储仿真状态和处理量的PyTorch张量。这包括机器人的根状态、关节状态、接触力等信息，这些信息对于控制和评估机器人的行为至关重要
- _prepare_reward_ function：准备奖励函数列表，这些函数将被调用来计算总奖励。这是训练过程中用于指导机器人学习的关键部分，通过奖励机器人达到期望的行为
- -create_ground_plane 和_create_heightfield：这两个方法用于在仿真环境中添加地面平面和高度场地开。这些地形是机器人训练和测试的场景，可以根据配置文件中的参数进行调整
第701-907行代码是一个用于模拟和控制机器人在不同地形上行走的Python类的一部分，特别是针对有腿的机器人。它使用了PyTorch库来处理数值计算，包括张量操作和自动微分
这个类包含了多个方法，用于设置环境、解析配置、绘制调试视觉效果、初始化测量点、计算地形高度、以及计算多种奖励函数。以下是主要功能的概述：
- _get_env_origins:根据地形类型(如高度场或三角网格）设置环境原点。对于粗糙地形，原点由地形平台定义；否则，创建一个网格布E
- -parse_cfg：解析配置文件，设置模拟的时间步长、观察值和奖励的缩放因子、命令范围等
- _draw_debug_vis:绘制调试视觉效果，如高度测量点，以帮助开发者理解模拟环境的状态
- _init_height_points:初始化用于采样地形高度的点
- _get_heights：根据机器人的位置和姿态，采样地形上特定点的高度
- reward*：一系列以-reward_ 开头的方法，计算不同类型的奖励 （或惩罚），包括线速度、角速度、姿态、基座高度、关节扭矩、关节速度、关节加速度、动作变化率、碰撞、终止条件、关节位置和速度限制、扭矩限制、线速度和角速度跟踪、脚步空中时间、绊倒、静止站立、脚部接触力等
这些方法共同工作，为机器人在复杂地形上行走提供仿真环境和控制策路的评估。通过计算奖励和惩罚，可以训练机器人学习如何有效地在不同地形上移动，同时保持平衡、避免碰撞，并尽可能高效地使用其动力系统
第二部分 机器人大脑HIT之detr
在读本部分之前，注意两点
- 关于HIT的原理，如本文开头所说，详见此文 《HumanPlus——斯坦福ALOHA团队开源的人形机器人：融合影子学习技术、RL、模仿学习》
- 关于目标检测器DETR，则详见此文《图像生成发展起源：从VAE、VQ-VAE、扩散模型DDPM、DETR到ViT、Swin transformer》的「第三部分 DETR：首次通过结合CNN+Transformer端对端解决object detection」
2.1 detr/models：涉及detr_vae、位置编码..
2.1.1 models/backbone.py
这段代码定义了一些类和函数，用于构建一个带有冻结批量归一化层的ResNet骨千网络，并将其与位置编码结合起来，最终构建一个用于特定任务的模型
- FrozenBatchNorm2d 类
FrozenBatchNorm2d 类继承自torch.nn.Module，实现了一个冻结的批量归一化层。与标准的批量归一化不同，这个类的批量统计数据和仿射参数是固定的。它的实现是从 torchvision.misc.ops 中复制的，并在rsqrt 之前添加了eps，以避免在使用某些 ResNet 模型时产生 NaN 值
分为
构造函数 _-init_——初始化了权重、偏置、运行均值和运行方差，并将它们注册为缓冲区
load_from_state_dict 方法在从状态字典加载模型参数时删除了num_ batches_tracked 键def __init__(self, n): super(FrozenBatchNorm2d, self).__init__() self.register_buffer("weight", torch.ones(n)) self.register_buffer("bias", torch.zeros(n)) self.register_buffer("running_mean", torch.zeros(n)) self.register_buffer("running_var", torch.ones(n))
forward 方法计算归一化后的输出，首先对权重、偏置、运行均值和运行方差进行 reshape，然后计算缩放和偏置，并应用于输入张量def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): num_batches_tracked_key = prefix + 'num_batches_tracked' if num_batches_tracked_key in state_dict: del state_dict[num_batches_tracked_key] super(FrozenBatchNorm2d, self)._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)def forward(self, x): # move reshapes to the beginning # to make it fuser-friendly w = self.weight.reshape(1, -1, 1, 1) b = self.bias.reshape(1, -1, 1, 1) rv = self.running_var.reshape(1, -1, 1, 1) rm = self.running_mean.reshape(1, -1, 1, 1) eps = 1e-5 scale = w * (rv + eps).rsqrt() bias = b - rm * scale return x * scale + bias - BackboneBase 类
BackboneBase 类继承自nn.Module，用于构建 ResNet 骨千网络的基础类
构造函数 --init--接收一个骨干网络、是否训练骨干网络、通道数和是否返回中间层作为参数，根据是否返回中间层，设置要返回的层，使用 IntermediiatelayerGetter 从骨干网络中获取指定的层
forward 方法将输入张量传递给骨千网络，并返回提取的特征def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool, name: str): super().__init__() # for name, parameter in backbone.named_parameters(): # only train later layers # TODO do we want this? # if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name: # parameter.requires_grad_(False) self.name = name if name.startswith('resnet'): if return_interm_layers: return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"} else: return_layers = {'layer4': "0"} self.body = IntermediateLayerGetter(backbone, return_layers=return_layers) self.num_channels = num_channelsdef forward(self, tensor): xs = self.body(tensor) return xs - Backbone 类
Backbone 类继承自 BackboneBase，实现了一个带有冻结批量归一化层的 ResNet 骨干网络
构造函数 --init-_接收模型名称、是否训练骨干网络、是否返回中问层和是否使用膨胀卷积作为参数
使用 getattr 从 torchvision.models 中获取指定的 ResNet 模型，并替换批量归一化层为FrozenBatchNorm2d
根据模型名称设置通道数
调用父类的构造函数初始化骨干网络 - Joiner 类
Joiner 类继承自nn. Sequential，将骨干网络和位置编码结合起来
构造函数-init_- 接收骨干网络和位置编码作为参数，并调用父类的构造函数
forward 方法将输入张量传递给骨干网络，获取特征，并将位置编码应用于特征，返回特征和位置编码 - build_backbone 函数
build_backbone 函数用于构建整个模型
接收一个包含配置参数的 args 对象
根据配置参数杓建位置编码
2.1.2 models/detr_vae.py：包含DETRVAE_ Decoder类、DETRVAE类、build
2.1.2.1 定义reparametrize、get_sinusoid_encoding_table
- 第一个函数reparametrize 用于从给定的均值（mu）和对数方差 （1ogvar) 中重新
参数化得到一个新的样本
具体来说，它通过计算标准差(std)，生成一个与标准差形状相同的正态分布随机变量 (eps)，然后返回 mu +std * eps 作为新的样本def reparametrize(mu, logvar): std = logvar.div(2).exp() eps = Variable(std.data.new(std.size()).normal_()) return mu + std * eps - 第二个函数 get_sinusoid_encoding_table 用于生成正弦编码表
它内部定义了一个辅助函数 get_position_angle_vec，该函数根据位置和隐藏维度计算角度向量def get_sinusoid_encoding_table(n_position, d_hid):
然后，get_sinusoid_encoding_table 使用这个辅助函数生成一个正弦编码表，其中偶数索引位置使用sin 函数，奇数索引位置使用cos 函数def get_position_angle_vec(position): return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]
最终返回一个PyTorch 的浮点张量sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)]) sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1return torch.FloatTensor(sinusoid_table).unsqueeze(0)
2.1.2.2 定义DETRVAE_ Decoder类、DETRVAE 类、CNNMLP 类
- 接下来是 DETRVAE_ Decoder 类，这是一个仅包含解码器的 Transformer 模型和多个线性层的深度学习模型
首先，构造函数接受多个参数，包括骨千网络(backbones )、Transformer 解码器(transformer_decoder）、状态维度 (state_dim）、查询数量 (num_queries)、相机名称class DETRVAE_Decoder(nn.Module):
（camera_names）和动作维度 (action_dim）。此外，还可以选择性地启用特征损失(feature_loss) .
在构造函数的开头，调用了父类的构造函数 superO•--init__（）。接着，初始化了一些基本属性，如num_queries、camera_names 和cam_num （相机数量)。然后，将传入的 Transformer 解码器赋值给 self.transformer_decoder，并设置状态维度和动作维度。def __init__(self, backbones, transformer_decoder, state_dim, num_queries, camera_names, action_dim, feature_loss=False):
接下来，计算 Transformer 解码器的隐藏维度 (hidden_dim ），并初始化多个线性层，包括action_head、proprio_head 和 is_pad_head，这些线性层分别用于动作预测、状态预测和填充标记预测super().__init__() self.num_queries = num_queries self.camera_names = camera_names self.cam_num = len(camera_names) self.transformer_decoder = transformer_decoder self.state_dim, self.action_dim = state_dim, action_dim
此外，还初始化了一个查询嵌入层 (query-embed )，用于生成查询向量
再其次hidden_dim = transformer_decoder.d_model self.action_head = nn.Linear(hidden_dim, action_dim) self.proprio_head = nn.Linear(hidden_dim, state_dim) self.is_pad_head = nn.Linear(hidden_dim, 1) self.query_embed = nn.Embedding(num_queries, hidden_dim)如果提供了骨干网络，则初始化一个卷积层（input_proj） 用于将骨干网络的输出投影到隐藏维度，并将骨千网络列表转换为 nn.Modulelist
同时，还初始化了一个线性层(input_proj-robot_state)用于将机器人状态投影到隐藏维度if backbones is not None: self.input_proj = nn.Conv2d(backbones[0].num_channels, hidden_dim, kernel_size=1) self.backbones = nn.ModuleList(backbones) self.input_proj_robot_state = nn.Linear(state_dim, hidden_dim)
最后，使用 上面定义的 get_sinusoid_encoding_table 函数生成一个正弦编码表，并将其注册为缓冲区(pos_table）else: # input_dim = 14 + 7 # robot_state + env_state self.input_proj_robot_state = nn.Linear(state_dim, hidden_dim) self.input_proj_env_state = nn.Linear(7, hidden_dim) self.pos = torch.nn.Embedding(2, hidden_dim) self.backbones = None
此外，还初始化了一个额外的位置嵌入层 (additional-pos_embed ），用于处理本体感知和潜在变量的嵌入。最后，将feature _1oss 参数赋值给 self.feature_1oss
至于，forward 方法则定义了模型的前向传播逻辑，包括处理图像特征、位置嵌入和前馈网络的计算# encoder extra parameters self.register_buffer('pos_table', get_sinusoid_encoding_table(1+1+num_queries, hidden_dim)) # [CLS], qpos, a_seq self.additional_pos_embed = nn.Embedding(1, hidden_dim) # learned position embedding for proprio and latent self.feature_loss = feature_loss - DETRVAE 类是一个完整的 DETR 模型，用于执行目标检测
它的构造函数与DETRVAE_Decoder 类似，但增加了编码器和潜在变量的处理逻辑encode 方法用于从动作序列中获取潜在变量，并根据是否使用向量量化 (VQ）来处理潜在变量class DETRVAE(nn.Module): """ This is the DETR module that performs object detection """
具体而言
该encode 方法用于对输入的关节位置(qpos）和动作序列(actions）进行编码，生成潜在变量 (latent variables)。该方法首先获取输入 qpos 的批次大小(bs）
如果编码器 (self.encoder） 为空，则生成一个全零的潜在样本，并通过线性层投影到潜在输入空间，同时将概率、二进制编码、均值和对数方差设置为 Nonedef encode(self, qpos, actions=None, is_pad=None, vq_sample=None): bs, _ = qpos.shape
如果编码器不为空，则根据是否提供了动作序列来判断当前是训练模式还是验证模式if self.encoder is None: latent_sample = torch.zeros([bs, self.latent_dim], dtype=torch.float32).to(qpos.device) latent_input = self.latent_out_proj(latent_sample) probs = binaries = mu = logvar = None
在训练模式下，首先将动作序列投影到嵌入维度，并与关节位置的嵌入向量和一个分类(CLS）标记向量进行拼接else: # cvae encoder is_training = actions is not None # train or val
然后，将拼接后的输入进行转置，以符合编码器的输入格式
接着，生成一个不包含填充标记的掩码### Obtain latent z from action sequence if is_training: # project action sequence to embedding dim, and concat with a CLS token action_embed = self.encoder_action_proj(actions) # (bs, seq, hidden_dim) qpos_embed = self.encoder_joint_proj(qpos) # (bs, hidden_dim) qpos_embed = torch.unsqueeze(qpos_embed, axis=1) # (bs, 1, hidden_dim) cls_embed = self.cls_embed.weight # (1, hidden_dim) cls_embed = torch.unsqueeze(cls_embed, axis=0).repeat(bs, 1, 1) # (bs, 1, hidden_dim) encoder_input = torch.cat([cls_embed, qpos_embed, action_embed], axis=1) # (bs, seq+1, hidden_dim) encoder_input = encoder_input.permute(1, 0, 2) # (seq+1, bs, hidden_dim)
并获取位置嵌入向量# do not mask cls token cls_joint_is_pad = torch.full((bs, 2), False).to(qpos.device) # False: not a padding is_pad = torch.cat([cls_joint_is_pad, is_pad], axis=1) # (bs, seq+1)
将这些输入传递给编码器，获取编码器的输出，并通过线性层投影到潜在信息空问# obtain position embedding pos_embed = self.pos_table.clone().detach() pos_embed = pos_embed.permute(1, 0, 2) # (seq+1, 1, hidden_dim)
如果使用向量量化 (VQ)，则对潜在信息进行重塑和软最大化(softmax）操作，生成概率分布# query model encoder_output = self.encoder(encoder_input, pos=pos_embed, src_key_padding_mask=is_pad) encoder_output = encoder_output[0] # take cls output only latent_info = self.latent_proj(encoder_output)
然后，通过多项式采样和独热编码 (one-hot encoding)生成二进制编码，并通过直通估计 (straight-throughestimation）生成潜在输入if self.vq: logits = latent_info.reshape([*latent_info.shape[:-1], self.vq_class, self.vq_dim]) probs = torch.softmax(logits, dim=-1)
否则，直接使用重新参数化技巧 (reparametrize） 从均值和对数方差生成潜在样本，并通过线性层投影到潜在输入空间binaries = F.one_hot(torch.multinomial(probs.view(-1, self.vq_dim), 1).squeeze(-1), self.vq_dim).view(-1, self.vq_class, self.vq_dim).float() binaries_flat = binaries.view(-1, self.vq_class * self.vq_dim) probs_flat = probs.view(-1, self.vq_class * self.vq_dim) straigt_through = binaries_flat - probs_flat.detach() + probs_flat latent_input = self.latent_out_proj(straigt_through) mu = logvar = None
在验证模式下，如果使用向量量化，则直接将提供的 vQ 样本投影到潜在输入空间else: probs = binaries = None mu = latent_info[:, :self.latent_dim] logvar = latent_info[:, self.latent_dim:] latent_sample = reparametrize(mu, logvar) latent_input = self.latent_out_proj(latent_sample)
否则，生成一个全零的潜在样本，并通过线性层投影到潜在输入空间else: mu = logvar = binaries = probs = None if self.vq: latent_input = self.latent_out_proj(vq_sample.view(-1, self.vq_class * self.vq_dim))
最终，该方法返回潜在输入、概率分布、二进制编码、均值和对数方差else: latent_sample = torch.zeros([bs, self.latent_dim], dtype=torch.float32).to(qpos.device) latent_input = self.latent_out_proj(latent_sample)
至于forward方法 和上面的类似，其定义了模型的前向传播逻辑，包括处理图像特征、位置嵌入和前馈网络的计算return latent_input, probs, binaries, mu, logvar - 最后是 CNNMLP 类，这是一个结合了卷积神经网络(CNN）和多层感知器 (MLP）的模型
它的构造函数初始化了多个卷积层和线性层，并根据提供的骨千网络设置不同的输入投影层class CNNMLP(nn.Module):
forward 方法定义了模型的前向传播逻辑，包括处理图像特征、位置嵌入和前馈网络的计算
这些类和函数共同构成了一个复杂的深度学习模型框架，能够处理图像和状态信息，并生成相应的动作预测
2.1.2.3 定义MLP、build_encoder、build、build_cnnmlp
- 首先，mLp 函数用于构建一个多层感知器（MLP)
该函数接受输入维度、隐藏层维度、输出维度和隐藏层深度作为参数。如果隐藏层深度为 0，则直接创建一个线性层将输入维度映射到输出维度
否则，首先创建一个线性层和一个 ReLU 激活函数，然后根据隐藏层深度添加多个线性层和 ReLU 激活函数def mlp(input_dim, hidden_dim, output_dim, hidden_depth): if hidden_depth == 0: mods = [nn.Linear(input_dim, output_dim)]
最后再添加一个线性层将隐藏层维度映射到输出维度else: mods = [nn.Linear(input_dim, hidden_dim), nn.ReLU(inplace=True)] for i in range(hidden_depth - 1): mods += [nn.Linear(hidden_dim, hidden_dim), nn.ReLU(inplace=True)]
所有这些层被封装在一个 nn. Sequential模块中并返回mods.append(nn.Linear(hidden_dim, output_dim))trunk = nn.Sequential(*mods) return trunk - 接下来，build_encoder 函数用于构建一个 Transformer 编码器
该函数从 args 参数中获取模型的超参数，如隐藏维度、dropout 率、注意力头数、前馈网络维度、编码器层数和是否在归一化之前进行计算
然后，创建一个 TransformerEncoderLayer 实例，并根据是否需要归一化创建一个nn.LayerNorm 实例def build_encoder(args): d_model = args.hidden_dim # 256 dropout = args.dropout # 0.1 nhead = args.nheads # 8 dim_feedforward = args.dim_feedforward # 2048 num_encoder_layers = args.enc_layers # 4 # TODO shared with VAE decoder normalize_before = args.pre_norm # False activation = "relu"
最后，使用这些层创建一个 TransformerEncoder 实例并返回encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation, normalize_before) encoder_norm = nn.LayerNorm(d_model) if normalize_before else Noneencoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) return encoder - build 函数用于根据提供的参数构建不同类型的模型
首先，从args 中获取状态维度和相机名称列表
如果 args.same_backbones 为真，则为所有相机使用相同的骨千网络；否则，为每个相机创建一个单独的骨干网络
接着，根据 args.no_encoder 判断是否需要创建编码器def build(args): state_dim = args.state_dim # TODO hardcode # From state # backbone = None # from state for now, no need for conv nets # From image backbones = [] if args.same_backbones: backbone = build_backbone(args) backbones = [backbone] else: for _ in args.camera_names: backbone = build_backbone(args) backbones.append(backbone)
如果args.model_type 为"ACT"，则创建一个 DETRVAE 模型if args.no_encoder: encoder = None else: encoder = build_transformer(args)
如果为 "HIT"，则创建一个DETRVAE_ Decoder 模型，就像此文《HumanPlus——斯坦福ALOHA团队开源的人形机器人：融合影子学习技术、RL、模仿学习》「1.2.2 机器人大脑——用于模仿学习的仅解码器Transformer(HIT)：预测机器人的目标姿态」一节中所说的：“把ACT原本的transformer编码器-解码器架构改成了transformer仅解码器架构”if args.model_type=="ACT": transformer = build_transformer(args) model = DETRVAE( backbones, transformer, encoder, state_dim=state_dim, num_queries=args.num_queries, camera_names=args.camera_names, vq=args.vq, vq_class=args.vq_class, vq_dim=args.vq_dim, action_dim=args.action_dim, )
最后，计算模型的参数数量并打印出来，然后返回模型实例elif args.model_type=="HIT": transformer_decoder = build_transformer_decoder(args) model = DETRVAE_Decoder( backbones, transformer_decoder, state_dim=state_dim, num_queries=args.num_queries, camera_names=args.camera_names, action_dim=args.action_dim, feature_loss= args.feature_loss if hasattr(args, 'feature_loss') else False, )n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) print("number of parameters: %.2fM" % (n_parameters/1e6,)) return model - 最后，build_cnnmlp函数用于构建一个 CNN-MLP 模型。该函数硬编码了状态维度为 14，为每个相机创建一个骨千网络
然后，使用这些骨干网络和状态维度创建一个CNNMLP 模型def build_cnnmlp(args): state_dim = 14 # TODO hardcode # From state # backbone = None # from state for now, no need for conv nets # From image backbones = [] for _ in args.camera_names: backbone = build_backbone(args) backbones.append(backbone)
最后，计算模型的参数数量并打印出来，然后返回模型实例model = CNNMLP( backbones, state_dim=state_dim, camera_names=args.camera_names, )n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) print("number of parameters: %.2fM" % (n_parameters/1e6,)) return model
2.1.3 models/position_encoding.py
这个代码片段定义了两种位置嵌入 (Position Embedding）方法以及一个构建位置嵌入的函数
2.1.3.1 第一种位置嵌入PositionEmbeddingSine
首先，PositionEmbeddingSine 类实现了一种基于正弦函数的位置嵌入方法。这种方法类似于
transformer论文中使用的位置嵌入，但进行了扩展以适用于图像数据
- 在初始化-init--方法中，首先调用了父类nn. Module 的初始化方法
然后，设置了几个关键参数：class PositionEmbeddingSine(nn.Module): """ This is a more standard version of the position embedding, very similar to the one used by the Attention is all you need paper, generalized to work on images. """ def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
num_pos_feats 表示位置特征的数量，默认为 64；
temperature 是一个温度参数，默认为 10000；
normalize 表示是否归一化，默认为 False ;
scale 是一个可选的缩放因子
如果提供了 scale 但未启用归一化，则会抛出一个 ValueError 异常，提示 "normalize should be True if scale is passed”def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): super().__init__() self.num_pos_feats = num_pos_feats self.temperature = temperature self.normalize = normalize
如果未提供 scale，则默认将其设置为 2* math.pi，然后将这些参数存储为类的属性
这个类的设计目的是为图像数据生成位置嵌入，以便在 Transformer 模型中使用。通过这些参数的配置，可以灵活地调整位置嵌入的特性，以适应不同的应用场景。if scale is not None and normalize is False: raise ValueError("normalize should be True if scale is passed") if scale is None: scale = 2 * math.pi self.scale = scale - 在forward 方法中，首先计算了沿着高度和宽度的累积和（cumsum ）
然后根据是否归一化对这些累积和进行缩放def forward(self, tensor): x = tensor # mask = tensor_list.mask # assert mask is not None # not_mask = ~mask not_mask = torch.ones_like(x[0, [0]]) y_embed = not_mask.cumsum(1, dtype=torch.float32) x_embed = not_mask.cumsum(2, dtype=torch.float32)
接着，计算位置嵌入的正弦和余弦值，并将它们拼接在一起，最终返回位置嵌入张量if self.normalize: eps = 1e-6 y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) pos_x = x_embed[:, :, :, None] / dim_t pos_y = y_embed[:, :, :, None] / dim_tpos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) return pos
对于上面的内容，还是值得我再好好解释说明下，比如上面说，这个位置编码方法和transformer论文中的位置编码是相似的，那到底相似的是什么，不相似的又是什么呢？
- 首先，对于transformer论文中的位置编码公式，如此文《一文通透位置编码：从标准位置编码、旋转位置编码RoPE到ALiBi、LLaMA 2 Long(含NTK-aware简介)》的第一部分所述
对于每个位置 pos 和维度 2i 和 2i+1，位置编码的计算为：
其中，指的是模型的维度
- 而上面代码中的实现而言，则相对transformer论文中原始的位置编码方法做了两个小改动
具体而言
首先，代码计算了累加和（cumsum），得到和
，并没有直接用 pos，且代码中对累加和进行了归一化处理，避免了直接使用累加和的绝对值。这是Transformer论文中没有的步骤
- 其次，代码使用了一个温度参数，影响了位置编码的频率
比如，dim_t的生成过程如下所示
这行代码生成一个张量：dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
假设 num_pos_feats=d_model，那么 dim_t 将是从 0 到 d_model −1 的整数序列
- 接下来，计算缩放因子
其中，右上角的分子dim_t // 2，意味着对 dim_t 中的每个元素除以 2。这是为了确保偶数和奇数维度的处理
而下面式子
将每个维度的索引映射到一个范围，这个范围是和特征维度 d_model相关的
总之，temperature 是一个超参数，通过调整它的值，模型可以学习到不同频率的位置信息。例如，较大的温度值会降低频率，使得位置编码变化较慢；较小的温度值会增加频率，使得位置编码变化更快 - 综上所述，dim_t 用来生成一个缩放因子，这个因子在计算位置编码时影响正弦和余弦函数的频率。最终生成的位置编码公式可以表示为：
总之，dim_t的计算通过引入温度参数和特征维度，帮助调整正弦和余弦函数的频率，从而影响模型对位置信息的学习能力。这种方法比直接使用固定频率的方式更灵活，可以更好地适应不同任务的需求
2.1.3.2 第二种位置嵌入PositionEmbeddingLearned
其次，PositionEmbeddingLearned 类实现了一种基于学习的位置嵌入方法
- 在初始化方法_-init.中，定义了行嵌入和列嵌入的嵌入层，并调用 reset_parameters 方法对嵌入层的权重进行初始化
- 在forward 方法中，首先获取输入张量的高度和宽度，然后生成行和列的索引，并通过嵌入层将这些索引映射到嵌入向量，且将行嵌入和列嵌入拼接在一起，返回位置嵌入张量
最后，build_position_encoding 函数根据提供的参数构建位置嵌入对象
- 首先，根据隐藏维度计算位置特征的数量(N_steps)
def build_position_encoding(args): N_steps = args.hidden_dim // 2 - 然后，根据 args. position_embedding 的值选择使用PositionEmbeddingSine 还是PositionEmbeddingLearned
如果 args.position_embedding 的值不在支持的范围内，则拋出一个 ValueErrorif args.position_embedding in ('v2', 'sine'): # TODO find a better way of exposing other arguments position_embedding = PositionEmbeddingSine(N_steps, normalize=True) elif args.position_embedding in ('v3', 'learned'): position_embedding = PositionEmbeddingLearned(N_steps)
最终，返回构建好的位置嵌入对象else: raise ValueError(f"not supported {args.position_embedding}")return position_embedding
2.1.4 基于Transformer实现简易GPT：来自karpathy/nanoGPT的实现
Transformer_Block 类是一个继承自nn.Module 的模块，用于实现 Transformer 模型中的一个基本块。其构造函数接受多个参数，包括 Latent_dim （潜在维度)、num head（多头注意力的头数）、dropout_rate (dropout 率)、self-attention（是否使用自注意力机制）和 query-num (查询数量）
class Transformer_Block(nn.Module):
def __init__(self, latent_dim, num_head, dropout_rate, self_attention=True, query_num=50) -> None:在初始化过程中，定义了多个层和模块，包括三个LayerNorm 层、一个多头注意力层和一个多层感知机（MLP)
super().__init__()
self.num_head = num_head
self.latent_dim = latent_dim
self.ln_1 = nn.LayerNorm(latent_dim)
self.attn = nn.MultiheadAttention(latent_dim, num_head, dropout=dropout_rate)
self.ln_2 = nn.LayerNorm(latent_dim)
# MLP 由两个线性层和一个 GELU 激活函数组成，并在最后添加了一个dropout 层
self.mlp = nn.Sequential(
# 第一个线性层将输入维度从 Latent_dim 扩展到 4*Latent_dim
nn.Linear(latent_dim, 4 * latent_dim),
nn.GELU(),
# 再通过第二个线性层将维度缩小回 Latent_dim
nn.Linear(4 * latent_dim, latent_dim),
nn.Dropout(dropout_rate),
)
# 此外，还定义了一个dropout 层 self.dropout1
self.ln_3 = nn.LayerNorm(latent_dim)
self.dropout1 = nn.Dropout(dropout_rate)
# 并将query-num 和 self_attention 参数分别赋值给实例变量 self.query-num 和 self.self_attention
# 这些初始化步骤为 Transformer 模块的前向传播过程提供了必要的组件和参数
self.query_num = query_num
self.self_attention = self_attention在 forward 方法中，根据 self_attention 的值决定执行不同的操作
- 如果 self_attention 为真，则对输入 × 进行 LayerNorm 归一化，然后通过多头注意力层计算注意力输出 x2，并将其与输入× 相加。接着，再次进行 LayerNorm 归一化，并通过 MLP处理，最后返回处理后的x
def forward(self, x):
if self.self_attention:
x = self.ln_1(x)
x2 = self.attn(x, x, x, need_weights=False)[0]
x = x + self.dropout1(x2)
x = self.ln_2(x)
x = x + self.mlp(x)
x = self.ln_3(x)
return x- 如果self-attention 为假，则将输入x 分为 xaction 和 x-condition 两部分，分别表示动作和条件。然后对 xaction 和 x_condition 进行多头注意力计算，并通过类似的步骤进行处理，最后将处理后的 x2 与 x-condition 拼接，返回拼接后的结果
else:
x = self.ln_1(x)
x_action = x[-self.query_num:].clone()
x_condition = x[:-self.query_num].clone()
x2 = self.attn(x_action, x_condition, x_condition, need_weights=False)[0]
x2 = x2 + self.dropout1(x2)
x2 = self.ln_2(x2)
x2 = x2 + self.mlp(x2)
x2 = self.ln_3(x2)
x = torch.cat((x_condition, x2), dim=0)
return xTransformer_BERT 类同样继承自nn. module，用于实现一个基于 Transformer 的 BERT 模型。其构造函数接受多个参数，包括context_len （上下文长度)、latent_dim 、mum_head 、mum_layer（层数）、dropout_rate、use_pos_embd_image（是否使用图像位置嵌入）、use-pos_embdaction（是否使用动作位置嵌入）和 query-num。在初始化过程中，根据use_pos_embd_image 和use-pos_embd_action 的值，决定是否初始化位置嵌入层weight-pos_embed
此外，还定义了一个由多个 Transformer_Block 组成的序列模块attention_blocks
在 forward 方法中，根据 use_ pos_embd _ image 和 use_ pos_embd_action 的值，对输入 × 进行相应的处理
- 如果两者都为假，则直接将位置嵌入加到 x 上
def forward(self, x, pos_embd_image=None, query_embed=None):
if not self.use_pos_embd_image and not self.use_pos_embd_action: #everything learned - severe overfitting
x = x + self.weight_pos_embed.weight[:, None]- 如果仅 use_pos_embd_image 为真，则对×的不同部分分别加上位置嵌入
elif self.use_pos_embd_image and not self.use_pos_embd_action: #use learned positional embedding for action
x[-self.query_num:] = x[-self.query_num:] + self.weight_pos_embed.weight[:, None]
x[:-self.query_num] = x[:-self.query_num] + pos_embd_image- 如果两者都为真，则使用正弦位置嵌入
elif self.use_pos_embd_action and self.use_pos_embd_image: #all use sinsoidal positional embedding
x[-self.query_num:] = x[-self.query_num:] + query_embed
x[:-self.query_num] = x[:-self.query_num] + pos_embd_image 处理完位置嵌入后，将x传递给 attention_blocks 进行进一步处理，最后返回处理后的x
x = self.attention_blocks(x)
# take the last token
return x2.1.5 models/transformer.py
Transformer_decoder 类是一个基于PyTorch的神经网络模块，用于实现Transformer解码器。它的构造函数接受多个参数，如上下文度、模型维度、注意力头数、解码层数、dropout率、是否使用图像位置嵌入、查询数量、是否使用动作位置嵌入和自注意力标志
- 构造函数中，首先调用父类的构造函数，然后初始化一个 Transformer_BERT 解码器实例，并调用_reset-parameters 方法来初始化模型参数。该方法使用Xavier均匀分布来初始化参数
- 在 forward方法中，首先检查输入的形状，如果输入是四维的(即包含高度和宽度)，则将其展平井调整维度顺序。然后，将查询嵌入和位置嵌入扩展到与批次大小匹配，并将额外的位置嵌入与位置嵌入拼接。接着，将输入张量与动作输入标记拼接，并传递给解码器。最后，返回解码器输出的转置结果
Transformer 类是一个完整的Transformer模型，包括编码器和解码器
- 构造函数中，首先初始化编码器层和解码器层，并使用_reset_parameters 方法初始化参数
- 在forward 方法中，同样检查输入的形状并进行相应的处理，然后将输入传递给编码器和解码器，最后返回解码器输出的转置结果
TransformerEncoder 类实现了 Transformer编码器，它包含多个编码器层
- 构造函数中，使用-get-clones 方法克隆多个编码器层，并初始化层数和归一化层
- 在forward 方法中，依次通过每一层编码器处理输入，并在最后应用归一化层（如果存在）
TransformerDecoder 类实现了Transformer解码器，它包含多个解码器层
- 构造函数中，使用-get-clones 方法克隆多个解码器层，并初始化层数、归一化层和是否返回中间结果的标志
- 在forward 方法中，依次通过每一层解码器处理输入，并在最后应用归一化层（如果存在）。如果设置了返回中问结果的标志，则返回所有中间结果的堆叠张量，否则返回解码器输出的扩展结果
TransformerEncoderLayer 类实现了单个Transformer编码器层，包括自注意力机制和前馈神经网络
- 构造函数中，初始化自注意力层、前馈神经网络层、归一化层和dropout层，并根据激活函数名称获取相应的激活函数
- 在forward 方法中，根据是否在归一化之前进行处理，选择不同的前向传播路径
TransformerDecoderLayer 类是一个基于PyTorch的神经网络模块，用于实现Transformer解码器层
- 它继承自nn.Module，并在构造函数中初始化了多个子模块，包括自注意力层、多头注意力层、前馈神经网络层、归一化层和dropout层
class TransformerDecoderLayer(nn.Module): def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu", normalize_before=False): - 构造函数接受多个参数，如模型维度 dmodel、注意力头数nhead 、前馈神经网络的维度dim_ feedforward 、dropout率 dropout、激活函数 activation 和是否在归一化之前进行处理的标志normalize_before
其中，自注意力层和多头注意力层使用 m.MultiheadAttention 进行初始化
前馈神经网络层由两个线性层和一个dropout层组成super().__init__() self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
归一化层使用nn.LayerNorm 进行初始化# Implementation of Feedforward model self.linear1 = nn.Linear(d_model, dim_feedforward) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model)
激活函数通过调用_get_activation_fn 函数获取self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.dropout3 = nn.Dropout(dropout)self.activation = _get_activation_fn(activation) self.normalize_before = normalize_before - with-pos_embed 方法用于将位置嵌入添加到输入张量中，如果位置嵌入为 None，则直接返回输入张量
def with_pos_embed(self, tensor, pos: Optional[Tensor]): return tensor if pos is None else tensor + pos - forward-post 方法实现了在归一化之后进行处理的前向传播路径
首先，将目标张量 tgt 和查询位置嵌入query-pos 相加，得到查询和键def forward_post(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None):
然后，通过自注意力层计算注意力输出，并进行dropout和归一化处理q = k = self.with_pos_embed(tgt, query_pos)
接着，通过多头注意力层计算目标张量和记忆张量 memory 之间的注意力输出，并进行dropout和归一化处理tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0] tgt = tgt + self.dropout1(tgt2) tgt = self.norm1(tgt)
最后，通过前馈神经网络层进行非线性变换，并进行dropout和归一化处理，返回最终的目标张量tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos), key=self.with_pos_embed(memory, pos), value=memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)[0] tgt = tgt + self.dropout2(tgt2) tgt = self.norm2(tgt)tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) tgt = tgt + self.dropout3(tgt2) tgt = self.norm3(tgt) return tgt - forward_pre 方法实现了在归一化之前进行处理的前向传播路径。与forward_post 方法类似，只是在每一步处理之前先进行归一化
def forward_pre(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None): tgt2 = self.norm1(tgt) q = k = self.with_pos_embed(tgt2, query_pos) tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0] tgt = tgt + self.dropout1(tgt2) tgt2 = self.norm2(tgt) tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos), key=self.with_pos_embed(memory, pos), value=memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)[0] tgt = tgt + self.dropout2(tgt2) tgt2 = self.norm3(tgt) tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) tgt = tgt + self.dropout3(tgt2) return tgt - forward 方法根据 normalize_before 标志选择调用 forward_pre 或forward_post 方法，完成前向传播
def forward(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None): if self.normalize_before: return self.forward_pre(tgt, memory, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) return self.forward_post(tgt, memory, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
然后这个detr/models/transformer.py代码文件还实现了以下4个函数
- -get-clones 函数用于克隆指定数量的模块，并返回一个包含这些模块的nn. Modulelist
- build_transformer 函数根据给定的参数构建并返回一个Transformer 实例
- build_transformer_decoder 函数根据给定的参数构建并返回一个 Transformer_decoder 实例
- -get_activation_fn 函数根据给定的激活函数名称返回相应的激活函数，如果名称无效，则抛出运行时错误
2.2 detr/main.py
2.2.1 main.py：get_args_parser
2.2.2 main.py：build_ACT_model_and_optimizer：最终调用build
这个函数 build_ACT_model_and_optimizer 的主要目的是构建一个ACT模型及其对应的优化器
该函数首先解析命令行参数，然后根据传入的覆盖参数 (args_override )更新这些参数，最后构建模型并配置优化器
- 首先，函数创建了一个命令行参数解析器 argparse.ArgumentParser，并将父解析器设置为 get_args_parser 函数返回的解析器
def build_ACT_model_and_optimizer(args_override): parser = argparse.ArgumentParser('DETR training and evaluation script', parents=[get_args_parser()]) - 然后，解析命令行参数并将结果存储在 args对象中
#存储 args = parser.parse_args() - 接下来，函数遍历 args_override 宇典，将每个键值对设置到 args 对象中。如果设置过程中出现错误，会打印错误信息
for k, v in args_override.items(): print(f"Setting {k} to {v}") try: setattr(args, k, v) except: print(f"Error setting {k} to {v}") - 然后，函数调用 build_ACT_model 函数构建模型，并将模型移动到GPU上
model = build_ACT_model(args) model.cuda() - 接下来，函数创建了一个参数字典列表 param_dicts，其中包含模型中所有需要梯度更新的参数
参数分为两类：一类是不包含"backbone"关键字的参数，另一类是包含"backbone"关键字的参数，并为后者设置了单独的学习率 args. lr_backboneparam_dicts = [ {"params": [p for n, p in model.named_parameters() if "backbone" not in n and p.requires_grad]}, { "params": [p for n, p in model.named_parameters() if "backbone" in n and p.requires_grad], "lr": args.lr_backbone, }, ] - 最后，函数使用 torch. optim.Adamw优化器，并将参数字典列表、学习率 args.lr和权重衰减 args.weight_decay 传递给优化器
optimizer = torch.optim.AdamW(param_dicts, lr=args.lr, weight_decay=args.weight_decay) - 最终，函数返回构建好的模型和优化器。这个函数的设计使得模型和优化器的构建过程高度灵活，可以根据不同的参数配置进行调整
return model, optimizer
可以看到，该函数调用了build_ACT_model
def build_ACT_model(args):
return build_vae(args)而build_vae的实现是来自上文2.1.2节中分析过的HIT/detr/models/detr_vae.py，虽然上面已经详细分析过了，不过 就权当复习一遍，再次分析下这个build的实现
这个函数build 的主要目的是根据传入的参数构建一个模型，该函数首先根据参数配置构建模型的各个组件，然后根据模型类型选择不同的模型架构，最后返回构建好的模型
- 首先，函数从参数中获取状态维度 state_dim
def build(args): state_dim = args.state_dim # TODO hardcode- 接下来，根据是否使用相同的骨千网络(backbone)，构建一个或多个骨千网络
如果参数args.same_backbones 为 True，则只构建一个骨千网络并将其添加到列表中；否则，为每个相机名称构建一个骨千网络并将其添加到列表中# From state # backbone = None # from state for now, no need for conv nets # From image backbones = []if args.same_backbones: backbone = build_backbone(args) backbones = [backbone] else: for _ in args.camera_names: backbone = build_backbone(args) backbones.append(backbone)- 然后，函数根据参数 args.no_encoder 决定是否构建编码器
如果参数为 True，则不构建编码器；否则，调用build_transformer 函数构建编码器if args.no_encoder: encoder = None else: encoder = build_transformer(args)- 接下来，根据模型类型 args.model-type，选择不同的模型架构
如果模型类型为“ACT"，则构建一个包含变压器 (transformer)和编码器的 DETRVAE 模型如果模型类型为“HIT"，则构建一个只包含变压器解码器的DETRVAE_Decoder 模型if args.model_type=="ACT": transformer = build_transformer(args) model = DETRVAE( backbones, transformer, encoder, state_dim=state_dim, num_queries=args.num_queries, camera_names=args.camera_names, vq=args.vq, vq_class=args.vq_class, vq_dim=args.vq_dim, action_dim=args.action_dim, )elif args.model_type=="HIT": transformer_decoder = build_transformer_decoder(args) model = DETRVAE_Decoder( backbones, transformer_decoder, state_dim=state_dim, num_queries=args.num_queries, camera_names=args.camera_names, action_dim=args.action_dim, feature_loss= args.feature_loss if hasattr(args, 'feature_loss') else False, )- 最后，函数计算模型中需要梯度更新的参数总数，并打印出来。然后返回构建好的模型。这个函数的设计使得模型的构建过程高度灵活，可以根据不同的参数配置构建不同的模型架构
n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) print("number of parameters: %.2fM" % (n_parameters/1e6,)) return model
2.2.3 main.py：build_CNNMLP_model_and_optimizer
// 待更
第三部分 机器人大脑HIT之其他：imitate_episodes_h1_train、model_util、policy
3.1 HIT/imitate_episodes_h1_train.py
forward_pass 函数负责将输入数据传递给策路模型进行前向传播。它接受两个参数：数据和策略模型
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()
return policy(qpos_data, image_data, action_data, is_pad) # TODO remove None- 数据包括图像数据、位置数据、动作数据和填充标志
- 函数首先将这些数据移动到GPU上，然后将它们传递给策略模型，并返回模型的输出
3.1.1 def train_bc(train_dataloader, val_dataloader, config)
train_bc函数是训练行为克隆 (Behavior Cloning）模型的主要函数，它接受训练数据加载器、验证数据加载器和配置参数
def train_bc(train_dataloader, val_dataloader, config):函数首先从配置中提取训练步骤数、检查点目录、随机种子、策略类、策路配置、验证间隔和保存问隔等参数
num_steps = config['num_steps']
ckpt_dir = config['ckpt_dir']
seed = config['seed']
policy_class = config['policy_class']
policy_config = config['policy_config']
validate_every = config['validate_every']
save_every = config['save_every']然后，它设置随机种子并创建策略模型，如果配置中指定了预训练模型或恢复检查点路径，则加载相应的模型权重
policy = make_policy(policy_class, policy_config)
if config['load_pretrain']:
loading_status = policy.deserialize(torch.load(f'{config["pretrained_path"]}/policy_last.ckpt', map_location='cuda'))
print(f'loaded! {loading_status}')
if config['resume_ckpt_path'] is not None:
loading_status = policy.deserialize(torch.load(config['resume_ckpt_path']))
print(f'Resume policy from: {config["resume_ckpt_path"]}, Status: {loading_status}')接着，将策略模型移动到GPU上，并创建优化器
policy.cuda()
optimizer = make_optimizer(policy_class, policy)如果加载了预训练模型，则加载相应的优化器状态。在训练过程中，函数使用repeater 函数将训练数据加载器转换为一个无限循环的数据加载器
if config['load_pretrain']:
optimizer.load_state_dict(torch.load(f'{config["pretrained_path"]}/optimizer_last.ckpt', map_location='cuda'))
min_val_loss = np.inf
best_ckpt_info = None
train_dataloader = repeater(train_dataloader)每个训练步骤中，首先进行验证(如果当前步骤是验证间隔的倍数）
for step in tqdm(range(num_steps+1)):
if step % validate_every == 0:
print('validating')在验证过程中，策略模型切换到评估模式，并对验证数据进行前向传播，计算验证损失
with torch.inference_mode():
policy.eval()
validation_dicts = []
for batch_idx, data in enumerate(val_dataloader):
forward_dict = forward_pass(data, policy)
validation_dicts.append(forward_dict)
if batch_idx > 50:
break
validation_summary = compute_dict_mean(validation_dicts)如果当前验证损失小于最小验证损失，则更新最小验证损失和最佳检查点信息
epoch_val_loss = validation_summary['loss']
if epoch_val_loss < min_val_loss:
min_val_loss = epoch_val_loss
best_ckpt_info = (step, min_val_loss, deepcopy(policy.serialize()))
for k in list(validation_summary.keys()):
validation_summary[f'val_{k}'] = validation_summary.pop(k)
if config['wandb']:
wandb.log(validation_summary, step=step)
print(f'Val loss: {epoch_val_loss:.5f}')
summary_string = ''
for k, v in validation_summary.items():
summary_string += f'{k}: {v.item():.3f} '
print(summary_string)然后，策略模型切换回训练模式，清零优化器梯度，对训练数据进行前向传播，计算损失并进行反向传播和优化器更新
# training
policy.train()
optimizer.zero_grad()
data = next(train_dataloader)
forward_dict = forward_pass(data, policy)
# backward
loss = forward_dict['loss']
loss.backward()
optimizer.step()
if config['wandb']:
wandb.log(forward_dict, step=step) # not great, make training 1-2% slower
如果当前步骤是保存间隔的倍数，则保存策略模型和优化器的状态
if step % save_every == 0:
ckpt_path = os.path.join(ckpt_dir, f'policy_step_{step}_seed_{seed}.ckpt')
torch.save(policy.serialize(), ckpt_path)
#save optimizer state
optimizer_ckpt_path = os.path.join(ckpt_dir, f'optimizer_step_{step}_seed_{seed}.ckpt')
torch.save(optimizer.state_dict(), optimizer_ckpt_path)
if step % 2000 == 0:
ckpt_path = os.path.join(ckpt_dir, f'policy_last.ckpt')
torch.save(policy.serialize(), ckpt_path)
optimizer_ckpt_path = os.path.join(ckpt_dir, f'optimizer_last.ckpt')
torch.save(optimizer.state_dict(), optimizer_ckpt_path)
ckpt_path = os.path.join(ckpt_dir, f'policy_last.ckpt')
torch.save(policy.serialize(), ckpt_path)
optimizer_ckpt_path = os.path.join(ckpt_dir, f'optimizer_last.ckpt')
torch.save(optimizer.state_dict(), optimizer_ckpt_path)
return best_ckpt_inforepeater 函数将数据加载器转换为一个无限循环的数据加载器。它使用repeat 函数不断重复数据加载器的内容，并在每个周期结束时打印当前周期数
3.1.2 def main_train(args)
main_train 是训练过程的入口点
- 它首先设置随机种子，并从命令行参数中提取各种配置参数
def main_train(args): set_seed(1) # command line parameters 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_steps = args['num_steps'] eval_every = args['eval_every'] validate_every = args['validate_every'] save_every = args['save_every'] resume_ckpt_path = args['resume_ckpt_path'] backbone = args['backbone'] same_backbones = args['same_backbones'] - 然后，根据任务名称和策略类等参数构建检查点目录，并检查目录是否已存在。如果目录已存在，则退出程序
ckpt_dir = f'{ckpt_dir}_{task_name}_{policy_class}_{backbone}_{same_backbones}' if os.path.isdir(ckpt_dir): print(f'ckpt_dir {ckpt_dir} already exists, exiting') return args['ckpt_dir'] = ckpt_dir - 接着，函数从任务配置中提取任务参数，并根据策略类创建策略配置
- 然后，函数创建训练和验证数据加载器，并保存数据集统计信息
- 最后，调用 train_bc 函数进行模型训练，并保存最佳检查点
在命令行参数解析部分，定义了各种命令行参数，包括
配置文件路径、评估标志、屏幕渲染标志、检查点目录、策略类、任务名称、批量大小、随机种子、训练步骤数、学习率、预训练模型路径、验证间隔、保存间隔、恢复检查点路径、跳过镜像数据标志、执行器网络目录、历史长度、未来长度、预测长度、KL权重、块大小、隐藏维度、前馈维度、编码器层数、解码器层数、注意力头数、图像位置嵌入标志、动作位置嵌入标志、特征损失权重、自注意力标志、骨千网络类型、相同骨千网络标志、使用掩码标志、图像宽度、图像高度、数据增强标志、归一化ResNet标志、灰度图像标志、随机颜色标志、随机数据标志、随机数据程度、WandB标志、模型类型和GPU ID
解析命令行参数后，设置CUDA设备并调用
main_train 函数开始训练
3.2 HIT/model_util.py
3.3 HIT/policy.py
这段代码定义了几个策略类 (HITPolicy、DiffusionPolicy 、ACTPolicy 和CNNMLPPolicy)，这
些类继承自 torch.nn.Module，用于不同的机器学习模型和训练策略
3.3.1 HITPolicy 类
HITPolicy 类是一个基于 HT模型的策略类。它的主要功能包括模型的初始化、前向传播、优化器配置、序列化和反序列化
- 初始化：这个_-init_—方法是一个类的构造函数，用于初始化类的实例。它接受一个参数 args_override，这是一个包含各种配置参数的字典
首先，调用 superO.--init_-（）来初始化父类。这是一个常见的做法，确保父类的初始化逻辑被执行class HITPolicy(nn.Module): def __init__(self, args_override):
接下来，将args_override 宇典中的model_type 键设置为"HIT"。这可能是为了确保模型类型的一致性super().__init__()
然后，调用 build_ACT_model-and_optimizer Cargs_override) 函数来构建模型和优化器。这个函数返回一个模型和一个优化器，并将它们分别赋值给实例的 self.model 和 self.optimizer 属性args_override['model_type'] = "HIT"
接下来，检查 args_override 宇典中是否包含feature_loss_weight 键。如果存在，则将其值赋给实例的 self. feature_loss_weight 属性；如果不存在，则将self. feature_Loss_weight 设置为0model, optimizer = build_ACT_model_and_optimizer(args_override) self.model = model self.optimizer = optimizer
最后，尝试从 args_override 宇典中获取state_idx 和action_idx 键的值，并将它们分别赋值给实例的 self.state_idx和 self.action_idx 属性。如果获取失败，则将这两个属性设置为 Noneself.feature_loss_weight = args_override['feature_loss_weight'] if 'feature_loss_weight' in args_override else 0.0try: self.state_idx = args_override['state_idx'] self.action_idx = args_override['action_idx'] except: self.state_idx = None self.action_idx = None - 前向传播：-call--方法用于前向传播。在训练时，它会计算损失，包括 L1 损失和特征损失。在推理时，它会返回预测的动作
具体而言，这个__call_.接受四个参数：qpos （位置数据)、image（图像数据)、actions（动作数据，可选）和is_pad（填充标志，可选）
首先，方法检查实例属性 state_idx和action_idx是否为 None。如果不是None，则根据这些索引对qpos和actions 进行切片操作，以提取特定的状态和动作数据def __call__(self, qpos, image, actions=None, is_pad=None):
接下来，方法使用 transforms. Normalize 对图像数据进行归一化处理。归一化的均值和标准差分别为[0.485， 0.456， 0.406] 和[0.229，0.224，0.225]if self.state_idx is not None: qpos = qpos[:, self.state_idx] if self.action_idx is not None: actions = actions[:, :, self.action_idx]
如果提供了 actions 参数，表示这是训练时间。方法对 actions 和 is_pad 进行切片操作，以匹配模型的查询数量normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) image = normalize(image)
然后，创建一个空的损失字典 loss_dict。调用模型的前向传播方法，获取预测的动作a_hat 和其他输出如hs_img_dictif actions is not None: # training time actions = actions[:, :self.model.num_queries] is_pad = is_pad[:, :self.model.num_queries]
计算所有动作的l1损失，并根据 is-pad 标志进行加权平均，得到最终的L1损失l1。将L1损失添加到损失字典中loss_dict = dict() a_hat, _, hs_img_dict = self.model(qpos, image)
如果模型启用了特征损失并且处于训练模式，方法还会计算特征损失，并将其加权后添加到总损失中all_l1 = F.l1_loss(actions, a_hat, reduction='none') l1 = (all_l1 * ~is_pad.unsqueeze(-1)).mean() loss_dict['l1'] = l1
最终的总损失存储在损失字典的 ，1oss”键中。方法返回损失字典
如果末提供 actions 参数，表示这是推理时间。方法调用模型的前向传播方法，获取预测的动作a_hat，并返回该预测结果if self.model.feature_loss and self.model.training: loss_dict['feature_loss'] = F.mse_loss(hs_img_dict['hs_img'], hs_img_dict['src_future']).mean() loss_dict['loss'] = loss_dict['l1'] + self.feature_loss_weight*loss_dict['feature_loss'] else: loss_dict['loss'] = loss_dict['l1'] return loss_dictelse: # inference time a_hat, _, _ = self.model(qpos, image) # no action, sample from prior return a_hat - 推理forward_inf
这个 forward_inf 方法是一个类的实例方法，用于在推理过程中处理输入数据并返回模型的预测结果。它接受两个参数：qpos(位置数据）和 image（图像数据）
首先，方法检查实例属性 state_idx 是否为None。如果不是None，则根据这些索引对qpos 进行切片操作，以提取特定的状态数据def forward_inf(self, qpos, image):
接下来，方法使用 transforms . Normalize对图像数据进行归一化处理。归一化的均值和标准差分别为 [0.485， 0.456， 0.4061和[0.229， 0.224， 0.225]if self.state_idx is not None: qpos = qpos[:, self.state_idx]
然后，方法调用模型的前向传播方法self.model(qpos, image)，荻取预测的动作a_hat 和其他输出。这里没有提供动作数据，因此模型将从先验分布中采样normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) image = normalize(image)
最后，方法返回预测的动作a_hata_hat, _, _ = self.model(qpos, image) # no action, sample from prior
这个方法的主要目的是在推理过程中处理输人数据，调用模型进行前向传播，并返回模return a_hat
型的预测结果 - 优化器配置：configure_optimizers 方法返回优化器
- 序列化和反序列化：serialize 方法返回模型的状态字典，deserialize 方法加载模型的状态字典
3.3.2 DiffusionPolicy 类：初始化、前向传播、序列化与反序列化
DiffusionPolicy 类是一个基于扩散模型的策略类。它的主要功能包括模型的初始化、前向传播、优化器配置、序列化和反序列化。
3.3.2.1 初始化
在初始化时，DiffusionPolicy 类会根据 args_override 参数设置一些属性，并构建模型的各个部分，包括骨千网络、池化层、线性层和噪声预测网络
def __init__(self, args_override):
super().__init__()
self.camera_names = args_override['camera_names']
self.observation_horizon = args_override['observation_horizon'] ### TODO TODO TODO DO THIS
self.action_horizon = args_override['action_horizon'] # apply chunk size
self.prediction_horizon = args_override['prediction_horizon'] # chunk size
self.num_inference_timesteps = args_override['num_inference_timesteps']
self.ema_power = args_override['ema_power']
self.lr = args_override['lr']
self.weight_decay = 0以及再初始化一些固定的参数，比如关键点数量、特征维度，以计算动作维度和观察维度
self.num_kp = 32
self.feature_dimension = 64
self.ac_dim = args_override['action_dim'] # 14 + 2
self.obs_dim = self.feature_dimension * len(self.camera_names) + 14 # camera features and proprio接下来，再构建模型组件
backbones = []
pools = []
linears = []
for _ in self.camera_names:
backbones.append(ResNet18Conv(**{'input_channel': 3, 'pretrained': False, 'input_coord_conv': False}))
pools.append(SpatialSoftmax(**{'input_shape': [512, 15, 20], 'num_kp': self.num_kp, 'temperature': 1.0, 'learnable_temperature': False, 'noise_std': 0.0}))
linears.append(torch.nn.Linear(int(np.prod([self.num_kp, 2])), self.feature_dimension))
backbones = nn.ModuleList(backbones)
pools = nn.ModuleList(pools)
linears = nn.ModuleList(linears)然后，替换BN为GN，以将ResNet18卷积网络中的批归一化层替换为组归一化层
backbones = replace_bn_with_gn(backbones) # TODO再初始化一个条件Unet1D网络，用于噪声预测
noise_pred_net = ConditionalUnet1D(
input_dim=self.ac_dim,
global_cond_dim=self.obs_dim*self.observation_horizon
)
继续，构建一个模型字典，即将所有模型组件存储在一个ModuleDict中，以便于管理和访问
nets = nn.ModuleDict({
'policy': nn.ModuleDict({
'backbones': backbones,
'pools': pools,
'linears': linears,
'noise_pred_net': noise_pred_net
})
})然后，启用EMA
nets = nets.float().cuda()
ENABLE_EMA = True
if ENABLE_EMA:
ema = EMAModel(model=nets, power=self.ema_power)
else:
ema = None
self.nets = nets
self.ema = ema再设置一个噪声调度器
# setup noise scheduler
self.noise_scheduler = DDIMScheduler(
num_train_timesteps=50,
beta_schedule='squaredcos_cap_v2',
clip_sample=True,
set_alpha_to_one=True,
steps_offset=0,
prediction_type='epsilon'
)最后，打印模型参数数量
n_parameters = sum(p.numel() for p in self.parameters())
print("number of parameters: %.2fM" % (n_parameters/1e6,))3.3.2.2 前向传播
_call_-方法用于前向传播。在训练时，它会计算L2损失。在推理时，它会通过护散过程生成动作
具体而言，-cal1_._方法，用于在训练和推理过程中处理输入数据并返回模型的输出或损失字典。它接受四个参数：qpos(位置数据)、image（图像数据)、actions（动作数据，可选）和is_pad(填充标志，可选)
def __call__(self, qpos, image, actions=None, is_pad=None):- 训练时间
如果提供了 actions 参数，表示这是训练时问
方法执行以下五个步骤：提取特征、添加噪声、预测噪声残差、计算损失、更新EMAB = qpos.shape[0] if actions is not None: # training time nets = self.nets all_features = []
1. 提取特征：
遍历所有相机，提取每个相机图像的特征，使用ResNet18卷积网络提取特征
使用空问软最大池化层进行池化for cam_id in range(len(self.camera_names)): cam_image = image[:, cam_id] cam_features = nets['policy']['backbones'][cam_id](cam_image)
使用线性层将池化后的特征转换为输出特征pool_features = nets['policy']['pools'][cam_id](cam_features) pool_features = torch.flatten(pool_features, start_dim=1)
将所有特征和位置数据拼接成一个观察条件向量out_features = nets['policy']['linears'][cam_id](pool_features) all_features.append(out_features)
2. 添加噪声：# concatenate all features obs_cond = torch.cat(all_features + [qpos], dim=1)
生成与动作数据形状相同的高斯噪声
为每个数据点采样一个扩散迭代步数# sample noise to add to actions noise = torch.randn(actions.shape, device=obs_cond.device)
根据每个扩散迭代步数的噪声幅度，将噪声添加到干净的动作数据中，生成噪声动作# sample a diffusion iteration for each data point timesteps = torch.randint( 0, self.noise_scheduler.config.num_train_timesteps, (B,), device=obs_cond.device ).long()
3.预测噪声残差：# 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)
使用噪声预测网络预测噪声残差
可能对于这个地方，有同学有疑问，即为何叫做预测噪声残差，而非叫做预测噪声呢？其实如此文《图像生成发展起源：从VAE、VQ-VAE、扩散模型DDPM、DETR到ViT、Swin transformer》此节「2.1.1 从扩散模型概念的提出到DDPM」中所说# predict the noise residual noise_pred = nets['policy']['noise_pred_net'](noisy_actions, timesteps, global_cond=obs_cond)
“这种操作就有点类似ResNet的残差结构。每次新增一些层，模型不是直接从去预测
，而是让新增的层去预测(
)。这样新增层不用全部重新学习，而是学习原来已经学习到的
4. 计算损失：
计算预测噪声和真实噪声之问的L2损失
根据填充标志对损失进行加杈平均# L2 loss all_l2 = F.mse_loss(noise_pred, noise, reduction='none')
将损失存储在损失字典中# mask out padding loss = (all_l2 * ~is_pad.unsqueeze(-1)).mean()
5. 更新EMA：loss_dict = {} loss_dict['l2_loss'] = loss loss_dict['loss'] = loss
如果启用了EMA并且模型处于训练模式，更新EMA模型if self.training and self.ema is not None: self.ema.step(nets) return loss_dict - 推理时间
如果未提供 actions 参数，表示这是推理时间
方法执行以下五个步骤：提取特征、初始化动作、初始化调度器、扩散过程、返回动作else: # inference time To = self.observation_horizon Ta = self.action_horizon Tp = self.prediction_horizon action_dim = self.ac_dim nets = self.nets if self.ema is not None: nets = self.ema.averaged_model all_features = []
1.提取特征：
遍历所有相机，提取每个相机图像的特征，使用ResNet18卷积网络提取特征
使用空间软最大池化层进行池化for cam_id in range(len(self.camera_names)): cam_image = image[:, cam_id] cam_features = nets['policy']['backbones'][cam_id](cam_image)
使用线性层将池化后的特征转换为输出特征pool_features = nets['policy']['pools'][cam_id](cam_features) pool_features = torch.flatten(pool_features, start_dim=1)
将所有特征和位置数据拼接成一个观察条件向量out_features = nets['policy']['linears'][cam_id](pool_features) all_features.append(out_features)
2. 初始化动作：# concatenate all features obs_cond = torch.cat(all_features + [qpos], dim=1)
从高斯噪声中初始化动作
3. 初始化调度器：# initialize action from Guassian noise noisy_action = torch.randn( (B, Tp, action_dim), device=obs_cond.device) naction = noisy_action
设置扩散调度器的时问步数
4. 扩散过程：# init scheduler self.noise_scheduler.set_timesteps(self.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 )
5. 返回去噪后的动作# inverse diffusion step (remove noise) naction = self.noise_scheduler.step( model_output=noise_pred, timestep=k, sample=naction ).prev_samplereturn naction
3.3.2.3 优化器配置、序列化、反序列化
优化器配置：configure_optimizers 方法返回优化器
def configure_optimizers(self):
optimizer = torch.optim.AdamW(self.nets.parameters(), lr=self.lr, weight_decay=self.weight_decay)
return optimizerserialize 方法返回模型和EMA 模型的状态字典
def serialize(self):
return {
"nets": self.nets.state_dict(),
"ema": self.ema.averaged_model.state_dict() if self.ema is not None else None,
}deserialize 方法加载模型和 EMA 模型的状态字典
def deserialize(self, model_dict):
status = self.nets.load_state_dict(model_dict["nets"])
print('Loaded model')
if model_dict.get("ema", None) is not None:
print('Loaded EMA')
status_ema = self.ema.averaged_model.load_state_dict(model_dict["ema"])
status = [status, status_ema]
return status3.3.3 ACTPolicy 类
ACTPolicy 类是一个基于 ACT 模型的策略类。它的主要功能包括模型的初始化、前向传播、优化器配置、序列化和反序列化
- 初始化
首先调用父类的初始化方法
然后，设置模型类型为 “ACT"，并通过调用 build_ACT_mode l-and_optimizer 函数构建模型和优化器class ACTPolicy(nn.Module): def __init__(self, args_override): super().__init__()
接着，将模型和优化器分别赋值给实例变量 self.model 和 self.optimizerargs_override['model_type'] = "ACT" model, optimizer = build_ACT_model_and_optimizer(args_override)
此外，还从args_override 中获取KL 散度权重 kL_weight 和向量量化标志 v9self.model = model # CVAE decoder self.optimizer = optimizer
最后，尝试从 args_override 中获取状态和动作的索引，如果获取失败，则将它们设置为 Noneself.kl_weight = args_override['kl_weight'] self.vq = args_override['vq'] print(f'KL Weight {self.kl_weight}')try: self.state_idx = args_override['state_idx'] self.action_idx = args_override['action_idx'] except: self.state_idx = None self.action_idx = None - 前向传播：_-cal1-- 方法用于前向传播。在训练时，它会计算 L1损失和 KL 散度损失。在推理时，它会返回预测的动作
- 优化器配置：configure_optimizers 方法返回优化器
- 序列化和反序列化：serialize 方法返回模型的状态字典，deserialize 方法加载模型的状态字典
3.3.4 CNNMLPPolicy 类
CNNMLPPolicy 类是一个基于 CNN 和MLP模型的策略类。它的主要功能包括模型的初始化、前向传播、优化器配置、序列化和反序列化
- 初始化：在初始化时，CNNMLPPolicy 类会调用 build_CNNMLP_model-and_optimizer 函数来构建模型和优化器
- 前向传播：--call---方法用于前向传播。在训练时，它会计算 MSE 损失。在推理时，它会返回预测的动作
- 优化器配置：configure_optimizers 方法返回优化器
辅助函数kl_divergence：计算KL 散度，用于衡量两个分布之间的差异
3.4 HIT/utils.py
// 待更
扫一扫
3059
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
