🤖 roboto_origin_03 Wiki
首页 / RSL-RL / 循环与注意力策略变体

标准的深度强化学习策略通常假设智能体能够观测到环境的完整状态(MDP 设定),但在机器人控制、导航和 partially observable 场景中,单帧观测往往不足以支持最优决策。RSL-RL 在基础 Actor-Critic 基础架构设计 之上,提供了两种重要的策略扩展:**循环策略(Recurrent Policy)**通过时序记忆建模历史信息,**注意力策略(Attention Policy)**通过空间注意力机制融合本体感觉与结构化感知观测。本文将系统剖析这两种变体的架构原理、隐藏状态管理、训练流程差异及适用场景。

循环策略架构设计

循环策略的核心动机源于 POMDP(Partially Observable Markov Decision Process)问题的现实需求:当机器人的当前观测仅包含关节位置、IMU 读数等低维状态,而缺少速度、摩擦或外部环境信息时,策略必须依赖历史序列进行推断。ActorCriticRecurrent 通过在 Actor 和 Critic 的输入端分别插入独立的 Memory 模块,将时序动态注入决策过程。

Memory 模块:GRU/LSTM 封装

循环策略不直接使用 PyTorch 原生 RNN,而是通过 Memory 类进行统一封装,以兼容推理与训练两种截然不同的执行模式。该模块支持 grulstm 两种细胞类型,并引入 HiddenState 类型别名来抽象 GRU 的单张量状态与 LSTM 的元组状态(hidden + cell)。在推理模式下,Memory 维护一个实例级别的 self.hidden_state,每步接收当前观测的 unsqueeze(0) 扩展序列维度,输出经过 RNN 后的特征,同时更新内部状态供下一步使用;在训练模式下,模块接收完整的轨迹张量、掩码 masks 和外部传入的 hidden_state,前向传播后通过 unpad_trajectories 去除填充,仅保留有效时间步的输出。这种设计使得同一模块可以无缝服务于环境交互(单步推理)和策略更新(批量训练)两种场景。

Sources: memory.py

隐藏状态生命周期与重置机制

隐藏状态的管理是循环策略工程实现中最精密的环节。在环境交互阶段,PPO 算法每次调用 act() 前通过 get_hidden_states() 捕获当前隐藏状态,并将其存入 RolloutStorage.Transition;当环境返回 dones 信号时,process_env_step() 调用 policy.reset(dones),将已完成 episode 的对应环境隐藏状态清零,避免历史信息跨 episode 泄漏。Memory.reset() 支持两种语义:当传入 dones 时,仅重置完成环境的隐藏状态;当 donesNone 时,全局重置所有状态。此外,detach_hidden_state() 方法用于在 BPTT 过程中截断梯度流,防止长期反向传播导致的数值不稳定。

Sources: actor_critic_recurrent.py, memory.py, ppo.py

轨迹切分与填充:Recurrent Mini-batch 生成

循环策略训练的最大挑战在于:向量化环境中不同环境的 episode 结束时刻不同,导致 rollout buffer 中的数据并非整齐的轨迹矩阵。RolloutStoragerecurrent_mini_batch_generator 使用 split_and_pad_trajectories 工具函数,依据 dones 信号将 [time, num_envs, ...] 的 rollout 数据切分为独立轨迹,填充零值至最长轨迹长度,生成 [max_traj_len, num_trajectories, ...] 的结构。与此同时,隐藏状态按照 trajectory 的起始点提取:通过 last_was_done 标记每个轨迹的第一个时间步,从保存的隐藏状态中提取对应的初始状态,重塑为 [num_layers, batch, hidden_dim] 传入 RNN。掩码 trajectory_masks 则用于在损失计算中屏蔽填充位置,确保梯度只作用于有效数据。

Sources: rollout_storage.py, utils.py

flowchart TD
    A[Rollout Buffer<br/>time x num_envs] --> B[split_and_pad_trajectories<br/>按 dones 切分]
    B --> C[Padded Trajectories<br/>max_len x num_traj]
    B --> D[Trajectory Masks<br/>有效时间步标记]
    E[Saved Hidden States<br/>time x num_layers x num_envs x hidden_dim] --> F[按轨迹起始提取<br/>num_layers x num_traj x hidden_dim]
    C --> G[RNN Forward<br/>batch_first=False]
    F --> G
    D --> H[unpad_trajectories<br/>去除填充]
    G --> H
    H --> I[MLP Actor/Critic]

