Skip to content

lit_sb3.vec_env.patch_gym

[docs] module lit_sb3.vec_env.patch_gym

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
from inspect import signature
from typing import Union
import warnings

import gymnasium

try:
    import gym  # type: ignore[import-untyped]

    gym_installed = True
except ImportError:
    gym_installed = False


def _patch_env(env: Union["gym.Env", gymnasium.Env]) -> gymnasium.Env:  # pragma: no cover
    """Adapted from https://github.com/thu-ml/tianshou.

    Takes an environment and patches it to return Gymnasium env.
    This function takes the environment object and returns a patched
    env, using shimmy wrapper to convert it to Gymnasium,
    if necessary.

    Args:
        env: A gym/gymnasium env

    Returns:
        Patched env (gymnasium env)
    """
    # Gymnasium env, no patching to be done
    if isinstance(env, gymnasium.Env):
        return env

    if not gym_installed or not isinstance(env, gym.Env):
        raise ValueError(
            f"The environment is of type {type(env)}, not a Gymnasium "
            f"environment. In this case, we expect OpenAI Gym to be "
            f"installed and the environment to be an OpenAI Gym environment."
        )

    try:
        import shimmy  # type: ignore[import-untyped]
    except ImportError as e:
        raise ImportError(
            "Missing shimmy installation. You provided an OpenAI Gym environment. "
            "Stable-Baselines3 (SB3) has transitioned to using Gymnasium internally. "
            "In order to use OpenAI Gym environments with SB3, you need to "
            "install shimmy (`pip install 'shimmy>=2.0'`)."
        ) from e

    warnings.warn(
        "You provided an OpenAI Gym environment. "
        "We strongly recommend transitioning to Gymnasium environments. "
        "Stable-Baselines3 is automatically wrapping your environments in a compatibility "
        "layer, which could potentially cause issues.",
        stacklevel=2,
    )

    if "seed" in signature(env.unwrapped.reset).parameters:
        # Gym 0.26+ env
        return shimmy.GymV26CompatibilityV0(env=env)
    # Gym 0.21 env
    return shimmy.GymV21CompatibilityV0(env=env)


def _convert_space(space: Union["gym.Space", gymnasium.Space]) -> gymnasium.Space:  # pragma: no cover
    """Takes a space and patches it to return Gymnasium Space.

    This function takes the space object and returns a patched
    space, using shimmy wrapper to convert it to Gymnasium,
    if necessary.

    Args:
        env: A gym/gymnasium Space

        space: The space parameter.

    Returns:
        Patched space (gymnasium Space)
    """
    # Gymnasium space, no conversion to be done
    if isinstance(space, gymnasium.Space):
        return space

    if not gym_installed or not isinstance(space, gym.Space):
        raise ValueError(
            f"The space is of type {type(space)}, not a Gymnasium "
            f"space. In this case, we expect OpenAI Gym to be "
            f"installed and the space to be an OpenAI Gym space."
        )

    try:
        import shimmy  # type: ignore[import-untyped]
    except ImportError as e:
        raise ImportError(
            "Missing shimmy installation. You provided an OpenAI Gym space. "
            "Stable-Baselines3 (SB3) has transitioned to using Gymnasium internally. "
            "In order to use OpenAI Gym space with SB3, you need to "
            "install shimmy (`pip install 'shimmy>=0.2.1'`)."
        ) from e

    warnings.warn(
        "You loaded a model that was trained using OpenAI Gym. "
        "We strongly recommend transitioning to Gymnasium by saving that model again.",
        stacklevel=2,
    )

    return shimmy.openai_gym_compatibility._convert_space(space)