ray-mario-main.py

def main():
    def env_creator_lambda(env_config):
        return env_creator(args.environment,
                           config,
                           args.dimension,
                           args.framestack)

    args = parse_args()
    config = {
        'env': 'super_mario_bros',
        'framework': 'torch',
        'rollout_fragment_length': 50,
        'train_batch_size': 500,
        'num_workers': args.workers,
        'num_envs_per_worker': 1,
        'num_gpus': args.gpus
    }
    ray.init()

    register_env('super_mario_bros', env_creator_lambda)
    trainer = ImpalaTrainer(config=config)

    if args.checkpoint:
        trainer.restore(args.checkpoint)

    for iteration in range(args.iterations):
        result = trainer.train()
        print_results(result, iteration)

        if iteration % 50 == 0:
            checkpoint = trainer.save()
            print('Checkpoint saved at', checkpoint)


if __name__ == "__main__":
    main()