Merge pull request #417 from kywch/g2048-65k

New baseline for G2048

Author: jsuarez5341

pufferlib/config/ocean/g2048.ini

@@ -1,68 +1,167 @@
 [base]
 package = ocean
 env_name = puffer_g2048
-policy_name = Policy
+policy_name = G2048
 rnn_name = Recurrent
 
 [policy]
-hidden_size = 256
+hidden_size = 512
 
 [rnn]
-input_size = 256
-hidden_size = 256
+input_size = 512
+hidden_size = 512
 
 [vec]
 num_envs = 4
 
 [env]
 num_envs = 4096
+reward_scaler = 0.67
+endgame_env_prob = 0.05
+scaffolding_ratio = 0.67
+use_heuristic_rewards = True
+snake_reward_weight = 0.0005
 
 [train]
-# https://wandb.ai/kywch/pufferlib/runs/n8xml0u9?nw=nwuserkywch
-total_timesteps = 3_000_000_000
+# 512 hidden: https://wandb.ai/kywch/pufferlib/runs/5thsjr61?nw=nwuserkywch
+total_timesteps = 6_767_676_767
 anneal_lr = True
+min_learning_rate = 0.00005
 batch_size = auto
 bptt_horizon = 64
-minibatch_size = 65536
+minibatch_size = 32768
 
-adam_beta1 = 0.99
-adam_beta2 = 0.96
-adam_eps = 1.0e-10
-clip_coef = 0.1
-ent_coef = 0.02
-gae_lambda = 0.6
-gamma = 0.985
-learning_rate = 0.001
-max_grad_norm = 1.0
-prio_alpha = 0.99
-prio_beta0 = 0.40
-vf_clip_coef = 0.1
+clip_coef = 0.067
+ent_coef = 0.0267
+gae_lambda = 0.67
+gamma = 0.99567
+vf_clip_coef = 0.167
 vf_coef = 2.0
-vtrace_c_clip = 4.3
-vtrace_rho_clip = 1.6
 
+learning_rate = 0.000467
+max_grad_norm = 0.5
+
+
+# These are newer puffer PPO params. Need more sweeping.
+adam_beta1 = 0.99
+adam_beta2 = 0.9999
+adam_eps = 0.0001
+prio_alpha = 0.8
+prio_beta0 = 0.1
+vtrace_c_clip = 2.0
+vtrace_rho_clip = 1.1
+
+
+### Targeted sweep
 
 [sweep]
 metric = score
 goal = maximize
+max_suggestion_cost = 7200
+sweep_only = endgame_env_prob, scaffolding_ratio, snake_reward_weight, learning_rate, max_grad_norm
+downsample = 1
 
-[sweep.train.total_timesteps]
-distribution = log_normal
-min = 3e8
-max = 1e10
-mean = 1e9
-scale = time
+[sweep.env.endgame_env_prob]
+distribution = uniform
+min = 0.0
+mean = 0.03
+max = 0.2
+scale = auto
+
+[sweep.env.scaffolding_ratio]
+distribution = uniform
+min = 0.1
+mean = 0.5
+max = 0.8
+scale = auto
+
+[sweep.env.snake_reward_weight]
+distribution = uniform
+min = 0.0001
+mean = 0.0007
+max = 0.0050
+scale = auto
 
 [sweep.train.learning_rate]
-distribution = log_normal
-min = 0.00001
-mean = 0.001
-max = 0.1
+distribution = uniform
+min = 0.0001
+mean = 0.0005
+max = 0.0030
 scale = 0.5
 
