code: def test_records_next_observations():
      """Recording next observations."""
      env = build_dummy_structured_env()
      rollout_generator = RolloutGenerator(
          env=env, record_next_observations=True
      )
      policy = RandomPolicy(env.action_spaces_dict)
      trajectory = rollout_generator.rollout(policy, n_steps=10)
  
      assert len(trajectory) == 10
  
      sub_step_keys = env.action_spaces_dict.keys()
      last_next_obs = None
      for record in trajectory.step_records:
          assert sub_step_keys == record.observations_dict.keys()
          assert sub_step_keys == record.next_observations_dict.keys()
          assert record.batch_shape is None
  
          for step_key in sub_step_keys:
              curr_obs = record.observations_dict[step_key]
  
              # Next obs from the previous sub-step should be equal to the current observation
              if last_next_obs:
                  assert list(curr_obs.keys()) == list(
                      last_next_obs.keys()
                  )
                  for obs_key in curr_obs.keys():
                      assert np.all(
                          curr_obs[obs_key] == last_next_obs[obs_key]
                      )
  
              last_next_obs = record.next_observations_dict[step_key]
  
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 6 for calling a global function with example usage: test_records_next_observations() without return types