code: def test_standard_rollout_with_logits_and_step_stats():
      """Recording logits and step statistics."""
      env = build_dummy_structured_env()
      rollout_generator = RolloutGenerator(
          env=env, record_step_stats=True, record_logits=True
      )
      policy = flatten_concat_probabilistic_policy_for_env(
          env
      )  # We need a torch policy to be able to record logits
      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 record.step_stats is not None
          assert sub_step_keys == record.logits_dict.keys()
  
          for step_key in sub_step_keys:
              assert (
                  record.logits_dict[step_key].keys()
                  == record.actions_dict[step_key].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 3 for calling a global function with example usage: test_standard_rollout_with_logits_and_step_stats() without return types