Fix the done flag setting in one-step episodes (#22)

cogment · Apr 10, 2024 · e80c032 · e80c032
1 parent 31bf263
commit e80c032
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/cogment_lab/envs/gymnasium.py b/cogment_lab/envs/gymnasium.py
@@ -198,7 +198,7 @@ async def reset(self, state: State):
 
         logging.info("Resetting environment")
 
-        obs, _info = state.env.reset(seed=state.session_cfg.seed, options=state.session_cfg.reset_args)  # THIS
+        obs, _info = state.env.reset(seed=state.session_cfg.seed, options=state.session_cfg.reset_args)
 
         state.observation_space = state.session_helper.get_observation_space(self.actor_name)
         frame = state.env.render() if state.session_cfg.render else None

diff --git a/cogment_lab/utils/trial_utils.py b/cogment_lab/utils/trial_utils.py
@@ -179,7 +179,8 @@ def extract_data_from_samples(
         data.rewards = initialize_buffer(None, sample_count - 1)  # type: ignore
     if "done" in fields:
         data.done = initialize_buffer(None, sample_count - 1)  # type: ignore
-        data.done[-1] = True  # type: ignore
+        if sample_count > 1:
+            data.done[-1] = True  # type: ignore
     if "next_observations" in fields:
         data.next_observations = initialize_buffer(observation_space, sample_count - 1)
     if "last_observation" in fields: