code: def test_handles_done_in_substep_with_recorded_episode_stats():
      """Recording episode stats and handling environments that return done during a (non-last) sub-step."""
      env = build_dummy_structured_env()
      env = _FiveSubstepsLimitWrapper.wrap(env)
      policy = RandomPolicy(env.action_spaces_dict)
  
      # -- Normal operation (should reset the env automatically and continue rollout) --
      rollout_generator = RolloutGenerator(
          env=env, record_episode_stats=True
      )
      trajectory = rollout_generator.rollout(policy, n_steps=10)
      assert len(trajectory) == 10
  
      # The done step records should have data for the first sub-step only
      dones = 0
      for step_record in trajectory.step_records:
          if step_record.is_done():
              assert [0] == list(step_record.observations_dict.keys())
              dones += 1
              assert step_record.episode_stats is not None
          else:
              assert [0, 1] == list(
                  step_record.observations_dict.keys()
              )
              assert step_record.episode_stats is None
      assert (
          dones == 3
      )  # Each episode is done after 5 sub-steps, i.e. 3 structured steps get recorded => 3 episodes fit
  
      # -- Terminate on done --
      rollout_generator = RolloutGenerator(
          env=env, terminate_on_done=True
      )
      trajectory = rollout_generator.rollout(policy, n_steps=10)
      assert len(trajectory) == 3
      assert trajectory.is_done()
      assert [0] == list(
          trajectory.step_records[-1].observations_dict.keys()
      )
  
imports: from typing import Dict, Any, Tuple
  import numpy
  from maze.core.agent.random_policy import RandomPolicy, DistributedRandomPolicy
  from maze.core.env.base_env import BaseEnv
  from maze.core.rollout.rollout_generator import RolloutGenerator
  from maze.core.wrappers.time_limit_wrapper import TimeLimitWrapper
  from maze.test.shared_test_utils.helper_functions import build_dummy_structured_env, build_dummy_maze_env, build_dummy_maze_env_with_structured_core_env
  from maze.test.shared_test_utils.helper_functions import flatten_concat_probabilistic_policy_for_env
  from maze.train.parallelization.vector_env.sequential_vector_env import SequentialVectorEnv
label_desc: Use image node_id 5 for calling a global function with example usage: test_handles_done_in_substep_with_recorded_episode_stats() without return types