RLLIB的GTrXL模型是否支持字典观测？如何接入自定义Dict空间？

问题背景 #

已定义包含多个Box子空间的gym Dict观测空间（如下代码），需要将这类复杂字典输入接入RLLIB的AttentionNet（即GTrXL模型），已参考Complex input nets示例，但不清楚二者的合理结合方式。

观测空间定义代码：

observation_space = None

values = {
    "left_arm_joint_position": gym.spaces.Box(low=self.robot.joint_limits_low, high=self.robots.joint_limits_high, shape=(7,)),
    "right_arm_joint_position": gym.spaces.Box(low=self.robot.joint_limits_low, high=self.robots.joint_limits_high, shape=(7,)),
    "left_arm_joint_velocity": gym.spaces.Box(low=-self.robot.joint_vel_limits, high=self.robot.joint_vel_limits, shape=(7,)),
    "right_arm_joint_velocity": gym.spaces.Box(low=-self.robot.joint_vel_limits, high=self.robot.joint_vel_limits, shape=(7,)),
    "left_arm_joint_torque": gym.spaces.Box(low=-self.robot.joint_torque_limits, high=self.robot.joint_torque_limits, shape=(7,)),
    "right_arm_joint_torque": gym.spaces.Box(low=-self.robot.joint_torque_limits, high=self.robot.joint_torque_limits, shape=(7,)),
    "depth_image": gym.spaces.Box(low=0, high=255, shape=(self._render_height, self._render_width, 1)),
}

observation_space = gym.spaces.Dict(values)

实现步骤与代码示例 #

1. 自定义输入预处理网络 #

针对字典中不同类型的观测（低维关节数据、高维深度图像）分别做特征提取，再将所有特征融合为GTrXL可接收的一维序列输入：

import torch
import torch.nn as nn
import gym

class DictInputProcessor(nn.Module):
    def __init__(self, obs_space):
        super().__init__()
        # 处理关节类低维数据：用全连接层映射到统一维度
        joint_feature_dim = 64
        self.joint_fc = nn.Sequential(
            nn.Linear(7, 32),
            nn.ReLU(),
            nn.Linear(32, joint_feature_dim),
            nn.ReLU()
        )
        # 处理深度图像：用卷积层提取特征
        img_h, img_w = obs_space["depth_image"].shape[:2]
        # 计算卷积后特征图尺寸
        conv1_h = (img_h - 3) // 2 + 1
        conv1_w = (img_w - 3) // 2 + 1
        conv2_h = (conv1_h - 3) // 2 + 1
        conv2_w = (conv1_w - 3) // 2 + 1
        self.image_conv = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, stride=2),
            nn.ReLU(),
            nn.Conv2d(16, 32, kernel_size=3, stride=2),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(32 * conv2_h * conv2_w, 128),
            nn.ReLU()
        )
        # 总特征维度：6个关节特征+1个图像特征
        self.total_feature_dim = 6 * joint_feature_dim + 128

    def forward(self, obs_dict):
        # 处理每个关节相关观测
        joint_features = []
        joint_keys = [
            "left_arm_joint_position", "right_arm_joint_position",
            "left_arm_joint_velocity", "right_arm_joint_velocity",
            "left_arm_joint_torque", "right_arm_joint_torque"
        ]
        for key in joint_keys:
            feat = self.joint_fc(obs_dict[key].float())
            joint_features.append(feat)
        # 处理深度图像：调整维度为(batch, channels, H, W)
        image_feat = self.image_conv(obs_dict["depth_image"].permute(0, 3, 1, 2).float())
        # 拼接所有特征
        all_features = torch.cat(joint_features + [image_feat], dim=1)
        return all_features

2. 扩展GTrXL模型（AttentionNet） #

继承RLLIB的AttentionNet，将自定义输入预处理整合到模型前向流程中：

from ray.rllib.models.torch.attention_net import AttentionNet

class CustomGTrXL(AttentionNet):
    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
        # 先初始化输入处理器，获取预处理后的特征空间
        self.input_processor = DictInputProcessor(obs_space)
        processed_obs_space = gym.spaces.Box(
            low=-float("inf"), high=float("inf"), 
            shape=(self.input_processor.total_feature_dim,)
        )
        # 用预处理后的空间初始化父类GTrXL模型
        super().__init__(processed_obs_space, action_space, num_outputs, model_config, name)

    def forward(self, input_dict, state, seq_lens):
        # 先处理字典观测，得到融合后的特征
        processed_obs = self.input_processor(input_dict["obs"])
        # 将处理后的特征传入父类的forward方法
        return super().forward(
            {"obs": processed_obs}, state, seq_lens
        )

3. 配置RLLIB训练参数 #

在训练配置中指定自定义模型，并设置GTrXL相关参数：

from ray.rllib.agents.ppo import PPOTrainer

config = {
    "env": "YourCustomEnv",  # 替换为你的自定义环境类名
    "model": {
        "custom_model": CustomGTrXL,
        # GTrXL核心配置
        "attention_num_heads": 4,
        "attention_layers": 3,
        "attention_dim": 256,
        "max_seq_len": 64,
        "memory_inference": 10,
        "memory_training": 10,
    },
    "framework": "torch",
    # 基础训练参数
    "lr": 5e-4,
    "train_batch_size": 4096,
    "sgd_minibatch_size": 256,
}

# 启动训练循环
trainer = PPOTrainer(config=config)
for iter in range(100):
    result = trainer.train()
    print(f"Iteration {iter}: 平均奖励={result['episode_reward_mean']}")

关键注意点 #

• 深度图像的卷积输出维度需要根据你的图像尺寸（_render_height和_render_width）计算，确保全连接层的输入维度匹配。
• 如果需要保留时序信息的局部关联性，可以将关节特征和图像特征拆分为独立token传入GTrXL（比如每个关节特征作为一个token，图像特征作为一个token，形成长度为7的序列），而非直接拼接成一维向量。
• 可根据任务复杂度调整GTrXL的注意力头数、层数、特征维度等参数，优化模型性能。

内容的提问来源于stack exchange，提问作者JS-FWR