人形运控部署框架汇总：rsl_rl 与 unitree_rl_gym 解析

import torch
import torch.nn as nn
import torch.optim as optim
from rsl_rl.modules import ActorCriticTransformer
from rsl_rl.storage import RolloutStorage

class PPO:
    actor_critic: ActorCriticTransformer

    def __init__(self,
                 actor_critic,
                 num_learning_epochs=1,
                 num_mini_batches=1,
                 clip_param=0.2,
                 gamma=0.998,
                 lam=0.95,
                 value_loss_coef=1.0,
                 entropy_coef=0.0,
                 learning_rate=1e-3,
                 max_grad_norm=1.0,
                 use_clipped_value_loss=True,
                 schedule="fixed",
                 desired_kl=0.01,
                 device='cpu'):

它初始化 PPO 的核心组件：

• device: 传入的设备参数。
• desired_kl: 目标 KL 散度，用于自适应学习率调整。
• schedule: 学习率调度策略。
• optimizer: 使用 Adam 优化器来更新 actor_critic 网络的参数。
• storage: 用于存储经验轨迹（transitions）的 RolloutStorage 对象。

1.1.2 存储初始化 init_storage

这个方法根据环境的数量、每个环境收集的转换（transitions）数量、观察空间形状和动作空间形状来创建 RolloutStorage 实例。

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)

1.1.3 模式切换 test_mode/train_mode

这些方法用于切换 actor_critic 网络到评估（测试）模式或训练模式。

def test_mode(self):
    self.actor_critic.test()

def train_mode(self):
    self.actor_critic.train()

1.1.4 动作选择 act

这是智能体与环境交互的核心。给定当前观察 obs 和 critic_obs，它执行以下操作：

def act(self, obs, critic_obs):
    # 获取动作均值和标准差
    self.transition.action_mean = self.actor_critic.action_mean.detach()
    self.transition.action_sigma = self.actor_critic.action_std.detach()
    # 计算动作
    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()
    # 如果是循环网络，获取隐藏状态
    if self.actor_critic.is_recurrent:
        self.transition.hidden_states = self.actor_critic.get_hidden_states()
    return self.transition.actions

1.1.5 处理环境步骤 process_env_step

在环境执行动作后调用此方法。如果 episode 因为达到时间限制而不是因为失败状态而结束，它会将最后一步的估计价值加到奖励中，这称为'超时引导 (Bootstrapping on time outs)'。

def process_env_step(self, rewards, dones, infos):
    self.transition.rewards = rewards.clone()
    self.transition.dones = dones
    # 重置隐藏状态
    self.actor_critic.reset(dones)
    # 超时引导
    if 'time_outs' in infos:
        self.transition.rewards += self.gamma * torch.squeeze(self.transition.values * infos['time_outs'].unsqueeze(1).to(self.device), 1)
    # 记录过渡
    self.storage.add_transitions(self.transition)
    self.transition.clear()

1.1.6 计算回报 compute_returns

在收集了足够多的 transitions 后调用。通常使用广义优势估计 (Generalized Advantage Estimation, GAE) 来计算每个时间步的回报 (returns) 和优势 (advantages)。

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)

1.1.7 update：策略和价值网络参数更新

这是执行策略和价值网络参数更新的核心循环。

def update(self):
    mean_value_loss = 0
    mean_surrogate_loss = 0
    
    # 检查是否为循环网络并选择生成器
    if self.actor_critic.is_recurrent:
        generator = self.storage.reccurent_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:
        # 计算重要性采样比率
        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()
        
        # 价值函数损失计算
        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()
        
        # 梯度更新
        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()

    # 自适应学习率调整 (KL 散度)
    if self.desired_kl is not None and self.schedule == "adaptive":
        with torch.inference_mode():
            kl = torch.sum(torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) / (2.0 * torch.square(sigma_batch)) - 0.5, axis=-1)
            kl_mean = torch.mean(kl)
            if kl_mean > self.desired_kl * 2.0:
                self.learning_rate = max(1e-5, self.learning_rate / 1.5)
            elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
                self.learning_rate = min(1e-2, self.learning_rate * 1.5)

    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.2 rsl_rl/env

1.3 rsl_rl/modules

1.3.0 modules/actor_critic_transformer.py

这段代码定义了一个基于 BERT 风格的 Transformer 模型，用于强化学习中的 Actor-Critic 方法。它由三个主要的类组成：Transformer_Block, Transformer、ActorCriticTransformer。

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_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)
        x = x[:, -1, :]
        x = self.output_layer(x)
        return x

Transformer_Block 类是构建 Transformer 模型的基础块，包含了多头注意力机制 (Multihead Attention) 和前馈神经网络 (FeedForward Neural Network)。

class Transformer_Block(nn.Module):
    def __init__(self, latent_dim, num_head, dropout_rate) -> None:
        super().__init__()
        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

1.3.1 modules/actor_critic_depth_cnn.py

1.3.2 modules/actor_critic_history.py

1.3.3 modules/actor_critic_recurrent.py

1.3.4 modules/actor_critic.py

这段代码定义了一个名为 ActorCritic 的类，它是一个用于强化学习的演员 - 评论家模型的实现。包含两个主要部分：一个用于决策的策略网络 (Actor) 和一个用于评估动作价值的价值网络 (Critic)。

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):
        super(ActorCritic, self).__init__()
        activation = get_activation(activation)
        mlp_input_dim_a = num_actor_obs
        mlp_input_dim_c = num_critic_obs
        
        # Policy actor
        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)
        
        # Value function critic
        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)

1.3.5 modules/depth_backbone.py

1.3.6 modules/normalizer.py

1.4 rsl_rl/runners