from gym.envs.registration import register
from d4rl.gym_mujoco import gym_envs

HOPPER_RANDOM_SCORE = -20.272305
HALFCHEETAH_RANDOM_SCORE = -280.178953
WALKER_RANDOM_SCORE = 1.629008
ANT_RANDOM_SCORE = -325.6

HOPPER_EXPERT_SCORE = 3234.3
HALFCHEETAH_EXPERT_SCORE = 12135.0
WALKER_EXPERT_SCORE = 4592.3
ANT_EXPERT_SCORE = 3879.7

# Single Policy datasets
register(
    id='hopper-medium-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_hopper_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': HOPPER_RANDOM_SCORE,
        'ref_max_score': HOPPER_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/hopper_medium.hdf5'
    }
)

register(
    id='halfcheetah-medium-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_cheetah_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': HALFCHEETAH_RANDOM_SCORE,
        'ref_max_score': HALFCHEETAH_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/halfcheetah_medium.hdf5'
    }
)

register(
    id='walker2d-medium-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_walker_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': WALKER_RANDOM_SCORE,
        'ref_max_score': WALKER_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/walker2d_medium.hdf5'
    }
)

register(
    id='hopper-expert-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_hopper_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': HOPPER_RANDOM_SCORE,
        'ref_max_score': HOPPER_EXPERT_SCORE,
        'dataset_url': 'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/hopper_expert.hdf5'
    }
)

register(
    id='halfcheetah-expert-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_cheetah_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': HALFCHEETAH_RANDOM_SCORE,
        'ref_max_score': HALFCHEETAH_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/halfcheetah_expert.hdf5'
    }
)

register(
    id='walker2d-expert-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_walker_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': WALKER_RANDOM_SCORE,
        'ref_max_score': WALKER_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/walker2d_expert.hdf5'
    }
)

register(
    id='hopper-random-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_hopper_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': HOPPER_RANDOM_SCORE,
        'ref_max_score': HOPPER_EXPERT_SCORE,
        'dataset_url': 'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/hopper_random.hdf5'
    }
)

register(
    id='halfcheetah-random-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_cheetah_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': HALFCHEETAH_RANDOM_SCORE,
        'ref_max_score': HALFCHEETAH_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/halfcheetah_random.hdf5'
    }
)

register(
    id='walker2d-random-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_walker_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': WALKER_RANDOM_SCORE,
        'ref_max_score': WALKER_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/walker2d_random.hdf5'
    }
)

# Mixed datasets
register(
    id='hopper-medium-replay-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_hopper_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': HOPPER_RANDOM_SCORE,
        'ref_max_score': HOPPER_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/hopper_mixed.hdf5'
    },
)

register(
    id='walker2d-medium-replay-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_walker_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': WALKER_RANDOM_SCORE,
        'ref_max_score': WALKER_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/walker_mixed.hdf5'
    }
)

register(
    id='halfcheetah-medium-replay-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_cheetah_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': HALFCHEETAH_RANDOM_SCORE,
        'ref_max_score': HALFCHEETAH_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/halfcheetah_mixed.hdf5'
    }
)

# Mixtures of random/medium and experts
register(
    id='walker2d-medium-expert-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_walker_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': WALKER_RANDOM_SCORE,
        'ref_max_score': WALKER_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/walker2d_medium_expert.hdf5'
    }
)

register(
    id='halfcheetah-medium-expert-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_cheetah_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': HALFCHEETAH_RANDOM_SCORE,
        'ref_max_score': HALFCHEETAH_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/halfcheetah_medium_expert.hdf5'
    }
)

register(
    id='hopper-medium-expert-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_hopper_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': HOPPER_RANDOM_SCORE,
        'ref_max_score': HOPPER_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/hopper_medium_expert_v1.hdf5'
    }
)

register(
    id='ant-medium-expert-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_ant_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': ANT_RANDOM_SCORE,
        'ref_max_score': ANT_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_medium_expert.hdf5'
    }
)

register(
    id='ant-medium-replay-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_ant_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': ANT_RANDOM_SCORE,
        'ref_max_score': ANT_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_mixed.hdf5'
    }
)

register(
    id='ant-medium-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_ant_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': ANT_RANDOM_SCORE,
        'ref_max_score': ANT_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_medium.hdf5'
    }
)

register(
    id='ant-random-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_ant_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': ANT_RANDOM_SCORE,
        'ref_max_score': ANT_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_random.hdf5'
    }
)

register(
    id='ant-expert-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_ant_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': ANT_RANDOM_SCORE,
        'ref_max_score': ANT_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_expert.hdf5'
    }
)

register(
    id='ant-random-expert-v0',
    entry_point='d4rl.gym_mujoco.gym_envs:get_ant_env',
    max_episode_steps=1000,
    kwargs={
        'ref_min_score': ANT_RANDOM_SCORE,
        'ref_max_score': ANT_EXPERT_SCORE,
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_random_expert.hdf5'
    }
)
