wehub-resource-sync
365 lines
13 KiB
Python
365 lines
13 KiB
Python
import gymnasium as gym
|
|
from gymnasium.wrappers import AtariPreprocessing
|
|
|
|
from ray import tune
|
|
from ray.rllib.algorithms.dqn.dqn import DQNConfig
|
|
from ray.rllib.connectors.env_to_module.frame_stacking import FrameStackingEnvToModule
|
|
from ray.rllib.connectors.learner.frame_stacking import FrameStackingLearner
|
|
from ray.rllib.utils.metrics import (
|
|
ENV_RUNNER_RESULTS,
|
|
EPISODE_RETURN_MEAN,
|
|
NUM_ENV_STEPS_SAMPLED_LIFETIME,
|
|
)
|
|
from ray.tune import Stopper
|
|
|
|
# Might need `gymnasium[atari, other]` to be installed.
|
|
|
|
# See the following links for becnhmark results of other libraries:
|
|
# Original paper: https://arxiv.org/abs/1812.05905
|
|
# CleanRL: https://wandb.ai/cleanrl/cleanrl.benchmark/reports/Mujoco--VmlldzoxODE0NjE
|
|
# AgileRL: https://github.com/AgileRL/AgileRL?tab=readme-ov-file#benchmarks
|
|
|
|
benchmark_envs = {
|
|
"AlienNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 6022.9,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"AmidarNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 202.8,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"AssaultNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 14491.7,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"AsterixNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 280114.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"AsteroidsNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2249.4,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"AtlantisNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 814684.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"BankHeistNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 826.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"BattleZoneNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 52040.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"BeamRiderNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 21768.5,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"BerzerkNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 1793.4,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"BowlingNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 39.4,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"BoxingNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 54.9,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"BreakoutNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 379.5,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"CentipedeNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 7160.9,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"ChopperCommandNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 10916.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"CrazyClimberNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 143962.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"DefenderNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 47671.3,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"DemonAttackNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 109670.7,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"DoubleDunkNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -0.6,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"EnduroNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2061.1,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"FishingDerbyNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 22.6,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"FreewayNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 29.1,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"FrostbiteNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 4141.1,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"GopherNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 72595.7,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"GravitarNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 567.5,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"HeroNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 50496.8,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"IceHockeyNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -11685.8,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"KangarooNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 10841.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"KrullNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 6715.5,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"KungFuMasterNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 28999.8,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"MontezumaRevengeNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 154.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"MsPacmanNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2570.2,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"NameThisGameNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 11686.5,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"PhoenixNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 103061.6,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"PitfallNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -37.6,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"PongNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 19.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"PrivateEyeNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 1704.4,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"QbertNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 18397.6,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"RoadRunnerNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 54261.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"RobotankNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 55.2,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"SeaquestNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 19176.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"SkiingNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -11685.8,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"SolarisNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 2860.7,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"SpaceInvadersNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 12629.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"StarGunnerNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 123853.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"SurroundNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 7.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"TennisNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -2.2,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"TimePilotNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 11190.5,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"TutankhamNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 126.9,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"VentureNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 45.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"VideoPinballNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 506817.2,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"WizardOfWorNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 14631.5,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"YarsRevengeNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 93007.9,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
"ZaxxonNoFrameskip-v4": {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 19658.0,
|
|
f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 200000000,
|
|
},
|
|
}
|
|
|
|
for env in benchmark_envs.keys():
|
|
tune.register_env(
|
|
env,
|
|
lambda ctx, e=env: AtariPreprocessing(
|
|
gym.make(e, **ctx), grayscale_newaxis=True, screen_size=84, noop_max=0
|
|
),
|
|
)
|
|
|
|
|
|
def _make_env_to_module_connector(env, spaces, device):
|
|
return FrameStackingEnvToModule(num_frames=4)
|
|
|
|
|
|
def _make_learner_connector(input_observation_space, input_action_space):
|
|
return FrameStackingLearner(num_frames=4)
|
|
|
|
|
|
# Define a `tune.Stopper` that stops the training if the benchmark is reached
|
|
# or the maximum number of timesteps is exceeded.
|
|
class BenchmarkStopper(Stopper):
|
|
def __init__(self, benchmark_envs):
|
|
self.benchmark_envs = benchmark_envs
|
|
|
|
def __call__(self, trial_id, result):
|
|
# Stop training if the mean reward is reached.
|
|
if (
|
|
result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]
|
|
>= self.benchmark_envs[result["env"]][
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}"
|
|
]
|
|
):
|
|
return True
|
|
# Otherwise check, if the total number of timesteps is exceeded.
|
|
elif (
|
|
result[f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}"]
|
|
>= self.benchmark_envs[result["env"]][f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}"]
|
|
):
|
|
return True
|
|
# Otherwise continue training.
|
|
else:
|
|
return False
|
|
|
|
# Note, this needs to implemented b/c the parent class is abstract.
|
|
def stop_all(self):
|
|
return False
|
|
|
|
|
|
# See Table 1 in the Rainbow paper for the hyperparameters.
|
|
config = (
|
|
DQNConfig()
|
|
.environment(
|
|
env=tune.grid_search(list(benchmark_envs.keys())),
|
|
env_config={
|
|
"max_episode_steps": 108000,
|
|
"obs_type": "grayscale",
|
|
# The authors actually use an action repetition of 4.
|
|
"repeat_action_probability": 0.25,
|
|
},
|
|
clip_rewards=True,
|
|
)
|
|
.env_runners(
|
|
# Every 4 agent steps a training update is performed.
|
|
rollout_fragment_length=4,
|
|
num_env_runners=1,
|
|
env_to_module_connector=_make_env_to_module_connector,
|
|
)
|
|
# TODO (simon): Adjust to new model_config_dict.
|
|
.training(
|
|
# Note, the paper uses also an Adam epsilon of 0.00015.
|
|
lr=0.0000625,
|
|
n_step=3,
|
|
tau=1.0,
|
|
train_batch_size=32,
|
|
target_network_update_freq=32000,
|
|
replay_buffer_config={
|
|
"type": "PrioritizedEpisodeReplayBuffer",
|
|
"capacity": 1000000,
|
|
"alpha": 0.5,
|
|
# Note the paper used a linear schedule for beta.
|
|
"beta": 0.4,
|
|
},
|
|
# Note, these are frames.
|
|
num_steps_sampled_before_learning_starts=80000,
|
|
noisy=True,
|
|
num_atoms=51,
|
|
v_min=-10.0,
|
|
v_max=10.0,
|
|
double_q=True,
|
|
dueling=True,
|
|
model={
|
|
"cnn_filter_specifiers": [[32, 8, 4], [64, 4, 2], [64, 3, 1]],
|
|
"fcnet_activation": "tanh",
|
|
"post_fcnet_hiddens": [512],
|
|
"post_fcnet_activation": "relu",
|
|
"post_fcnet_weights_initializer": "orthogonal_",
|
|
"post_fcnet_weights_initializer_config": {"gain": 0.01},
|
|
},
|
|
learner_connector=_make_learner_connector,
|
|
)
|
|
.reporting(
|
|
metrics_num_episodes_for_smoothing=10,
|
|
min_sample_timesteps_per_iteration=1000,
|
|
)
|
|
.evaluation(
|
|
evaluation_duration="auto",
|
|
evaluation_interval=1,
|
|
evaluation_num_env_runners=1,
|
|
evaluation_parallel_to_training=True,
|
|
evaluation_config={
|
|
"explore": False,
|
|
},
|
|
)
|
|
)
|
|
|
|
tuner = tune.Tuner(
|
|
"DQN",
|
|
param_space=config,
|
|
run_config=tune.RunConfig(
|
|
stop=BenchmarkStopper(benchmark_envs=benchmark_envs),
|
|
name="benchmark_dqn_atari",
|
|
),
|
|
)
|
|
|
|
tuner.fit()
|