ActorCriticRecurrent 接口契约

ActorCriticRecurrent 在 API 层面通过 is_recurrent = True 标识自身属性,act()evaluate() 方法均扩展了 maskshidden_state 参数,与基础 ActorCritic 保持前向兼容。Actor 路径上,观测先经过归一化,再通过 memory_a 提取时序特征,最终输入 Actor MLP 生成动作分布;Critic 路径对称地使用独立的 memory_c 模块。两个路径的 RNN 相互隔离,确保值函数估计与策略决策拥有各自的时序表征空间。值得注意的是,循环策略当前与对称性增强(Symmetry Augmentation)不兼容,PPO 在初始化时会显式检查并抛出错误。

Sources: actor_critic_recurrent.py, ppo.py

注意力策略架构设计

注意力策略面向另一类感知增强场景:机器人除了本体感觉(proprioception)外,还接收结构化空间观测,如高度图、占据栅格或激光扫描。直接展平此类空间数据输入 MLP 会破坏局部空间结构,而纯 CNN 又难以动态聚焦与当前决策相关的空间区域。ActorCriticAttnEnc 采用本体感觉为 Query、空间特征为 Key/Value 的多头注意力范式,实现自适应的空间感知融合。

AttentionEncoder:CNN + MHA 空间编码

AttentionEncoder 是注意力策略的核心组件,其输入分为两路:本体感觉向量 proprioception 和地图扫描 map_scans。地图扫描首先通过 CNN 提取空间特征,输出通道数为 embedding_dim - 3;同时模块注册预计算的位置编码缓冲区,包含网格的物理坐标 xy。在特征融合阶段,位置编码、原始高度图(1 通道)和 CNN 特征在通道维度拼接为 embedding_dim 维的点级特征,重塑为 [B, W*L, embedding_dim] 的序列表示。本体感觉则经 MLP 投影为 [B, 1, embedding_dim] 的 Query 向量。随后,Query 与 Key/Value 分别经过层归一化,输入 nn.MultiheadAttention,输出空间注意力编码和注意力权重图。这种设计的本质是让策略"看向"地图中与当前本体状态最相关的空间位置。

Sources: attn_encoder.py, cnn.py

flowchart LR
    A[Map Scans<br/>B x W x L x 1] --> B[CNN<br/>提取局部特征]
    C[Position Encoding<br/>2 x W x L] --> D[拼接: pos + z + cnn_feat]
    B --> D
    D --> E[Pointwise Features<br/>B x W*L x embedding_dim]
    F[Proprioception<br/>B x d_obs] --> G[MLP Projection<br/>B x 1 x embedding_dim]
    E --> H[LayerNorm KV]
    G --> I[LayerNorm Q]
    H --> J[MultiheadAttention]
    I --> J
    J --> K[Map Encoding<br/>B x embedding_dim]
    J --> L[Attention Weights<br/>B x W x L]

时序观测与历史长度处理

ActorCriticAttnEnc 支持时序观测输入,通过 actor_history_lengthcritic_history_length 定义历史帧数。模块在初始化时计算 single_actor_obs_dim = num_actor_obs // actor_history_length,在前向传播中提取最后一帧的观测作为注意力编码器的本体感觉输入。这意味着注意力机制仅对当前帧的空间感知进行编码,而历史信息通过观测向量的拼接在 MLP 阶段保留。这种分层设计将时序聚合委托给观测构造层,将空间聚合委托给注意力层,各司其职。

Sources: actor_critic_attn_enc.py

可选扩展:观测编码器与 Critic 估计

除核心注意力编码外,ActorCriticAttnEnc 提供两项可选的辅助功能。当启用 enable_obs_encoder 时,模块为 Actor 和 Critic 各自附加一个 MLP 观测编码器,将完整历史观测压缩为 latent_dim 维的隐变量,与单帧观测拼接后输入注意力编码器。当启用 enable_critic_estimation 时,模块引入一个估计器网络,尝试从 Actor 侧观测预测 Critic 侧的特定状态分量(由 estimation_slice 指定),并通过 get_aux_loss() 返回预测值与真实值之间的 MSE 辅助损失。该辅助损失在 PPO 更新阶段以 aux_loss_coef 为系数加入总损失,可作为表征学习的正则化手段。

Sources: actor_critic_attn_enc.py, actor_critic_attn_enc.py, ppo.py

