feature(whl/pu): add jericho environments and related configs (#307)

* init commit

* add bert encoding

* debug

* debug and polish env

* polish config

* update

* polish(pu): polish jericho config and pipeline

* polish(pu): polish HFLanguageRepresentationNetwork

* fix(pu): fix padded action_mask

* sync code

* feature(pu): vllm server for embedding model inference

* fix(pu): fix update_per_collect in ddp train

* fix(pu): fix train_entry_time to the same in ddp train

* polish(pu): polish jericho uz config

* feature(pu): add remove_stuck_actions option

* polish(pu): set projection layer train after language embed model

* polish(pu): polish jericho_env and add jericho_ppo_config

* sync code

* fix(pu): fix gradient accumulation_steps option

* fix(pu): fix action_mask all-zero bug

* fix(pu): fix lr target_model_update bug when  accumulation_steps>1

* polish(pu): polish jericho env and its unizero/ppo config

* polish(pu): polish jericho unizero config

* polish(pu): polish jericho unizero config

* polish(pu): polish comments

* polish(pu): polish jericho configs

* polish(pu): polish jericho configs

---------

Co-authored-by: ‘whl’ <‘[email protected]’>
Co-authored-by: dyyoungg <[email protected]>
Co-authored-by: puyuan <[email protected]>
Co-authored-by: PaParaZz1 <[email protected]>

Author: 5 people

Diff for: docs/source/_templates/layout.html renamed to docs/source/_templates/layout.html

@@ -21,6 +21,8 @@
 from lzero.worker import MuZeroEvaluator as Evaluator
 from lzero.worker import MuZeroCollector as Collector
 from .utils import random_collect, calculate_update_per_collect
+import torch.distributed as dist
+from ding.utils import set_pkg_seed, get_rank, get_world_size
 
 
 def train_unizero(
@@ -33,167 +35,201 @@ def train_unizero(
 ) -> 'Policy':
     """
     Overview:
-        The train entry for UniZero, proposed in our paper UniZero: Generalized and Efficient Planning with Scalable Latent World Models.
+        This function serves as the training entry point for UniZero, as proposed in our paper "UniZero: Generalized and Efficient Planning with Scalable Latent World Models".
         UniZero aims to enhance the planning capabilities of reinforcement learning agents by addressing the limitations found in MuZero-style algorithms,
-        particularly in environments requiring the capture of long-term dependencies. More details can be found in https://arxiv.org/abs/2406.10667.
+        particularly in environments that require capturing long-term dependencies. More details can be found in https://arxiv.org/abs/2406.10667.
+    
     Arguments:
-        - input_cfg (:obj:`Tuple[dict, dict]`): Config in dict type.
-            ``Tuple[dict, dict]`` type means [user_config, create_cfg].
-        - seed (:obj:`int`): Random seed.
-        - model (:obj:`Optional[torch.nn.Module]`): Instance of torch.nn.Module.
-        - model_path (:obj:`Optional[str]`): The pretrained model path, which should
-            point to the ckpt file of the pretrained model, and an absolute path is recommended.
-            In LightZero, the path is usually something like ``exp_name/ckpt/ckpt_best.pth.tar``.
-        - max_train_iter (:obj:`Optional[int]`): Maximum policy update iterations in training.
-        - max_env_step (:obj:`Optional[int]`): Maximum collected environment interaction steps.
+        - input_cfg (:obj:`Tuple[dict, dict]`): Configuration in dictionary format.
+            ``Tuple[dict, dict]`` indicates [user_config, create_cfg].
+        - seed (:obj:`int`): Random seed for reproducibility.
+        - model (:obj:`Optional[torch.nn.Module]`): Instance of a PyTorch model.
+        - model_path (:obj:`Optional[str]`): Path to the pretrained model, which should
+            point to the checkpoint file of the pretrained model. An absolute path is recommended.
+            In LightZero, the path typically resembles ``exp_name/ckpt/ckpt_best.pth.tar``.
+        - max_train_iter (:obj:`Optional[int]`): Maximum number of policy update iterations during training.
+        - max_env_step (:obj:`Optional[int]`): Maximum number of environment interaction steps to collect.
+    
     Returns:
-        - policy (:obj:`Policy`): Converged policy.
+        - policy (:obj:`Policy`): The converged policy after training.
     """
 
     cfg, create_cfg = input_cfg
 
     # Ensure the specified policy type is supported
-    assert create_cfg.policy.type in ['unizero', 'sampled_unizero'], "train_unizero entry now only supports the following algo.: 'unizero', 'sampled_unizero'"
+    assert create_cfg.policy.type in ['unizero', 'sampled_unizero'], "train_unizero only supports the following algorithms: 'unizero', 'sampled_unizero'"
+    logging.info(f"Using policy type: {create_cfg.policy.type}")
 
-    # Import the correct GameBuffer class based on the policy type
+    # Import the appropriate GameBuffer class based on the policy type
     game_buffer_classes = {'unizero': 'UniZeroGameBuffer', 'sampled_unizero': 'SampledUniZeroGameBuffer'}
-
     GameBuffer = getattr(__import__('lzero.mcts', fromlist=[game_buffer_classes[create_cfg.policy.type]]),
                          game_buffer_classes[create_cfg.policy.type])
 
-    # Set device based on CUDA availability
+    # Check for GPU availability and set the device accordingly
     cfg.policy.device = cfg.policy.model.world_model_cfg.device if torch.cuda.is_available() else 'cpu'
-    logging.info(f'cfg.policy.device: {cfg.policy.device}')
+    logging.info(f"Device set to: {cfg.policy.device}")
 
-    # Compile the configuration
+    # Compile the configuration file
     cfg = compile_config(cfg, seed=seed, env=None, auto=True, create_cfg=create_cfg, save_cfg=True)
 
-    # Create main components: env, policy
+    # Create environment manager
     env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env)
     collector_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in collector_env_cfg])
     evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg])
 
