[RLlib] Add agent index on first step in env runner v2 (ray-project#3…

…7544) Signed-off-by: Artur Niederfahrenhorst <attaismyname@googlemail.com> Signed-off-by: e428265 <arvind.chandramouli@lmco.com>
lmco · Aug 31, 2023 · 886c92a · 886c92a
1 parent 480cd26
commit 886c92a
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/rllib/evaluation/env_runner_v2.py b/rllib/evaluation/env_runner_v2.py
@@ -781,6 +781,7 @@ def __process_resetted_obs_for_eval(
                         SampleBatch.NEXT_OBS: obs,
                         SampleBatch.INFOS: infos,
                         SampleBatch.T: episode.length,
+                        SampleBatch.AGENT_INDEX: episode.agent_index(agent_id),
                     },
                 )
                 for agent_id, obs in agents_obs