-[sweep.train.gae_lambda]
-distribution = logit_normal
-min = 0.01
-mean = 0.6
-max = 0.995
-scale = auto
+[sweep.train.max_grad_norm]
+distribution = uniform
+min = 0.1
+mean = 0.5
+max = 2.0
+scale = 0.5
+
+[sweep.train.vf_clip_coef]
+distribution = uniform
+min = 0.05
+max = 0.5
+mean = 0.2
+scale = auto
+
+
+### Broad sweep
+
+; [sweep]
+; metric = score
+; goal = maximize
+
+; [sweep.env.reward_scaler]
+; distribution = uniform
+; min = 0.1
+; mean = 0.5
+; max = 1.0
+; scale = auto
+
+; [sweep.env.scaffolding_ratio]
+; distribution = uniform
+; min = 0.0
+; mean = 0.5
+; max = 0.8
+; scale = auto
+
+; [sweep.env.snake_reward_weight]
+; distribution = uniform
+; min = 0.00001
+; mean = 0.00005
+; max = 0.0002
+; scale = auto
+
+; [sweep.train.total_timesteps]
+; distribution = log_normal
+; min = 3e8
+; max = 1e10
+; mean = 1e9
+; scale = time
+
+; [sweep.train.learning_rate]
+; distribution = log_normal
+; min = 0.00001
+; mean = 0.001
+; max = 0.1
+; scale = 0.5
+
+; [sweep.train.gamma]
+; distribution = logit_normal
+; min = 0.8
+; mean = 0.995
+; max = 0.9999
+; scale = auto
+
+; [sweep.train.gae_lambda]
+; distribution = logit_normal
+; min = 0.01
+; mean = 0.7
+; max = 0.995
+; scale = auto
+
+; [sweep.train.clip_coef]
+; distribution = log_normal
+; min = 0.001
+; max = 0.5
+; mean = 0.05
+; scale = auto

pufferlib/ocean/g2048/binding.c

@@ -3,9 +3,15 @@
 #define Env Game
 #include "../env_binding.h"
 
-// g2048.h does not have a 'size' field, so my_init can just return 0
 static int my_init(Env* env, PyObject* args, PyObject* kwargs) {
-    // No custom initialization needed for 2048
+    env->can_go_over_65536 = unpack(kwargs, "can_go_over_65536");
+    env->reward_scaler = unpack(kwargs, "reward_scaler");
+    env->endgame_env_prob = unpack(kwargs, "endgame_env_prob");
+    env->scaffolding_ratio = unpack(kwargs, "scaffolding_ratio");
+    env->use_heuristic_rewards = unpack(kwargs, "use_heuristic_rewards");
+    env->snake_reward_weight = unpack(kwargs, "snake_reward_weight");
+    env->use_sparse_reward = unpack(kwargs, "use_sparse_reward");
+    init(env);
     return 0;
 }
 
@@ -15,5 +21,11 @@ static int my_log(PyObject* dict, Log* log) {
     assign_to_dict(dict, "merge_score", log->merge_score);
     assign_to_dict(dict, "episode_return", log->episode_return);
     assign_to_dict(dict, "episode_length", log->episode_length);
+    assign_to_dict(dict, "lifetime_max_tile", log->lifetime_max_tile);
+    assign_to_dict(dict, "reached_32768", log->reached_32768);
+    assign_to_dict(dict, "reached_65536", log->reached_65536);
+    assign_to_dict(dict, "monotonicity_reward", log->monotonicity_reward);
+    assign_to_dict(dict, "snake_state", log->snake_state);
+    assign_to_dict(dict, "snake_reward", log->snake_reward);
     return 0;
 }

pufferlib/ocean/g2048/eval.py

