From b4f3e2ec2eabdb028843534f28b459ff9af9466c Mon Sep 17 00:00:00 2001 From: nshoman Date: Tue, 5 May 2026 08:12:40 -0400 Subject: [PATCH] fix pz action mask check to see if the agent is present in the observation and info dicts before accessing the keys --- torchrl/envs/libs/pettingzoo.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/torchrl/envs/libs/pettingzoo.py b/torchrl/envs/libs/pettingzoo.py index 29f4febd0b5..fffb220a010 100644 --- a/torchrl/envs/libs/pettingzoo.py +++ b/torchrl/envs/libs/pettingzoo.py @@ -804,6 +804,8 @@ def _update_action_mask(self, td, observation_dict, info_dict): group_mask = td.get((group, "action_mask")) group_mask += True for index, agent in enumerate(agents): + if agent not in observation_dict or agent not in info_dict: + continue agent_obs = observation_dict[agent] agent_info = info_dict[agent] if isinstance(agent_obs, dict) and "action_mask" in agent_obs: