TorchRL中为PettingZoo环境定义RewardSum实例报错求助

问题原因 #

你使用的是AEC模式（单智能体轮流行动）的多智能体环境，每个智能体的奖励以(agent_id, 'reward')的元组形式作为键存储在TensorDict中。而默认的RewardSum转换期望处理单键的'reward'，且其默认的reset_keys与环境的reset_keys不匹配，导致键长度不一致的报错。

解决方案 #

我们需要先将所有智能体的单步奖励求和得到总奖励，再用RewardSum累计总奖励到["next", "episode_reward"]键中。具体步骤如下：

1. 定义总奖励转换类 #

创建一个自定义Transform，聚合所有智能体的单步奖励：

from torchrl.envs.transforms import Transform
import torch

class TotalReward(Transform):
    def __init__(self, agents):
        # 输入键为所有智能体的奖励键
        super().__init__(in_keys=[(agent, 'reward') for agent in agents], out_keys=['reward'])
    
    def _apply_transform(self, tensordict):
        # 求和所有智能体的单步奖励
        total_reward = sum(tensordict[key] for key in self.in_keys)
        tensordict['reward'] = total_reward
        return tensordict

2. 组合转换并初始化环境 #

先通过TotalReward聚合奖励，再用RewardSum累计总奖励到目标键：

from torchrl.envs import TransformedEnv
from torchrl.envs.transforms import RewardSum
from torchrl.envs.libs.pettingzoo import PettingZooEnv

# 初始化原始环境
env = PettingZooEnv(
    task="mpe/simple_spread_v3",
    parallel=False,
    use_mask=True,
    group_map=None
)

# 获取所有智能体ID
agents = env.agents

# 应用转换
env = TransformedEnv(
    env,
    TotalReward(agents),
    RewardSum(out_key=("next", "episode_reward"))
)

# 验证环境规格
check_env_specs(env)

3. 验证结果 #

运行后，你可以通过tensordict["next", "episode_reward"]获取所有智能体的累计奖励总和。

内容的提问来源于stack exchange，提问作者sakinah