feature(pu): add UniZero algo. and related configs/utils/envs/models

.gitignore

@@ -1421,7 +1421,7 @@ log*
 default*
 events.*
 
-# DI-engine special key
+# LightZero special key
 *default_logger.txt
 *default_tb_logger
 *evaluate.txt
@@ -1444,4 +1444,5 @@ events.*
 !/lzero/mcts/**/lib/*.h
 **/tb/*
 **/mcts/ctree/tests_cpp/*
-**/*tmp*
+**/*tmp*
+lzero/mcts/ctree/ctree_alphazero/pybind11

lzero/entry/__init__.py

@@ -1,8 +1,10 @@
-from .train_alphazero import train_alphazero
 from .eval_alphazero import eval_alphazero
-from .train_muzero import train_muzero
-from .train_muzero_with_reward_model import train_muzero_with_reward_model
 from .eval_muzero import eval_muzero
 from .eval_muzero_with_gym_env import eval_muzero_with_gym_env
+from .train_alphazero import train_alphazero
+from .train_muzero import train_muzero
+from .train_muzero_with_gym_env import train_muzero_with_gym_env
 from .train_muzero_with_gym_env import train_muzero_with_gym_env
+from .train_muzero_with_reward_model import train_muzero_with_reward_model
 from .train_rezero import train_rezero
+from .train_unizero import train_unizero

lzero/entry/eval_muzero.py

@@ -13,6 +13,7 @@
 from ding.utils import set_pkg_seed
 from ding.worker import BaseLearner
 from lzero.worker import MuZeroEvaluator
+from lzero.entry.utils import initialize_zeros_batch
 
 
 def eval_muzero(
@@ -25,7 +26,7 @@ def eval_muzero(
 ) -> 'Policy':  # noqa
     """
     Overview:
-        The eval entry for MCTS+RL algorithms, including MuZero, EfficientZero, Sampled EfficientZero.
+        The eval entry for MCTS+RL algorithms, including MuZero, EfficientZero, Sampled EfficientZero, StochasticMuZero, GumbelMuZero, UniZero, etc.
     Arguments:
         - input_cfg (:obj:`Tuple[dict, dict]`): Config in dict type.
             ``Tuple[dict, dict]`` type means [user_config, create_cfg].
@@ -38,8 +39,8 @@ def eval_muzero(
         - policy (:obj:`Policy`): Converged policy.
     """
     cfg, create_cfg = input_cfg
-    assert create_cfg.policy.type in ['efficientzero', 'muzero', 'stochastic_muzero', 'gumbel_muzero', 'sampled_efficientzero'], \
-        "LightZero now only support the following algo.: 'efficientzero', 'muzero', 'stochastic_muzero', 'gumbel_muzero', 'sampled_efficientzero'"
+    assert create_cfg.policy.type in ['efficientzero', 'muzero', 'muzero_context', 'muzero_rnn_full_obs', 'stochastic_muzero', 'gumbel_muzero', 'sampled_efficientzero', 'unizero'], \
+        "LightZero now only support the following algo.: 'efficientzero', 'muzero', 'muzero_context', 'muzero_rnn_full_obs', 'stochastic_muzero', 'gumbel_muzero', 'sampled_efficientzero', 'unizero'"
 
     if cfg.policy.cuda and torch.cuda.is_available():
         cfg.policy.device = 'cuda'
@@ -85,6 +86,13 @@ def eval_muzero(
     # Learner's before_run hook.
     learner.call_hook('before_run')
 
+    policy.last_batch_obs = initialize_zeros_batch(
+        cfg.policy.model.observation_shape,
+        len(evaluator_env_cfg),
+        cfg.policy.device
+    )
+    policy.last_batch_action = [-1 for _ in range(len(evaluator_env_cfg))]
+
     while True:
         # ==============================================================
         # eval trained model

lzero/entry/train_muzero.py

@@ -18,7 +18,7 @@
 from lzero.policy.random_policy import LightZeroRandomPolicy
 from lzero.worker import MuZeroCollector as Collector
 from lzero.worker import MuZeroEvaluator as Evaluator
-from .utils import random_collect
+from .utils import random_collect, initialize_zeros_batch
 
 
 def train_muzero(
@@ -47,10 +47,10 @@ def train_muzero(
     """
 
     cfg, create_cfg = input_cfg
-    assert create_cfg.policy.type in ['efficientzero', 'muzero', 'sampled_efficientzero', 'gumbel_muzero', 'stochastic_muzero'], \
+    assert create_cfg.policy.type in ['efficientzero', 'muzero', 'muzero_context', 'muzero_rnn_full_obs', 'sampled_efficientzero', 'gumbel_muzero', 'stochastic_muzero'], \
         "train_muzero entry now only support the following algo.: 'efficientzero', 'muzero', 'sampled_efficientzero', 'gumbel_muzero', 'stochastic_muzero'"
 
-    if create_cfg.policy.type == 'muzero':
+    if create_cfg.policy.type in ['muzero', 'muzero_context', 'muzero_rnn_full_obs']:
         from lzero.mcts import MuZeroGameBuffer as GameBuffer
     elif create_cfg.policy.type == 'efficientzero':
         from lzero.mcts import EfficientZeroGameBuffer as GameBuffer

lzero/entry/train_unizero.py

@@ -0,0 +1,207 @@
+import logging
+import os
+from functools import partial
+from typing import Tuple, Optional
+
+import torch
+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 set_pkg_seed, get_rank
+from ding.worker import BaseLearner
+from tensorboardX import SummaryWriter
+from torch.utils.tensorboard import SummaryWriter
+
+from lzero.entry.utils import log_buffer_memory_usage
+from lzero.policy import visit_count_temperature
+from lzero.policy.random_policy import LightZeroRandomPolicy
+from lzero.worker import MuZeroCollector as Collector
+from lzero.worker import MuZeroEvaluator as Evaluator
+from .utils import random_collect, initialize_zeros_batch
+
+
+def train_unizero(
+        input_cfg: Tuple[dict, dict],
+        seed: int = 0,
+        model: Optional[torch.nn.Module] = None,
+        model_path: Optional[str] = None,
+        max_train_iter: Optional[int] = int(1e10),
+        max_env_step: Optional[int] = int(1e10),
+) -> 'Policy':
+    """
+    Overview:
+        The train entry for UniZero, 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.
+    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.
+    Returns:
+        - policy (:obj:`Policy`): Converged policy.
+    """
+
+    cfg, create_cfg = input_cfg
+
+    # Ensure the specified policy type is supported
+    assert create_cfg.policy.type in ['unizero'], "train_unizero entry now only supports the following algo.: 'unizero'"
+
+    # Import the correct GameBuffer class based on the policy type
+    game_buffer_classes = {'unizero': 'UniZeroGameBuffer'}
+
+    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
+    cfg.policy.device = cfg.policy.model.world_model.device if torch.cuda.is_available() else 'cpu'
+    logging.info(f'cfg.policy.device: {cfg.policy.device}')
+
+    # Compile the configuration
+    cfg = compile_config(cfg, seed=seed, env=None, auto=True, create_cfg=create_cfg, save_cfg=True)
+
+    # Create main components: env, policy
+    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])
+
+    collector_env.seed(cfg.seed)
+    evaluator_env.seed(cfg.seed, dynamic_seed=False)
+    set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda)
+
+    policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'collect', 'eval'])
+
+    # Load pretrained model if specified
+    if model_path is not None:
+        logging.info(f'Loading model from {model_path} begin...')
+        policy.learn_mode.load_state_dict(torch.load(model_path, map_location=cfg.policy.device))
+        logging.info(f'Loading model from {model_path} end!')
+
+    # Create worker components: learner, collector, evaluator, replay buffer, commander
+    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)
+    collector = Collector(env=collector_env, policy=policy.collect_mode, tb_logger=tb_logger, exp_name=cfg.exp_name,
+                          policy_config=policy_config)
+    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)
+
+    # Learner's before_run hook
+    learner.call_hook('before_run')
+
+    # Collect random data before training
+    if cfg.policy.random_collect_episode_num > 0:
+        random_collect(cfg.policy, policy, LightZeroRandomPolicy, collector, collector_env, replay_buffer)
+
+    policy.last_batch_obs = initialize_zeros_batch(cfg.policy.model.observation_shape, len(evaluator_env_cfg),
+                                                   cfg.policy.device)
+    policy.last_batch_action = [-1 for _ in range(len(evaluator_env_cfg))]
+    batch_size = policy._cfg.batch_size
+
+    # TODO: for visualize
+    stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep)
+
+    while True:
+        # Log buffer memory usage
+        log_buffer_memory_usage(learner.train_iter, replay_buffer, tb_logger)
+
+        # Set temperature for visit count distributions
+        collect_kwargs = {
+            'temperature': visit_count_temperature(
+                policy_config.manual_temperature_decay,
+                policy_config.fixed_temperature_value,
+                policy_config.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:
+            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
+            )
+            collect_kwargs['epsilon'] = epsilon_greedy_fn(collector.envstep)
+
+        # Evaluate policy performance
+        if evaluator.should_eval(learner.train_iter):
+            policy.last_batch_obs = initialize_zeros_batch(
+                cfg.policy.model.observation_shape, len(evaluator_env_cfg), cfg.policy.device
+            )
+            policy.last_batch_action = [-1] * len(evaluator_env_cfg)
+            stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep)
+            if stop:
+                break
+
+        # Collect new data
+        policy.last_batch_obs = initialize_zeros_batch(
+            cfg.policy.model.observation_shape, len(collector_env_cfg), cfg.policy.device
+        )
+        policy.last_batch_action = [-1] * len(collector_env_cfg)
+        new_data = collector.collect(train_iter=learner.train_iter, policy_kwargs=collect_kwargs)
+
+        # Determine updates per collection
+        update_per_collect = cfg.policy.update_per_collect
+        if update_per_collect is None:
+            collected_transitions_num = sum(len(game_segment) for game_segment in new_data[0])
+            update_per_collect = int(collected_transitions_num * cfg.policy.model_update_ratio)
+
+        # 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 collector.envstep > cfg.policy.train_start_after_envsteps:
+            # data_sufficient = replay_buffer.get_num_of_game_segments() > batch_size  # For 'episode' sample type
+            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'batch_size: {batch_size}, replay_buffer: {replay_buffer}. Continue to collect now ....'
+                )
+                continue
+
+            for i in range(update_per_collect):
+                if data_sufficient:
+                    train_data = replay_buffer.sample(batch_size, policy)
+                    if cfg.policy.reanalyze_ratio > 0 and i % 20 == 0:
+                        # Clear caches
+                        for model in [policy._collect_model, policy._target_model]:
+                            # model.world_model.precompute_pos_emb_diff_kv()  # TODO
+                            model.world_model.clear_caches()
+
+                        torch.cuda.empty_cache()
+
+                    train_data.append({'train_which_component': 'transformer'})
+                    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'])
+
+        # Clear caches and precompute positional embedding matrices
+        for model in [policy._collect_model, policy._target_model]:
+            model.world_model.precompute_pos_emb_diff_kv()
+            model.world_model.clear_caches()
+
+        torch.cuda.empty_cache()
+
+        # Check stopping criteria
+        if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter:
+            break
+
+    learner.call_hook('after_run')
+    return policy

lzero/entry/utils.py

@@ -4,7 +4,29 @@
 import psutil
 from pympler.asizeof import asizeof
 from tensorboardX import SummaryWriter
+from typing import Optional, Callable
+import torch
+
 
+def initialize_zeros_batch(observation_shape, batch_size, device):
+    """
+    Overview:
+        Initialize a zeros tensor for batch observations based on the shape. This function is used to initialize the UniZero model input.
+    Arguments:
+        - observation_shape (:obj:`Union[int, List[int]]`): The shape of the observation tensor.
+        - batch_size (:obj:`int`): The batch size.
+        - device (:obj:`str`): The device to store the tensor.
+    Returns:
+        - zeros (:obj:`torch.Tensor`): The zeros tensor.
+    """
+    if isinstance(observation_shape, list):
+        shape = [batch_size, *observation_shape]
+    elif isinstance(observation_shape, int):
+        shape = [batch_size, observation_shape]
+    else:
+        raise TypeError("observation_shape must be either an int or a list")
+
+    return torch.zeros(shape).to(device)
 
 def random_collect(
         policy_cfg: 'EasyDict',  # noqa

lzero/mcts/buffer/__init__.py

@@ -1,4 +1,5 @@
 from .game_buffer_muzero import MuZeroGameBuffer
+from .game_buffer_unizero import UniZeroGameBuffer
 from .game_buffer_efficientzero import EfficientZeroGameBuffer
 from .game_buffer_sampled_efficientzero import SampledEfficientZeroGameBuffer
 from .game_buffer_gumbel_muzero import GumbelMuZeroGameBuffer