feature(pu/zt): add 2048 env and Stochastic MuZero

lzero/entry/eval_muzero.py

@@ -38,8 +38,8 @@ def eval_muzero(
         - policy (:obj:`Policy`): Converged policy.
     """
     cfg, create_cfg = input_cfg
-    assert create_cfg.policy.type in ['efficientzero', 'muzero', 'sampled_efficientzero'], \
-        "LightZero now only support the following algo.: 'efficientzero', 'muzero', 'sampled_efficientzero'"
+    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'"
 
     if cfg.policy.cuda and torch.cuda.is_available():
         cfg.policy.device = 'cuda'

lzero/entry/train_muzero.py

@@ -47,7 +47,7 @@ def train_muzero(
     """
 
     cfg, create_cfg = input_cfg
-    assert create_cfg.policy.type in ['efficientzero', 'muzero', 'sampled_efficientzero', 'gumbel_muzero'], \
+    assert create_cfg.policy.type in ['efficientzero', 'muzero', 'sampled_efficientzero', 'gumbel_muzero', 'stochastic_muzero'], \
         "train_muzero entry now only support the following algo.: 'efficientzero', 'muzero', 'sampled_efficientzero', 'gumbel_muzero'"
 
     if create_cfg.policy.type == 'muzero':
@@ -58,6 +58,8 @@ def train_muzero(
         from lzero.mcts import SampledEfficientZeroGameBuffer as GameBuffer
     elif create_cfg.policy.type == 'gumbel_muzero':
         from lzero.mcts import GumbelMuZeroGameBuffer as GameBuffer
+    elif create_cfg.policy.type == 'stochastic_muzero':
+        from lzero.mcts import StochasticMuZeroGameBuffer as GameBuffer
 
     if cfg.policy.cuda and torch.cuda.is_available():
         cfg.policy.device = 'cuda'

lzero/mcts/buffer/__init__.py

@@ -2,3 +2,4 @@
 from .game_buffer_efficientzero import EfficientZeroGameBuffer
 from .game_buffer_sampled_efficientzero import SampledEfficientZeroGameBuffer
 from .game_buffer_gumbel_muzero import GumbelMuZeroGameBuffer
+from .game_buffer_stochastic_muzero import StochasticMuZeroGameBuffer

lzero/mcts/buffer/game_segment.py

@@ -73,6 +73,9 @@ def __init__(self, action_space: int, game_segment_length: int = 200, config: Ea
 
         if self.config.sampled_algo:
             self.root_sampled_actions = []
+        if self.config.use_ture_chance_label_in_chance_encoder:
+            self.chance_segment = []
+
 
     def get_unroll_obs(self, timestep: int, num_unroll_steps: int = 0, padding: bool = False) -> np.ndarray:
         """
@@ -128,7 +131,8 @@ def append(
             obs: np.ndarray,
             reward: np.ndarray,
             action_mask: np.ndarray = None,
-            to_play: int = -1
+            to_play: int = -1,
+            chance: np.ndarray = 0,
     ) -> None:
         """
         Overview:
@@ -140,10 +144,12 @@ def append(
 
         self.action_mask_segment.append(action_mask)
         self.to_play_segment.append(to_play)
+        if self.config.use_ture_chance_label_in_chance_encoder:
+            self.chance_segment.append(chance)
 
     def pad_over(
             self, next_segment_observations: List, next_segment_rewards: List, next_segment_root_values: List,
-            next_segment_child_visits: List, next_segment_improved_policy: List = None
+            next_segment_child_visits: List, next_segment_improved_policy: List = None, next_chances: List = None,
     ) -> None:
         """
         Overview:
@@ -184,6 +190,9 @@ def pad_over(
         if self.config.gumbel_algo:
             for improved_policy in next_segment_improved_policy:
                 self.improved_policy_probs.append(improved_policy)
+        if self.config.use_ture_chance_label_in_chance_encoder:
+            for chances in next_chances:
+                self.chance_segment.append(chances)
 
     def get_targets(self, timestep: int) -> Tuple:
         """
@@ -253,6 +262,8 @@ def game_segment_to_array(self) -> None:
 
         self.action_mask_segment = np.array(self.action_mask_segment)
         self.to_play_segment = np.array(self.to_play_segment)
+        if self.config.use_ture_chance_label_in_chance_encoder:
+            self.chance_segment = np.array(self.chance_segment)
 
     def reset(self, init_observations: np.ndarray) -> None:
         """
@@ -271,6 +282,8 @@ def reset(self, init_observations: np.ndarray) -> None:
 
         self.action_mask_segment = []
         self.to_play_segment = []
+        if self.config.use_ture_chance_label_in_chance_encoder:
+            self.chance_segment = []
 
         assert len(init_observations) == self.frame_stack_num