Add league training

scripts/cleanba.py

@@ -43,8 +43,8 @@ class Args:
     """seed of the experiment"""
     log_frequency: int = 10
     """the logging frequency of the model performance (in terms of `updates`)"""
-    time_log_freq: int = 1000
-    """the logging frequency of the deck time statistics"""
+    time_log_freq: int = 0
+    """the logging frequency of the deck time statistics, 0 to disable"""
     save_interval: int = 400
     """the frequency of saving the model (in terms of `updates`)"""
     checkpoint: Optional[str] = None
@@ -360,7 +360,7 @@ def rollout(
         if args.concurrency:
             if update != 2:
                 params = params_queue.get()
-                params["params"]["Encoder_0"]['Embed_0']["embedding"].block_until_ready()
+                # params["params"]["Encoder_0"]['Embed_0']["embedding"].block_until_ready()
                 actor_policy_version += 1
         else:
             params = params_queue.get()
@@ -416,20 +416,21 @@ def rollout(
                             t.next_dones[idx] = True
                             t.rewards[idx] = -next_reward[idx]
                             break
-
-                for i in range(2):
-                    deck_time = info['step_time'][idx][i]
-                    deck_name = deck_names[info['deck'][idx][i]]
-
-                    time_count = deck_time_count[deck_name]
-                    avg_time = deck_avg_times[deck_name]
-                    avg_time = avg_time * (time_count / (time_count + 1)) + deck_time / (time_count + 1)
-                    max_time = max(deck_time, deck_max_times[deck_name])
-                    deck_avg_times[deck_name] = avg_time
-                    deck_max_times[deck_name] = max_time
-                    deck_time_count[deck_name] += 1
-                    if deck_time_count[deck_name] % args.time_log_freq == 0:
-                        print(f"Deck {deck_name}, avg: {avg_time * 1000:.2f}, max: {max_time * 1000:.2f}")
+                
+                if args.time_log_freq:
+                    for i in range(2):
+                        deck_time = info['step_time'][idx][i]
+                        deck_name = deck_names[info['deck'][idx][i]]
+
+                        time_count = deck_time_count[deck_name]
+                        avg_time = deck_avg_times[deck_name]
+                        avg_time = avg_time * (time_count / (time_count + 1)) + deck_time / (time_count + 1)
+                        max_time = max(deck_time, deck_max_times[deck_name])
+                        deck_avg_times[deck_name] = avg_time
+                        deck_max_times[deck_name] = max_time
+                        deck_time_count[deck_name] += 1
+                        if deck_time_count[deck_name] % args.time_log_freq == 0:
+                            print(f"Deck {deck_name}, avg: {avg_time * 1000:.2f}, max: {max_time * 1000:.2f}")
 
                 episode_reward = info['r'][idx] * (1 if cur_main else -1)
                 win = 1 if episode_reward > 0 else 0
@@ -474,14 +475,12 @@ def rollout(
         else:
             eval_stats = None
 
-        learn_opponent = False
         payload = (
             global_step,
             update,
             sharded_storage,
             *sharded_data,
             np.mean(params_queue_get_time),
-            learn_opponent,
             eval_stats,
         )
         rollout_queue.put(payload)
@@ -758,7 +757,6 @@ def main():
         sharded_next_inputs: List,
         sharded_next_main: List,
         key: jax.random.PRNGKey,
-        learn_opponent: bool = False,
     ):
         storage = jax.tree.map(lambda *x: jnp.hstack(x), *sharded_storages)
         # TODO: rstate will be out-date after the first update, maybe consider R2D2
@@ -862,7 +860,6 @@ def main():
         single_device_update,
         axis_name="local_devices",
         devices=global_learner_decices,
-        static_broadcasted_argnums=(7,),
     )
 
     params_queues = []
@@ -906,7 +903,6 @@ def main():
                     update,
                     *sharded_data,
                     avg_params_queue_get_time,
-                    learn_opponent,
                     eval_stats,
                 ) = rollout_queues[d_idx * args.num_actor_threads + thread_id].get()
                 sharded_data_list.append(sharded_data)
@@ -929,13 +925,12 @@ def main():
             agent_state,
             *list(zip(*sharded_data_list)),
             learner_keys,
-            learn_opponent,
         )
         unreplicated_params = flax.jax_utils.unreplicate(agent_state.params)
         params_queue_put_time = 0
         for d_idx, d_id in enumerate(args.actor_device_ids):
             device_params = jax.device_put(unreplicated_params, local_devices[d_id])
-            # device_params["params"]["Encoder_0"]['Embed_0']["embedding"].block_until_ready()
+            device_params["params"]["Encoder_0"]['Embed_0']["embedding"].block_until_ready()
             params_queue_put_start = time.time()
             for thread_id in range(args.num_actor_threads):
                 params_queues[d_idx * args.num_actor_threads + thread_id].put(device_params)