+    # Initialize environment and random seed
     collector_env.seed(cfg.seed)
     evaluator_env.seed(cfg.seed, dynamic_seed=False)
     set_pkg_seed(cfg.seed, use_cuda=torch.cuda.is_available())
 
+    # Initialize wandb if specified
     if cfg.policy.use_wandb:
-        # Initialize wandb
+        logging.info("Initializing wandb...")
         wandb.init(
             project="LightZero",
             config=cfg,
             sync_tensorboard=False,
             monitor_gym=False,
             save_code=True,
         )
+        logging.info("wandb initialization completed!")
 
+    # Create policy
+    logging.info("Creating policy...")
     policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'collect', 'eval'])
+    logging.info("Policy created successfully!")
 
     # Load pretrained model if specified
     if model_path is not None:
-        logging.info(f'Loading model from {model_path} begin...')
+        logging.info(f"Loading pretrained model from {model_path}...")
         policy.learn_mode.load_state_dict(torch.load(model_path, map_location=cfg.policy.device))
-        logging.info(f'Loading model from {model_path} end!')
+        logging.info("Pretrained model loaded successfully!")
 
-    # Create worker components: learner, collector, evaluator, replay buffer, commander
+    # Create core components for training
     tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) if get_rank() == 0 else None
     learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name)
-
-    # MCTS+RL algorithms related core code
-    policy_config = cfg.policy
-    replay_buffer = GameBuffer(policy_config)
+    replay_buffer = GameBuffer(cfg.policy)
     collector = Collector(env=collector_env, policy=policy.collect_mode, tb_logger=tb_logger, exp_name=cfg.exp_name,
-                          policy_config=policy_config)
+                          policy_config=cfg.policy)
     evaluator = Evaluator(eval_freq=cfg.policy.eval_freq, n_evaluator_episode=cfg.env.n_evaluator_episode,
                           stop_value=cfg.env.stop_value, env=evaluator_env, policy=policy.eval_mode,
-                          tb_logger=tb_logger, exp_name=cfg.exp_name, policy_config=policy_config)
+                          tb_logger=tb_logger, exp_name=cfg.exp_name, policy_config=cfg.policy)
 
-    # Learner's before_run hook
+    # Execute the learner's before_run hook
     learner.call_hook('before_run')
-    if policy_config.use_wandb:
+
+    if cfg.policy.use_wandb:
         policy.set_train_iter_env_step(learner.train_iter, collector.envstep)
 
-    # Collect random data before training
+    # Randomly collect data if specified
     if cfg.policy.random_collect_episode_num > 0:
+        logging.info("Collecting random data...")
         random_collect(cfg.policy, policy, LightZeroRandomPolicy, collector, collector_env, replay_buffer)
+        logging.info("Random data collection completed!")
 
     batch_size = policy._cfg.batch_size
 
+    if cfg.policy.multi_gpu:
+        # Get current world size and rank
+        world_size = get_world_size()
+        rank = get_rank()
+    else:
+        world_size = 1
+        rank = 0
+
     while True:
-        # Log buffer memory usage
+        # Log memory usage of the replay buffer
         log_buffer_memory_usage(learner.train_iter, replay_buffer, tb_logger)
 
-        # Set temperature for visit count distributions
+        # Set temperature parameter for data collection
         collect_kwargs = {
             'temperature': visit_count_temperature(
-                policy_config.manual_temperature_decay,
-                policy_config.fixed_temperature_value,
-                policy_config.threshold_training_steps_for_final_temperature,
+                cfg.policy.manual_temperature_decay,
+                cfg.policy.fixed_temperature_value,
+                cfg.policy.threshold_training_steps_for_final_temperature,
                 trained_steps=learner.train_iter
             ),
             'epsilon': 0.0  # Default epsilon value
         }
 
