Unity-Technologies/ml-agents · error · KeyError

agent_id {agent_id} is not present in the DecisionSteps

Error message

agent_id {agent_id} is not present in the DecisionSteps

What it means

DecisionSteps.__getitem__ raises KeyError when the requested agent_id is not among the agents that produced a decision this step. agent_id_to_index only contains ids present in the current batch, so any id not stepping/deciding now is invalid. This is a plain Python KeyError (the message is the formatted string).

Source

Thrown at ml-agents-envs/mlagents_envs/base_env.py:118

        this DecisionSteps.
        """
        if self._agent_id_to_index is None:
            self._agent_id_to_index = {}
            for a_idx, a_id in enumerate(self.agent_id):
                self._agent_id_to_index[a_id] = a_idx
        return self._agent_id_to_index

    def __len__(self) -> int:
        return len(self.agent_id)

    def __getitem__(self, agent_id: AgentId) -> DecisionStep:
        """
        returns the DecisionStep for a specific agent.
        :param agent_id: The id of the agent
        :returns: The DecisionStep
        """
        if agent_id not in self.agent_id_to_index:
            raise KeyError(f"agent_id {agent_id} is not present in the DecisionSteps")
        agent_index = self._agent_id_to_index[agent_id]  # type: ignore
        agent_obs = []
        for batched_obs in self.obs:
            agent_obs.append(batched_obs[agent_index])
        agent_mask = None
        if self.action_mask is not None:
            agent_mask = []
            for mask in self.action_mask:
                agent_mask.append(mask[agent_index])
        group_id = self.group_id[agent_index]
        return DecisionStep(
            obs=agent_obs,
            reward=self.reward[agent_index],
            agent_id=agent_id,
            action_mask=agent_mask,
            group_id=group_id,
            group_reward=self.group_reward[agent_index],
        )

View on GitHub (pinned to 3ecb446f75)

Solutions

  1. Check membership first: only index d[agent_id] if agent_id in d.agent_id_to_index.
  2. Iterate over d.agent_id instead of a stale external id list so you only touch present agents.
  3. Handle agents found in TerminalSteps separately — get them from the TerminalSteps object returned by env.get_steps().

Example fix

# before
step = decision_steps[agent_id]  # KeyError when agent is done
# after
if agent_id in decision_steps:
    step = decision_steps[agent_id]
else:
    step = terminal_steps[agent_id]  # agent finished this step
Defensive patterns

Strategy: type-guard

Validate before calling

if agent_id not in decision_steps.agent_id_to_index:
    # agent is done or absent this step — handle via TerminalSteps

Type guard

def in_decision(ds, agent_id: int) -> bool:
    return agent_id in ds.agent_id_to_index

Try / catch

try:
    step = decision_steps[agent_id]
except KeyError:
    step = None  # agent not deciding this step

Prevention

When it happens

Trigger: Indexing d = DecisionSteps() or the DecisionSteps from env.get_steps() with an agent id that is done/absent this step: d[agent_id] where agent_id not in d.agent_id_to_index — commonly an agent that terminated (moved to TerminalSteps) or hasn't been spawned yet.

Common situations: Looping over a cached list of all agent ids while some have finished their episode (they're in TerminalSteps, not DecisionSteps); querying before the first env step; agent ids reset between episodes with new ids for new spawns.

Related errors

AI-assisted analysis of Unity-Technologies/ml-agents@3ecb446f75 (2026-09-02). Data as JSON: /api/errors/9585e70c4bdd5c1e. Report an issue: GitHub.