策略变体对比与训练差异

维度 ActorCritic(基础) ActorCriticRecurrent ActorCriticAttnEnc
is_recurrent False True False
核心网络组件 MLP Memory(GRU/LSTM) + MLP AttentionEncoder(CNN+MHA) + MLP
时序建模方式 RNN 隐藏状态 观测拼接历史帧
空间感知能力 多头注意力聚焦
训练 batch 生成 全局随机打乱 轨迹切分+填充 全局随机打乱
隐藏状态管理 每步保存/按轨迹传入
辅助损失支持 Critic Estimation MSE
对称性增强兼容
关键新增参数 rnn_type, rnn_hidden_dim, rnn_num_layers embedding_dim, head_num, map_size, history_length, enable_critic_estimation

Sources: actor_critic.py, actor_critic_recurrent.py, actor_critic_attn_enc.py, ppo.py

在训练循环中,两种变体对 PPO 的侵入性极小。PPO 通过 policy.is_recurrent 自动分支选择 mini-batch 生成器:对于循环策略调用 recurrent_mini_batch_generator,对于非循环策略(含基础 MLP 和注意力)调用 mini_batch_generator。在策略更新阶段,act()evaluate() 统一接收 maskshidden_states 参数;对于非循环策略,这两个参数为 None,不会干扰前向逻辑。这种基于布尔标志的分支策略使得 PPO 算法层无需感知具体网络结构,实现了算法与网络架构的解耦。

Sources: ppo.py, ppo.py

参数配置与使用场景

ActorCriticRecurrent 核心参数

参数名 类型 默认值 说明
rnn_type str "lstm" 循环细胞类型,可选 "gru""lstm"
rnn_hidden_dim int 256 RNN 隐藏层维度,替代已弃用的 rnn_hidden_size
rnn_num_layers int 1 RNN 堆叠层数

推荐场景:观测具有显著的时间依赖性而单帧信息不充分的任务,例如需要估计速度但观测中不包含速度读数、需要记忆过去接触状态进行步态规划,或环境存在时间上的模式需要识别。

Sources: actor_critic_recurrent.py

ActorCriticAttnEnc 核心参数

参数名 类型 默认值 说明
embedding_dim int 64 注意力嵌入维度,必须大于 3 且可被 head_num 整除
head_num int 8 多头注意力头数
map_size tuple[int] (17, 11) 输入地图扫描的空间尺寸 (L, W)
map_resolution float 0.1 地图网格物理分辨率,用于位置编码
actor_history_length int 5 Actor 侧观测历史帧数
critic_history_length int 5 Critic 侧观测历史帧数
enable_critic_estimation bool False 是否启用 Critic 状态分量估计辅助任务
estimation_slice list[int] [78,79,80] 待估计的状态分量索引
enable_obs_encoder bool False 是否启用额外观测编码器生成隐变量
latent_dim int 0 观测编码器输出维度

推荐场景:机器人携带结构化空间传感器(如深度相机、激光雷达、高度图)的导航与避障任务,需要模型自动学习"看向"地图中的关键区域(如障碍物、目标点、地形变化处),而非手动设计特征提取。

Sources: actor_critic_attn_enc.py

架构设计模式总结

从软件工程视角审视,两种变体均遵循组合优于继承的设计哲学。ActorCriticRecurrent 不修改 MLP 的实现,而是在其前置入 Memory 模块;ActorCriticAttnEnc 同样保持 MLP 不变,将空间编码委托给 AttentionEncoder。这种模块化使得基础组件可以独立演进、复用和测试——Memory 可被蒸馏框架的 StudentTeacherRecurrent 直接复用,AttentionEncoder 的核心 CNN 和 MHA 组合亦可迁移至其他架构。此外,两种变体均完整保留了基础架构的观测分组(obs_groups)、归一化(EmpiricalNormalization)和动作分布(state_dependent_stdnoise_std_type)机制,确保用户在不同策略变体间迁移时,大部分超参数无需调整。

Sources: modules/init.py

如需深入理解这些策略变体在完整训练流程中的集成方式,可继续阅读 PPO 算法实现与训练流程 了解 actprocess_env_stepupdate 的协作细节;若关注数据如何在循环策略中经过切分与填充进入网络,请参阅 经验采样与小批量生成。对于需要将空间感知能力迁移至 CNN 方案的场景,可参考 CNN 观测编码与特征提取