code: def test_vectorized_rollout():
      """Rollout with a vector env."""
      concurrency = 3
      env = SequentialVectorEnv(
          [build_dummy_structured_env] * concurrency
      )
      rollout_generator = RolloutGenerator(env=env)
      policy = DistributedRandomPolicy(
          env.action_spaces_dict, concurrency=concurrency
      )
      trajectory = rollout_generator.rollout(policy, n_steps=10)
  
      assert len(trajectory) == 10
  
      sub_step_keys = env.action_spaces_dict.keys()
      for record in trajectory.step_records:
          assert sub_step_keys == record.actions_dict.keys()
          assert sub_step_keys == record.observations_dict.keys()
          assert sub_step_keys == record.rewards_dict.keys()
  
          assert record.batch_shape == [concurrency]
          # The first dimension of the observations should correspond to the distributed env concurrency
          # (We just check the very first array present in the first observation)
          first_sub_step_obs: Dict = list(
              record.observations_dict.values()
          )[0]
          first_obs_value = list(first_sub_step_obs.values())[0]
          assert first_obs_value.shape[0] == concurrency
  
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 2 for calling a global function with example usage: test_vectorized_rollout() without return types