code: def wrap_env(
      gym_env: gym.Env,
      discount: types.Float = 1.0,
      max_episode_steps: Optional[types.Int] = None,
      gym_env_wrappers: Sequence[types.GymnasiumEnvWrapper] = (),
      time_limit_wrapper: TimeLimitWrapperType = wrappers.TimeLimit,
      env_wrappers: Sequence[types.PyEnvWrapper] = (),
      auto_reset: bool = True,
  ) -> py_environment.PyEnvironment:
      """Wraps given gymnasium environment with TF Agent's GymnasiumWrapper.
  
      Note that by default a TimeLimit wrapper is used to limit episode lengths
      to the default benchmarks defined by the registered environments.
  
      Args:
        gym_env: An instance of OpenAI gym environment.
        discount: Discount to use for the environment.
        max_episode_steps: Used to create a TimeLimitWrapper. No limit is applied if
          set to None or 0. Usually set to `gym_spec.max_episode_steps` in `load`.
        gym_env_wrappers: Iterable with references to wrapper classes to use
          directly on the gym environment.
        time_limit_wrapper: Wrapper that accepts (env, max_episode_steps) params to
          enforce a TimeLimit. Usuaully this should be left as the default,
          wrappers.TimeLimit.
        env_wrappers: Iterable with references to wrapper classes to use on the
          gym_wrapped environment.
        auto_reset: If True (default), reset the environment automatically after a
          terminal state is reached.
  
      Returns:
        A PyEnvironment instance.
      """
  
      for wrapper in gym_env_wrappers:
          gym_env = wrapper(gym_env)
      env = gym_wrapper.GymnasiumWrapper(
          gym_env,
          discount=discount,
          auto_reset=auto_reset,
      )
  
      if max_episode_steps is not None and max_episode_steps > 0:
          env = time_limit_wrapper(env, max_episode_steps)
  
      for wrapper in env_wrappers:
          env = wrapper(env)
  
      return env
  
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 2 for calling a global function with example usage: wrap_env(gym_env, discount, max_episode_steps, gym_env_wrappers, time_limit_wrapper, env_wrappers, auto_reset) and returns: env