DQN Pong Bolts Example

DQN Pong Bolts Example.py

from pl_bolts.models.rl.common import wrappers, cli
from pl_bolts.models.rl.dqn_model import DQN

parser = argparse.ArgumentParser(add_help=False)

# Trainer args
parser = pl.Trainer.add_argparse_args(parser)

# Model args
parser = DQN.add_model_specific_args(parser)
args = parser.parse_args()

model = DQN(**args.__dict__)

# Saving model
checkpoint_callback = ModelCheckpoint(
    save_top_k=1,
    monitor='avg_reward',
    mode='max',
    period=100
)

# Setup Trainer
seed_everything(123)
trainer = pl.Trainer.from_argparse_args(args, checkpoint_callback=checkpoint_callback)

# Train model
trainer.fit(model)