code: def load(
      environment_name: Text,
      discount: types.Int = 1.0,
      max_episode_steps: Optional[types.Int] = None,
      gym_env_wrappers: Sequence[
          types.GymEnvWrapper
      ] = DEFAULT_ATARI_GYM_WRAPPERS,
      env_wrappers: Sequence[types.PyEnvWrapper] = (),
      spec_dtype_map: Optional[Dict[gym.Space, np.dtype]] = None,
      gym_kwargs: Optional[Dict[str, Any]] = None,
  ) -> py_environment.PyEnvironment:
      """Loads the selected environment and wraps it with the specified wrappers."""
      if spec_dtype_map is None:
          spec_dtype_map = {gym.spaces.Box: np.uint8}
  
      gym_kwargs = gym_kwargs if gym_kwargs else {}
      gym_spec = gym.spec(environment_name)
      gym_env = gym_spec.make(**gym_kwargs)
  
      if (
          max_episode_steps is None
          and gym_spec.max_episode_steps is not None
      ):
          max_episode_steps = gym_spec.max_episode_steps
  
      return suite_gym.wrap_env(
          gym_env,
          discount=discount,
          max_episode_steps=max_episode_steps,
          gym_env_wrappers=gym_env_wrappers,
          time_limit_wrapper=atari_wrappers.AtariTimeLimit,
          env_wrappers=env_wrappers,
          spec_dtype_map=spec_dtype_map,
          auto_reset=False,
      )
  
file_path: agents/tf_agents/environments/suite_atari.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Dict, Optional, Sequence, Text, Any
  import ale_py
  import gin
  import gym
  import numpy
  from tf_agents.environments import atari_preprocessing
  from tf_agents.environments import atari_wrappers
  from tf_agents.environments import py_environment
  from tf_agents.environments import suite_gym
  from tf_agents.typing import types
label_desc: Use image node_id 2 for calling a global function with example usage: load(environment_name, discount, max_episode_steps, gym_env_wrappers, env_wrappers, spec_dtype_map, gym_kwargs) and returns: suite_gym