[RLlib] Add correct terminated and truncated batch sizes on zero-leng…

…th episodes (#46721)
ray-project · Sep 25, 2024 · 4a5207d · 4a5207d
1 parent a6cf9d7
commit 4a5207d
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/rllib/connectors/learner/add_columns_from_episodes_to_train_batch.py b/rllib/connectors/learner/add_columns_from_episodes_to_train_batch.py
@@ -117,6 +117,8 @@ def __call__(
                     Columns.TERMINATEDS,
                     items_to_add=(
                         [False] * (len(sa_episode) - 1) + [sa_episode.is_terminated]
+                        if len(sa_episode) > 0
+                        else []
                     ),
                     num_items=len(sa_episode),
                     single_agent_episode=sa_episode,
@@ -132,6 +134,8 @@ def __call__(
                     Columns.TRUNCATEDS,
                     items_to_add=(
                         [False] * (len(sa_episode) - 1) + [sa_episode.is_truncated]
+                        if len(sa_episode) > 0
+                        else []
                     ),
                     num_items=len(sa_episode),
                     single_agent_episode=sa_episode,