@@ -0,0 +1,105 @@
+from pufferlib import pufferl
+
+def evaluate(env_name, load_model_path):
+    args = pufferl.load_config(env_name)
+    args['vec']['num_envs'] = 1
+    args['env']['num_envs'] = 4096
+    args['load_model_path'] = load_model_path
+    # Turn off endgame_envs and scaffolding episodes, which do not report results
+    args['env']['endgame_env_prob'] = 0
+    args['env']['scaffolding_ratio'] = 0
+    args['env']['can_go_over_65536'] = True
+
+    vecenv = pufferl.load_env(env_name, args)
+    policy = pufferl.load_policy(args, vecenv, env_name)
+    trainer = pufferl.PuffeRL(args['train'], vecenv, policy)
+
+    # Each evaluate runs for 64 ticks. NOTE: bppt horizon might be short for g2048?
+    # Avg episode length from the current model is ~18000, so it takes ~300 epochs for an avg episode.
+    # It's hard to get the single best score because stats are already averaged across done envs.
+    for i in range(10000):
+        stats = trainer.evaluate()
+
+        trainer.epoch += 1
+        if i % 20 == 0:
+            trainer.print_dashboard()
+
+    trainer.close()
+
+    # Get the estimates
+    num_episodes = sum(stats['n'])
+    episode_lengths = sum(n * l for n, l in zip(stats['n'], stats['episode_length'])) / num_episodes
+    max_tiles = sum(n * m for n, m in zip(stats['n'], stats['score'])) / num_episodes
+    merge_scores = sum(n * s for n, s in zip(stats['n'], stats['merge_score'])) / num_episodes
+    reached_32768 = sum(n * s for n, s in zip(stats['n'], stats['reached_32768'])) / num_episodes
+    reached_65536 = sum(n * s for n, s in zip(stats['n'], stats['reached_65536'])) / num_episodes
+
+    print(f"Num episodes: {int(num_episodes)}")
+    print(f"Max tile avg: {max_tiles:.1f}")
+    # The stats from vecenv are averaged across envs that were done in the same tick. Cannot get the single max.
+    print(f"Episode length -- Avg: {episode_lengths:.1f}, Max: {max(stats['episode_length']):.1f}")
+    print(f"Merge score -- Avg: {merge_scores:.1f}, Max: {max(stats['merge_score']):.1f}")
+    print(f"Reached 32768 prob: {reached_32768*100:.2f} %")
+    print(f"Reached 65536 prob: {reached_65536*100:.2f} %")
+
+    """
+    # hidden 256: https://wandb.ai/kywch/pufferlib/runs/nvd0pfuj?nw=nwuserkywch
+    Num episodes: 154406
+    Max tile avg: 22532.9
+    Episode length -- Avg: 16667.2, Max: 26659.1
+    Merge score -- Avg: 462797.9, Max: 744224.9
+    Reached 32768 prob: 46.08 %
+    Reached 65536 prob: 3.53 %
+
+    # hidden 512: https://wandb.ai/kywch/pufferlib/runs/2ch3my60?nw=nwuserkywch
+    Num episodes: 119243
+    Max tile avg: 30662.2
+    Episode length -- Avg: 21539.7, Max: 29680.3
+    Merge score -- Avg: 618011.8, Max: 918755.8
+    Reached 32768 prob: 68.25 %
+    Reached 65536 prob: 13.09 %
+
+    # hidden 512 (replication): https://wandb.ai/kywch/pufferlib/runs/5thsjr61?nw=nwuserkywch
+    Num episodes: 115652
+    Max tile avg: 31773.2
+    Episode length -- Avg: 22196.4, Max: 30316.5
+    Merge score -- Avg: 639395.6, Max: 909969.8
+    Reached 32768 prob: 71.22 %
+    Reached 65536 prob: 14.75 %
+    """
+
+def finetune(env_name, load_model_path):
+    args = pufferl.load_config(env_name)
+    args['load_model_path'] = load_model_path
+    # args['env']['use_sparse_reward'] = True
+    args['env']['scaffolding_ratio'] = 0.85
+
+    # args['policy']['hidden_size'] = 512
+    # args['rnn']['input_size'] = 512
+    # args['rnn']['hidden_size'] = 512
+
+    args['train']['total_timesteps'] = 1_000_000_000
+    args['train']['learning_rate'] = 0.00005
+    args['train']['anneal_lr'] = False
+
+    args['wandb'] = True
+    args['tag'] = 'pg2048'
+
+    pufferl.train(env_name, args)
+
+if __name__ == '__main__':
+    import os
+    import wandb
+
+    # https://wandb.ai/kywch/pufferlib/runs/5thsjr61?nw=nwuserkywch
+    wandb_run_id = '5thsjr61'
+    wandb.init(id=wandb_run_id, project='pufferlib', entity='kywch')
+
+    artifact = wandb.use_artifact(f'{wandb_run_id}:latest')
+    data_dir = artifact.download()
+    model_file = max(os.listdir(data_dir))
+    model_path = f'{data_dir}/{model_file}'
+    wandb.finish()
+
+    evaluate('puffer_g2048', load_model_path=model_path)
+    # finetune('puffer_g2048', load_model_path='puffer_g2048_256_base.pt')