import gym

from transfer.envs.utils.wrappers import SuccessCounter


def get_gym_env(env_name):
    env = gym.make(env_name)
    env = SuccessCounter(env)
    env.name = env_name
    env.num_envs = 1
    return env


def get_reacher_env():
    from transfer.envs.reacher_2d import Reacher2dEnv

    env = Reacher2dEnv()
    env = SuccessCounter(env)
    env.name = "Reacher2dEnv"
    env.num_envs = 1
    return env