-        # Configure epsilon for epsilon-greedy exploration
-        if policy_config.eps.eps_greedy_exploration_in_collect:
+        # Configure epsilon-greedy exploration
+        if cfg.policy.eps.eps_greedy_exploration_in_collect:
             epsilon_greedy_fn = get_epsilon_greedy_fn(
-                start=policy_config.eps.start,
-                end=policy_config.eps.end,
-                decay=policy_config.eps.decay,
-                type_=policy_config.eps.type
+                start=cfg.policy.eps.start,
+                end=cfg.policy.eps.end,
+                decay=cfg.policy.eps.decay,
+                type_=cfg.policy.eps.type
             )
             collect_kwargs['epsilon'] = epsilon_greedy_fn(collector.envstep)
 
         # Evaluate policy performance
         if evaluator.should_eval(learner.train_iter):
+            logging.info(f"Training iteration {learner.train_iter}: Starting evaluation...")
             stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep)
+            logging.info(f"Training iteration {learner.train_iter}: Evaluation completed, stop condition: {stop}, current reward: {reward}")
             if stop:
+                logging.info("Stopping condition met, training ends!")
                 break
 
         # Collect new data
         new_data = collector.collect(train_iter=learner.train_iter, policy_kwargs=collect_kwargs)
+        logging.info(f"Rank {rank}, Training iteration {learner.train_iter}: New data collection completed!")
 
         # Determine updates per collection
-        update_per_collect = calculate_update_per_collect(cfg, new_data)
+        update_per_collect = cfg.policy.update_per_collect
+        if update_per_collect is None:
+            update_per_collect = calculate_update_per_collect(cfg, new_data, world_size)
 
         # Update replay buffer
         replay_buffer.push_game_segments(new_data)
         replay_buffer.remove_oldest_data_to_fit()
 
-        # Train the policy if sufficient data is available
+        if world_size > 1:
+            # Synchronize all ranks before training
+            try:
+                dist.barrier()
+            except Exception as e:
+                logging.error(f'Rank {rank}: Synchronization barrier failed, error: {e}')
+                break
+
+        # Check if there is sufficient data for training
         if collector.envstep > cfg.policy.train_start_after_envsteps:
             if cfg.policy.sample_type == 'episode':
                 data_sufficient = replay_buffer.get_num_of_game_segments() > batch_size
             else:
                 data_sufficient = replay_buffer.get_num_of_transitions() > batch_size
+            
             if not data_sufficient:
                 logging.warning(
-                    f'The data in replay_buffer is not sufficient to sample a mini-batch: '
+                    f'Rank {rank}: The data in replay_buffer is not sufficient to sample a mini-batch: '
                     f'batch_size: {batch_size}, replay_buffer: {replay_buffer}. Continue to collect now ....'
                 )
                 continue
 
+            # Execute multiple training rounds
             for i in range(update_per_collect):
                 train_data = replay_buffer.sample(batch_size, policy)
                 if cfg.policy.reanalyze_ratio > 0 and i % 20 == 0:
-                    # Clear caches and precompute positional embedding matrices
-                    policy.recompute_pos_emb_diff_and_clear_cache()  # TODO
-
-                if policy_config.use_wandb:
+                    policy.recompute_pos_emb_diff_and_clear_cache()
+                
+                if cfg.policy.use_wandb:
                     policy.set_train_iter_env_step(learner.train_iter, collector.envstep)
 
-                train_data.append({'train_which_component': 'transformer'})
-                log_vars = learner.train(train_data, collector.envstep)
+                train_data.append(learner.train_iter)
 
+                log_vars = learner.train(train_data, collector.envstep)
                 if cfg.policy.use_priority:
                     replay_buffer.update_priority(train_data, log_vars[0]['value_priority_orig'])
 
         policy.recompute_pos_emb_diff_and_clear_cache()
 
         # Check stopping criteria
         if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter:
+            logging.info("Stopping condition met, training ends!")
             break
 
     learner.call_hook('after_run')
-    wandb.finish()
-    return policy
+    if cfg.policy.use_wandb:
+        wandb.finish()
+    logging.info("===== Training Completed =====")
+    return policy

Diff for: lzero/entry/train_unizero.py

@@ -4,13 +4,14 @@
 from typing import Tuple, Optional
 
 import torch
+import wandb
 from ding.config import compile_config
 from ding.envs import create_env_manager
 from ding.envs import get_vec_env_setting
 from ding.policy import create_policy
 from ding.rl_utils import get_epsilon_greedy_fn
 from ding.utils import EasyTimer
