code: def load(
      environment_name: Text,
      discount: types.Float = 1.0,
      max_episode_steps: Optional[types.Int] = None,
      gym_env_wrappers: Sequence[types.GymnasiumEnvWrapper] = (),
      env_wrappers: Sequence[types.PyEnvWrapper] = (),
      gym_kwargs: Optional[Dict[str, Any]] = None,
      render_kwargs: Optional[Dict[str, Any]] = None,
  ) -> py_environment.PyEnvironment:
      """Loads the selected environment and wraps it with the specified wrappers.
  
      Note that by default a TimeLimit wrapper is used to limit episode lengths
      to the default benchmarks defined by the registered environments.
  
      Args:
        environment_name: Name for the environment to load.
        discount: Discount to use for the environment.
        max_episode_steps: If None the max_episode_steps will be set to the default
          step limit defined in the environment's spec. No limit is applied if set
          to 0 or if there is no max_episode_steps set in the environment's spec.
        gym_env_wrappers: Iterable with references to wrapper classes to use
          directly on the gym environment.
        env_wrappers: Iterable with references to wrapper classes to use on the
          gym_wrapped environment.
        gym_kwargs: Optional kwargs to pass to the Gym environment class.
        render_kwargs: Optional kwargs for rendering to pass to `make()` of the
          gymnasium_wrapped environment.
  
      Returns:
        A PyEnvironment instance.
      """
      gym_kwargs = gym_kwargs if gym_kwargs else {}
      render_kwargs = render_kwargs if render_kwargs else {}
      make_args = {**gym_kwargs, **render_kwargs}
      gym_spec = gym.spec(environment_name)
      gym_env = gym_spec.make(**make_args)
  
      if (
          max_episode_steps is None
          and gym_spec.max_episode_steps is not None
      ):
          max_episode_steps = gym_spec.max_episode_steps
  
      return wrap_env(
          gym_env,
          discount=discount,
          max_episode_steps=max_episode_steps,
          gym_env_wrappers=gym_env_wrappers,
          env_wrappers=env_wrappers,
      )
  
file_path: agents/tf_agents/environments/suite_gymnasium.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Any, Callable, Dict, Optional, Sequence, Text
  import gin
  import gymnasium
  from tf_agents.environments import gymnasium_wrapper
  from tf_agents.environments import py_environment
  from tf_agents.environments import wrappers
  from tf_agents.typing import types
label_desc: Use image node_id 1 for calling a global function with example usage: load(environment_name, discount, max_episode_steps, gym_env_wrappers, env_wrappers, gym_kwargs, render_kwargs) and returns: wrap_env