opendilab · puyuan1996 · Apr 29, 2025 · Apr 29, 2025 · Apr 29, 2025 · May 1, 2025
diff --git a/README.md b/README.md
@@ -28,7 +28,7 @@
 [![GitHub license](https://img.shields.io/github/license/opendilab/LightZero)](https://github.com/opendilab/LightZero/blob/master/LICENSE)
 [![discord badge](https://dcbadge.vercel.app/api/server/dkZS2JF56X?style=flat)](https://discord.gg/dkZS2JF56X)
 
-Updated on 2025.04.09 LightZero-v0.2.0
+Updated on 2025.06.03 LightZero-v0.2.0
 
 English | [简体中文(Simplified Chinese)](https://github.com/opendilab/LightZero/blob/main/README.zh.md) | [Documentation](https://opendilab.github.io/LightZero) | [LightZero Paper](https://arxiv.org/abs/2310.08348) | [🔥UniZero Paper](https://arxiv.org/abs/2406.10667) | [🔥ReZero Paper](https://arxiv.org/abs/2404.16364)
 

diff --git a/README.zh.md b/README.zh.md
@@ -27,7 +27,7 @@
 [![Contributors](https://img.shields.io/github/contributors/opendilab/LightZero)](https://github.com/opendilab/LightZero/graphs/contributors)
 [![GitHub license](https://img.shields.io/github/license/opendilab/LightZero)](https://github.com/opendilab/LightZero/blob/master/LICENSE)
 
-最近更新于 2025.04.09 LightZero-v0.2.0
+最近更新于 2025.06.03 LightZero-v0.2.0
 
 [English](https://github.com/opendilab/LightZero/blob/main/README.md) | 简体中文 | [文档](https://opendilab.github.io/LightZero) | [LightZero 论文](https://arxiv.org/abs/2310.08348) | [🔥UniZero 论文](https://arxiv.org/abs/2406.10667) | [🔥ReZero 论文](https://arxiv.org/abs/2404.16364)
 

diff --git a/docs/source/tutorials/algos/customize_algos.md b/docs/source/tutorials/algos/customize_algos.md
@@ -119,16 +119,17 @@ Here is an example of unit testing in LightZero. In this example, we test the `i
 ```Python
 import pytest
 import torch
-from lzero.policy.scaling_transform import inverse_scalar_transform, InverseScalarTransform
+from lzero.policy.scaling_transform import DiscreteSupport, inverse_scalar_transform, InverseScalarTransform
 
 @pytest.mark.unittest
 def test_scaling_transform():
     import time
     logit = torch.randn(16, 601)
+    discrete_support = DiscreteSupport(-300., 301., 1.)
     start = time.time()
-    output_1 = inverse_scalar_transform(logit, 300)
+    output_1 = inverse_scalar_transform(logit, discrete_support)
     print('t1', time.time() - start)
-    handle = InverseScalarTransform(300)
+    handle = InverseScalarTransform(discrete_support)
     start = time.time()
     output_2 = handle(logit)
     print('t2', time.time() - start)

diff --git a/docs/source/tutorials/algos/customize_algos_zh.md b/docs/source/tutorials/algos/customize_algos_zh.md
@@ -120,16 +120,17 @@ if timestep.done:
 ```Python
 import pytest
 import torch
-from lzero.policy.scaling_transform import inverse_scalar_transform, InverseScalarTransform
+from lzero.policy.scaling_transform import DiscreteSupport, inverse_scalar_transform, InverseScalarTransform
 
 @pytest.mark.unittest
 def test_scaling_transform():
     import time
     logit = torch.randn(16, 601)
+    discrete_support = DiscreteSupport(-300., 301., 1.)
     start = time.time()
-    output_1 = inverse_scalar_transform(logit, 300)
+    output_1 = inverse_scalar_transform(logit, discrete_support)
     print('t1', time.time() - start)
-    handle = InverseScalarTransform(300)
+    handle = InverseScalarTransform(discrete_support)
     start = time.time()
     output_2 = handle(logit)
     print('t2', time.time() - start)

diff --git a/docs/source/tutorials/config/config.md b/docs/source/tutorials/config/config.md
@@ -44,7 +44,8 @@ The `main_config` dictionary contains the main parameter settings for running th
     - `downsample`: Whether to downsample the input.
     - `norm_type`: The type of normalization used.
     - `num_channels`: The number of channels in the convolutional layers (number of features extracted).
-    - `support_scale`: The range of the value support set (`-support_scale` to `support_scale`).
+    - `reward_support_range`: The range of the reward support set (`(start, stop, step)`).
+    - `value_support_range`: The range of the value support set (`(start, stop, step)`).
     - `bias`: Whether to use bias terms in the layers.
     - `discrete_action_encoding_type`: How discrete actions are encoded.
     - `self_supervised_learning_loss`: Whether to use a self-supervised learning loss (as in EfficientZero).

diff --git a/docs/source/tutorials/config/config_zh.md b/docs/source/tutorials/config/config_zh.md
@@ -43,7 +43,8 @@
     - `downsample`: 是否进行降采样。
     - `norm_type`: 归一化使用的方法。
     - `num_channels`: 卷积层提取的特征个数。
-    - `support_scale`: 价值支持集的范围 (-support_scale, support_scale)。
+    - `reward_support_range`: 价值支持集的范围 (`(start, stop, step)`)。<!-- TODO : ADAPT THIS DESCRIPTION, I DON'T SPEAK CHINESE -->
+    - `value_support_range`: 价值支持集的范围 (`(start, stop, step)`)。<!-- TODO : ADAPT THIS DESCRIPTION, I DON'T SPEAK CHINESE -->
     - `bias`: 是否使用偏置。
     - `discrete_action_encoding_type`: 离散化动作空间使用的编码类型。
     - `self_supervised_learning_loss`: 是否使用自监督学习损失（参照EfficientZero的实现）。

diff --git a/lzero/agent/config/gumbel_muzero/gomoku_play_with_bot.py b/lzero/agent/config/gumbel_muzero/gomoku_play_with_bot.py
@@ -44,9 +44,8 @@
                 image_channel=3,
                 num_res_blocks=1,
                 num_channels=32,
-                support_scale=10,
-                reward_support_size=21,
-                value_support_size=21,
+                reward_support_range=(-10., 11., 1.),
+                value_support_range=(-10., 11., 1.),
             ),
             cuda=True,
             env_type='board_games',

diff --git a/lzero/agent/config/gumbel_muzero/tictactoe_play_with_bot.py b/lzero/agent/config/gumbel_muzero/tictactoe_play_with_bot.py
@@ -38,9 +38,8 @@
                 reward_head_hidden_channels=[8],
                 value_head_hidden_channels=[8],
                 policy_head_hidden_channels=[8],
-                support_scale=10,
-                reward_support_size=21,
-                value_support_size=21,
+                reward_support_range=(-10., 11., 1.),
+                value_support_range=(-10., 11., 1.),
             ),
             cuda=True,
             env_type='board_games',

diff --git a/lzero/agent/config/muzero/gomoku_play_with_bot.py b/lzero/agent/config/muzero/gomoku_play_with_bot.py
@@ -44,9 +44,8 @@
                 image_channel=3,
                 num_res_blocks=1,
                 num_channels=32,
-                support_scale=10,
-                reward_support_size=21,
-                value_support_size=21,
+                reward_support_range=(-10., 11., 1.),
+                value_support_range=(-10., 11., 1.),
             ),
             cuda=True,
             env_type='board_games',

diff --git a/lzero/agent/config/muzero/tictactoe_play_with_bot.py b/lzero/agent/config/muzero/tictactoe_play_with_bot.py
@@ -38,9 +38,8 @@
                 reward_head_hidden_channels=[8],
                 value_head_hidden_channels=[8],
                 policy_head_hidden_channels=[8],
-                support_scale=10,
-                reward_support_size=21,
-                value_support_size=21,
+                reward_support_range=(-10., 11., 1.),
+                value_support_range=(-10., 11., 1.),
                 norm_type='BN', 
             ),
             cuda=True,

diff --git a/lzero/entry/__init__.py b/lzero/entry/__init__.py
@@ -1,5 +1,6 @@
 from .eval_alphazero import eval_alphazero
 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
@@ -12,4 +13,5 @@
 from .train_muzero_multitask_segment_ddp import train_muzero_multitask_segment_ddp
 from .train_unizero_multitask_segment_ddp import train_unizero_multitask_segment_ddp
 from .train_unizero_multitask_segment_eval import train_unizero_multitask_segment_eval
-from .utils import *
+from .train_unizero_multitask_balance_segment_ddp import train_unizero_multitask_balance_segment_ddp
+from .utils import *
diff --git a/lzero/entry/compute_task_weight.py b/lzero/entry/compute_task_weight.py
diff --git a/lzero/entry/eval_muzero.py b/lzero/entry/eval_muzero.py
@@ -1,6 +1,7 @@
 import os
 from functools import partial
 from typing import Optional, Tuple
+import logging
 
 import numpy as np
 import torch
@@ -51,7 +52,7 @@ def eval_muzero(
     # Create main components: env, policy
     env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env)
     evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg])
-
+    # print(f"cfg.seed:{cfg.seed}")
     evaluator_env.seed(cfg.seed, dynamic_seed=False)
     set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda)