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

…7544) Signed-off-by: Artur Niederfahrenhorst <attaismyname@googlemail.com>
eugo-inc · Jul 28, 2023 · 9a131db · 9a131db
1 parent 7c07659
commit 9a131db
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