-from ding.utils import set_pkg_seed, get_rank
+from ding.utils import set_pkg_seed, get_rank, get_world_size
 from ding.worker import BaseLearner
 from tensorboardX import SummaryWriter
 from torch.utils.tensorboard import SummaryWriter
@@ -103,6 +104,9 @@ def train_unizero_segment(
     # Learner's before_run hook
     learner.call_hook('before_run')
 
+    if cfg.policy.use_wandb:
+        policy.set_train_iter_env_step(learner.train_iter, collector.envstep)
+
     # Collect random data before training
     if cfg.policy.random_collect_episode_num > 0:
         random_collect(cfg.policy, policy, LightZeroRandomPolicy, collector, collector_env, replay_buffer)
@@ -116,6 +120,14 @@ def train_unizero_segment(
     train_epoch = 0
     reanalyze_batch_size = cfg.policy.reanalyze_batch_size
 
+    if cfg.policy.multi_gpu:
+        # Get current world size and rank
+        world_size = get_world_size()
+        rank = get_rank()
+    else:
+        world_size = 1
+        rank = 0
+
     while True:
         # Log buffer memory usage
         log_buffer_memory_usage(learner.train_iter, replay_buffer, tb_logger)
@@ -151,7 +163,7 @@ def train_unizero_segment(
         new_data = collector.collect(train_iter=learner.train_iter, policy_kwargs=collect_kwargs)
 
         # Determine updates per collection
-        update_per_collect = calculate_update_per_collect(cfg, new_data)
+        update_per_collect = calculate_update_per_collect(cfg, new_data, world_size)
 
         # Update replay buffer
         replay_buffer.push_game_segments(new_data)
@@ -196,8 +208,10 @@ def train_unizero_segment(
                         logging.info(f'Buffer reanalyze time: {timer.value}')
 
                 train_data = replay_buffer.sample(batch_size, policy)
+                if cfg.policy.use_wandb:
+                    policy.set_train_iter_env_step(learner.train_iter, collector.envstep)
 
-                train_data.append({'train_which_component': 'transformer'})
+                train_data.append(learner.train_iter)
                 log_vars = learner.train(train_data, collector.envstep)
 
                 if cfg.policy.use_priority:
@@ -211,4 +225,6 @@ def train_unizero_segment(
             break
 
     learner.call_hook('after_run')
+    if cfg.policy.use_wandb:
+        wandb.finish()
     return policy

Diff for: lzero/entry/train_unizero_segment.py

@@ -26,7 +26,7 @@ def ddp_synchronize():
     if is_ddp_enabled():
         dist.barrier()
 
-def ddp_all_reduce_sum(tensor):
+def ddp_all_reduce_sum(tensor: torch.Tensor) -> torch.Tensor:
     """
     Perform an all-reduce operation (sum) on the given tensor across
     all processes in DDP mode. Returns the reduced tensor.
@@ -41,15 +41,16 @@ def ddp_all_reduce_sum(tensor):
         dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
     return tensor
 
-def calculate_update_per_collect(cfg, new_data):
+def calculate_update_per_collect(cfg: 'EasyDict', new_data: List[List[torch.Tensor]], world_size: int = 1) -> int:
     """
     Calculate the number of updates to perform per data collection in a
     Distributed Data Parallel (DDP) setting. This ensures that all GPUs
     compute the same `update_per_collect` value, synchronized across processes.
 
     Arguments:
         - cfg: Configuration object containing policy settings.
-        - new_data (list): The newly collected data segments.
+        - new_data (List[List[torch.Tensor]]): The newly collected data segments.
+        - world_size (int): The total number of processes.
 
     Returns:
         - int: The number of updates to perform per collection.
@@ -68,7 +69,7 @@ def calculate_update_per_collect(cfg, new_data):
             for game_segment in new_data[0]
         )
 
-        if torch.cuda.is_available():
+        if torch.cuda.is_available() and world_size > 1:
             # Convert the collected transitions count to a GPU tensor for DDP operations.
             collected_transitions_tensor = torch.tensor(
                 collected_transitions_num, dtype=torch.int64, device='cuda'

Diff for: lzero/entry/utils.py

@@ -711,7 +711,7 @@ def _compute_target_policy_non_reanalyzed(
             ]
         else:
             legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
-
+            
         with torch.no_grad():
             policy_index = 0
             # 0 -> Invalid target policy for padding outside of game segments,
@@ -730,7 +730,7 @@ def _compute_target_policy_non_reanalyzed(
                             # for atari/classic_control/box2d environments that only have one player.
                             target_policies.append(distributions)
                         else:
-                            # for board games that have two players.
+                            # for board games that have two players or envs that have varied action space.
                             policy_tmp = [0 for _ in range(policy_shape)]
                             for index, legal_action in enumerate(legal_actions[policy_index]):
                                 # only the action in ``legal_action`` the policy logits is nonzero