diff --git a/lzero/entry/train_unizero_multitask_segment_ddp.py b/lzero/entry/train_unizero_multitask_segment_ddp.py
index 0eb8d9606..240f2c67e 100644
--- a/lzero/entry/train_unizero_multitask_segment_ddp.py
+++ b/lzero/entry/train_unizero_multitask_segment_ddp.py
@@ -26,6 +26,11 @@
     symlog,
     inv_symlog,
 )
+from lzero.entry.utils import (
+    collect_and_log_moe_statistics,
+    TemperatureScheduler,
+    log_buffer_memory_usage
+)
 # NOTE: The following imports are for type hinting purposes.
 # The actual GameBuffer is selected dynamically based on the policy type.
 from lzero.mcts import UniZeroGameBuffer
@@ -52,7 +57,8 @@ def train_unizero_multitask_segment_ddp(
         model_path: Optional[str] = None,
         max_train_iter: Optional[int] = int(1e10),
         max_env_step: Optional[int] = int(1e10),
-        benchmark_name: str = "atari"
+        benchmark_name: str = "atari",
+        cal_moe_profile: bool = True
 ) -> 'Policy':
     """
     Overview:
@@ -68,6 +74,7 @@ def train_unizero_multitask_segment_ddp(
         - max_train_iter (:obj:`Optional[int]`): The maximum number of policy update iterations during training.
         - max_env_step (:obj:`Optional[int]`): The maximum number of environment interaction steps to collect.
         - benchmark_name (:obj:`str`): The name of the benchmark, e.g., "atari" or "dmc".
+        - cal_moe_profile (:obj:`bool`, optional): Whether to enable MoE expert selection statistics and heatmap logging. Default is True.
 
     Returns:
         - policy (:obj:`Policy`): The converged policy.
@@ -201,6 +208,8 @@ def train_unizero_multitask_segment_ddp(
         # Create a TensorBoard logger.
         log_dir = os.path.join('./{}/log'.format(cfg.exp_name), f'serial_rank_{rank}')
         tb_logger = SummaryWriter(log_dir)
+        # Inject TensorBoard logger into policy for gradient conflict and MoE statistics logging
+        policy.logger = tb_logger
 
         # Create the shared learner.
         learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name)
@@ -498,6 +507,12 @@ def train_unizero_multitask_segment_ddp(
                     # DDP automatically synchronizes gradients and parameters during training.
                     log_vars = learner.train(train_data_multi_task, envstep_multi_task, policy_kwargs=learn_kwargs)
 
+                    if cal_moe_profile and cfg.policy.model.world_model_cfg.multiplication_moe_in_transformer and cfg.policy.model.world_model_cfg.num_experts_of_moe_in_transformer:
+                        # Collect and log MoE expert selection statistics (heatmaps, distribution divergence)
+                        moe_log_interval = getattr(cfg.policy, 'moe_log_interval', 1)
+                        if learner.train_iter % moe_log_interval == 0:
+                            collect_and_log_moe_statistics(policy, tb_logger, learner.train_iter, world_size, rank)   
+            
                     # Check if task_exploitation_weight needs to be calculated.
                     if i == 0:
                         # Calculate task weights.
diff --git a/lzero/entry/utils.py b/lzero/entry/utils.py
index dc8dacf0f..e6e0705a6 100644
--- a/lzero/entry/utils.py
+++ b/lzero/entry/utils.py
@@ -35,6 +35,14 @@
     - `log_buffer_memory_usage` - Log buffer memory usage
     - `log_buffer_run_time` - Log buffer runtime
 
+  - **MoE Statistics Utilities**:
+    - `merge_expert_stats_across_ranks` - Merge expert selection stats from distributed ranks
+    - `create_heatmap_with_values_fast` - Fast Task-Expert heatmap generation
+    - `collect_and_log_moe_statistics` - End-to-end MoE stats collection and TensorBoard logging
+    - `jensen_shannon_divergence_batch_gpu` - GPU batch JS divergence for task distributions
+    - `wasserstein_distance_batch_gpu` - GPU batch Wasserstein distance
+    - `compute_distribution_divergences_optimized` - Optimized inter-task distribution divergence
+
 - **`__init__.py`** - Package initialization file
   - Exports all training and evaluation entry functions
   - Exports commonly used functions from utility modules
@@ -61,7 +69,14 @@
 import torch.nn.functional as F
 from pympler.asizeof import asizeof
 from tensorboardX import SummaryWriter
-
+import time
+from typing import Optional, Callable, Union, List, Tuple, Dict
+from io import BytesIO
+import concurrent.futures
+import seaborn as sns
+from PIL import Image
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
 # ==============================================================================
 # Placeholder Types for External Dependencies
 #
@@ -1044,3 +1059,990 @@ def __init__(self):
     print("\nParameter status after un-freezing:")
     log_module_trainable_status(model, "DummyModel", logging.getLogger())
     
+_GLOBAL_HEATMAP_FIG = None
+_GLOBAL_HEATMAP_AX = None
+
+
+def merge_expert_stats_across_ranks(all_expert_stats):
+    """
+    Overview:
+        Merge expert selection statistics data from all distributed training ranks.
+        Combines statistics from different GPU processes for comprehensive analysis.
+
+    Arguments:
+        - all_expert_stats (:obj:`list`): List of expert statistics from all ranks.
+            Each element is a dict: {task_id: {window_type: {frequencies, total_selections, data_points}}}.
+
+    Returns:
+        - merged_stats (:obj:`dict`): Merged statistics dictionary with structure
+            {task_id: {window_type: {frequencies, total_selections, data_points}}}.
+            Frequencies are converted to numpy arrays for serialization.
+
+    Notes:
+        Only processes statistics with total_selections > 0.
+
+    Examples:
+        >>> stats_list = [rank0_stats, rank1_stats, rank2_stats]
+        >>> merged = merge_expert_stats_across_ranks(stats_list)
+        >>> print(f"Merged {len(merged)} tasks")
+    """
+    merged_stats = {}  # {task_id: {window_type: stats}}
+
+    for rank_expert_stats in all_expert_stats:
+        if rank_expert_stats:
+            for task_id, task_stats in rank_expert_stats.items():
+                if task_id not in merged_stats:
+                    merged_stats[task_id] = {}
+
+                for window_type, stats in task_stats.items():
+                    # Only process statistics with actual data (tasks handled by current GPU)
+                    if stats and stats.get('total_selections', 0) > 0:
+                        merged_stats[task_id][window_type] = {
+                            'frequencies': np.array(stats['frequencies']),
+                            'total_selections': stats['total_selections'],
+                            'data_points': stats['data_points']
+                        }
+    return merged_stats
+
+
+def _get_or_create_heatmap_figure(figsize):
+    """
+    Overview:
+        Get or create a reusable heatmap figure for memory efficiency.
+        Maintains global figure cache to reduce memory allocation overhead.
+    Arguments:
+        - figsize (:obj:`tuple`): Figure size as (width, height).
+    Returns:
+        - fig (:obj:`matplotlib.figure.Figure`): Matplotlib figure object.
+        - ax (:obj:`matplotlib.axes.Axes`): Matplotlib axes object.
+    Examples:
+        >>> fig, ax = _get_or_create_heatmap_figure((10, 8))
+        >>> ax.plot([1, 2, 3], [4, 5, 6])
+    """
+    global _GLOBAL_HEATMAP_FIG, _GLOBAL_HEATMAP_AX
+    if _GLOBAL_HEATMAP_FIG is None:
+        _GLOBAL_HEATMAP_FIG, _GLOBAL_HEATMAP_AX = plt.subplots(figsize=figsize)
+    else:
+        # Clear previous content
+        _GLOBAL_HEATMAP_AX.clear()
+        # Adjust image size
+        _GLOBAL_HEATMAP_FIG.set_size_inches(figsize)
+    return _GLOBAL_HEATMAP_FIG, _GLOBAL_HEATMAP_AX
+
+
+def create_heatmap_with_values_fast(matrix, task_ids, title="Task-Expert Selection Frequencies"):
+    """
+    Overview:
+        Efficiently create annotated blue-themed heatmap with performance optimizations.
+        Optimizations include matplotlib figure reuse, selective value annotations,
+        optimized image conversion pipeline, and reduced DPI for faster computation.
+    Arguments:
+        - matrix (:obj:`numpy.ndarray`): Input matrix for heatmap visualization.
+        - task_ids (:obj:`list`): List of task identifiers for y-axis labels.
+        - title (:obj:`str`, optional): Heatmap title. Default is "Task-Expert Selection Frequencies".
+    Returns:
+        - img_array (:obj:`numpy.ndarray`): Image array in CHW format for TensorBoard logging.
+    Shapes:
+        - matrix: :math:`(N_{tasks}, N_{experts})` where N_tasks and N_experts are dimensions.
+        - img_array: :math:`(3, H, W)` where H and W are image height and width.
+    Examples:
+        >>> import numpy as np
+        >>> matrix = np.random.rand(5, 8)
+        >>> task_ids = [0, 1, 2, 3, 4]
+        >>> heatmap = create_heatmap_with_values_fast(matrix, task_ids)
+        >>> print(f"Heatmap shape: {heatmap.shape}")  # (3, height, width)
+    """
+    try:
+        figsize = (max(6, matrix.shape[1]), max(4, matrix.shape[0]))
+        fig, ax = _get_or_create_heatmap_figure(figsize)
+
+        # Intelligently choose whether to display value annotations
+        show_annot = matrix.size <= 64  # Only display values for 8x8 or smaller matrices
+
+        # Use matplotlib directly to avoid seaborn overhead
+        im = ax.imshow(matrix, cmap='Blues', aspect='auto')
+
+        # Selectively add value annotations
+        if show_annot:
+            for i in range(matrix.shape[0]):
+                for j in range(matrix.shape[1]):
+                    value = matrix[i, j]
+                    color = 'white' if value > 0.5 else 'black'
+                    ax.text(j, i, f'{value:.3f}', ha='center', va='center',
+                           color=color, fontsize=8)
+
+        # Set labels and title
+        ax.set_xticks(range(matrix.shape[1]))
+        ax.set_yticks(range(matrix.shape[0]))
+        ax.set_xticklabels([f'E{i}' for i in range(matrix.shape[1])], fontsize=10)
+        ax.set_yticklabels([f'T{tid}' for tid in task_ids], fontsize=10)
+        ax.set_title(title, fontsize=12, pad=15)
+        ax.set_xlabel('Experts', fontsize=10)
+        ax.set_ylabel('Tasks', fontsize=10)
+
+        # Simplified colorbar
+        if not hasattr(fig, '_colorbar_created'):
+            plt.colorbar(im, ax=ax, label='Frequency')
+            fig._colorbar_created = True
+
+        # Optimized image conversion: using lower DPI and simplified pipeline
+        fig.canvas.draw()
+        try:
+            # Get RGB data directly from canvas
+            if hasattr(fig.canvas, 'buffer_rgba'):
+                buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
+                buf = buf.reshape(fig.canvas.get_width_height()[::-1] + (4,))
+                img_array = buf[:, :, :3]  # Remove alpha channel
+            else:
+                buf = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+                img_array = buf.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+
+            # Convert to CHW format
+            img_array = img_array.transpose(2, 0, 1)
+
+        except Exception:
+            # Fallback: create simple blue gradient matrix
+            h, w = matrix.shape
+            img_array = np.zeros((3, h*20, w*20), dtype=np.uint8)
+            # Simple matrix upscaling and mapping to blue channel
+            matrix_resized = np.repeat(np.repeat(matrix, 20, axis=0), 20, axis=1)
+            img_array[2] = (matrix_resized * 255).astype(np.uint8)
+
+        return img_array
+
+    except Exception as e:
+        print(f"Warning: Heatmap generation failed: {e}, using fallback")
+        # Ultimate fallback: return blank image
+        return np.zeros((3, 100, 100), dtype=np.uint8)
+
+
+def create_heatmap_with_values(matrix, task_ids, title="Task-Expert Selection Frequencies"):
+    """
+    Overview:
+        Create annotated blue-themed heatmap using seaborn - original version for fallback.
+        This function serves as a backup when the optimized version encounters issues.
+    Arguments:
+        - matrix (:obj:`numpy.ndarray`): Input matrix for heatmap visualization.
+        - task_ids (:obj:`list`): List of task identifiers for y-axis labels.
+        - title (:obj:`str`, optional): Heatmap title. Default is "Task-Expert Selection Frequencies".
+    Returns:
+        - img_array (:obj:`numpy.ndarray`): Image array in CHW format for TensorBoard logging.
+    Shapes:
+        - matrix: :math:`(N_{tasks}, N_{experts})` where N_tasks and N_experts are dimensions.
+        - img_array: :math:`(3, H, W)` where H and W are image height and width.
+    Examples:
+        >>> import numpy as np
+        >>> matrix = np.random.rand(5, 8)
+        >>> task_ids = [0, 1, 2, 3, 4]
+        >>> heatmap = create_heatmap_with_values(matrix, task_ids)
+        >>> print(f"Heatmap shape: {heatmap.shape}")  # (3, height, width)
+    """
+    fig, ax = plt.subplots(figsize=(max(8, matrix.shape[1]), max(6, matrix.shape[0])))
+
+    # Use blue color scheme
+    sns.heatmap(matrix,
+                annot=True,  # Display values
+                fmt='.3f',   # Value format
+                cmap='Blues',  # Blue theme
+                ax=ax,
+                cbar_kws={'label': 'Selection Frequency'},
+                xticklabels=[f'Expert{i}' for i in range(matrix.shape[1])],
+                yticklabels=[f'Task{tid}' for tid in task_ids])
+
+    ax.set_title(title, fontsize=14, pad=20)
+    ax.set_xlabel('Experts', fontsize=12)
+    ax.set_ylabel('Tasks', fontsize=12)
+
+    plt.tight_layout()
+
+    # Save to BytesIO
+    buf = BytesIO()
+    plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
+    buf.seek(0)
+
+    # Convert to numpy array for tensorboard
+    img = Image.open(buf)
+    img_array = np.array(img)
+    buf.close()
+    plt.close(fig)
+
+    # Convert to CHW format (Channel, Height, Width)
+    if len(img_array.shape) == 3:
+        img_array = img_array.transpose(2, 0, 1)
+
+    return img_array
+
+
+def log_expert_selection_details(tb_logger, merged_stats, valid_task_ids, matrix, window_type, train_iter):
+    """
+    Overview:
+        Log detailed expert selection statistics for each task.
+        Records frequency entropy, variance, and total selections for analysis.
+    Arguments:
+        - tb_logger (:obj:`SummaryWriter`): TensorBoard logger for metric recording.
+        - merged_stats (:obj:`dict`): Merged expert selection statistics across ranks.
+        - valid_task_ids (:obj:`list`): List of valid task identifiers.
+        - matrix (:obj:`numpy.ndarray`): Expert selection frequency matrix.
+        - window_type (:obj:`str`): Time window type (immediate, short, medium, long).
+        - train_iter (:obj:`int`): Current training iteration for logging.
+    Examples:
+        >>> log_expert_selection_details(tb_logger, stats, [0,1,2], matrix, 'immediate', 1000)
+    """
+    for i, task_id in enumerate(valid_task_ids):
+        frequencies = matrix[i]
+        stats = merged_stats[task_id][window_type]
+
+        # Calculate and record task expert selection entropy (uniformity metric)
+        task_frequencies = np.array(frequencies)
+        task_frequencies = task_frequencies + 1e-8  # Avoid log(0)
+        task_entropy = -np.sum(task_frequencies * np.log(task_frequencies))
+        tb_logger.add_scalar(
+            f'MOE_Details/Task{task_id}_{window_type}/ExpertSelectionEntropy',
+            task_entropy, global_step=train_iter
+        )
+
+        # Record task expert selection variance (dispersion)
+        expert_variance = np.var(task_frequencies)
+        tb_logger.add_scalar(
+            f'MOE_Details/Task{task_id}_{window_type}/ExpertSelectionVariance',
+            expert_variance, global_step=train_iter
+        )
+
+        # Record task-level summary statistics
+        tb_logger.add_scalar(
+            f'MOE_Details/Task{task_id}_{window_type}/TotalSelections',
+            stats['total_selections'], global_step=train_iter
+        )
+        tb_logger.add_scalar(
+            f'MOE_Details/Task{task_id}_{window_type}/DataPoints',
+            stats['data_points'], global_step=train_iter
+        )
+
+
+def log_global_moe_statistics(tb_logger, matrix, window_type, valid_task_ids, train_iter):
+    """
+    Overview:
+        Log global MOE statistics including expert usage uniformity and extremes.
+        Provides system-wide view of expert utilization patterns.
+    Arguments:
+        - tb_logger (:obj:`SummaryWriter`): TensorBoard logger for metric recording.
+        - matrix (:obj:`numpy.ndarray`): Expert selection frequency matrix.
+        - window_type (:obj:`str`): Time window type (immediate, short, medium, long).
+        - valid_task_ids (:obj:`list`): List of valid task identifiers.
+        - train_iter (:obj:`int`): Current training iteration for logging.
+    Examples:
+        >>> log_global_moe_statistics(tb_logger, matrix, 'immediate', [0,1,2], 1000)
+    """
+    # Record basic information
+    tb_logger.add_scalar(
+        f'MOE_Global/{window_type}/NumActiveTasks',
+        len(valid_task_ids), global_step=train_iter
+    )
+    tb_logger.add_scalar(
+        f'MOE_Global/{window_type}/NumExperts',
+        matrix.shape[1], global_step=train_iter
+    )
+
+    # Calculate expert usage uniformity
+    expert_avg_usage = np.mean(matrix, axis=0)  # Average usage frequency per expert
+    usage_entropy = -np.sum(expert_avg_usage * np.log(expert_avg_usage + 1e-8))
+    tb_logger.add_scalar(
+        f'MOE_Global/{window_type}/ExpertUsageEntropy',
+        usage_entropy, global_step=train_iter
+    )
+
+    # Record most and least used experts
+    most_used_expert = np.argmax(expert_avg_usage)
+    least_used_expert = np.argmin(expert_avg_usage)
+    tb_logger.add_scalar(
+        f'MOE_Global/{window_type}/MostUsedExpert',
+        most_used_expert, global_step=train_iter
+    )
+    tb_logger.add_scalar(
+        f'MOE_Global/{window_type}/LeastUsedExpert',
+        least_used_expert, global_step=train_iter
+    )
+
+
+def process_and_log_moe_heatmaps_fast(tb_logger, merged_stats, window_type, train_iter):
+    """
+    Overview:
+        Efficiently process and log MOE heatmaps with performance optimizations.
+        Includes vectorized data processing, conditional heatmap generation,
+        and batch statistical processing.
+    Arguments:
+        - tb_logger (:obj:`SummaryWriter`): TensorBoard logger for metric recording.
+        - merged_stats (:obj:`dict`): Merged expert selection statistics across ranks.
+        - window_type (:obj:`str`): Time window type (immediate, short, medium, long).
+        - train_iter (:obj:`int`): Current training iteration for logging.
+    Examples:
+        >>> process_and_log_moe_heatmaps_fast(tb_logger, stats, 'immediate', 1000)
+    """
+    # Quick filtering of valid tasks
+    valid_task_data = [(tid, stats[window_type]['frequencies'])
+                      for tid, stats in merged_stats.items()
+                      if window_type in stats]
+
+    if not valid_task_data:
+        return
+
+    # Vectorized matrix construction
+    valid_task_ids, frequencies_list = zip(*valid_task_data)
+    matrix = np.array(frequencies_list)
+
+    # Conditional heatmap generation: only for small matrices
+    if matrix.size <= 200:  # Only generate heatmap when tasks*experts <= 200
+        try:
+            heatmap_img = create_heatmap_with_values_fast(
+                matrix, valid_task_ids,
+                f'MOE {window_type} Task-Expert Selection'
+            )
+
+            # Log heatmap to tensorboard
+            tb_logger.add_image(
+                f'MOE_Heatmap/{window_type}_TaskExpert_Heatmap',
+                heatmap_img,
+                global_step=train_iter,
+                dataformats='CHW'
+            )
+        except Exception as e:
+            print(f"Warning: Heatmap generation failed: {e}")
+
+    # Always log statistical data (lightweight operation)
+    log_expert_selection_details(tb_logger, merged_stats, valid_task_ids, matrix, window_type, train_iter)
+    log_global_moe_statistics(tb_logger, matrix, window_type, valid_task_ids, train_iter)
+
+
+def process_and_log_moe_heatmaps(tb_logger, merged_stats, window_type, train_iter):
+    """
+    Overview:
+        Process and log MOE heatmaps - original version for fallback.
+        This function serves as a backup when the optimized version encounters issues.
+    Arguments:
+        - tb_logger (:obj:`SummaryWriter`): TensorBoard logger for metric recording.
+        - merged_stats (:obj:`dict`): Merged expert selection statistics across ranks.
+        - window_type (:obj:`str`): Time window type (immediate, short, medium, long).
+        - train_iter (:obj:`int`): Current training iteration for logging.
+    Examples:
+        >>> process_and_log_moe_heatmaps(tb_logger, stats, 'immediate', 1000)
+    """
+    all_task_ids = sorted(merged_stats.keys())
+    task_expert_matrix = []
+    valid_task_ids = []
+
+    # Collect frequency data from valid tasks
+    for task_id in all_task_ids:
+        if window_type in merged_stats[task_id]:
+            frequencies = merged_stats[task_id][window_type]['frequencies']
+            task_expert_matrix.append(frequencies)
+            valid_task_ids.append(task_id)
+
+    if not task_expert_matrix:
+        return
+
+    # Convert to numpy matrix (num_tasks, num_experts)
+    matrix = np.array(task_expert_matrix)
+
+    # Create annotated blue-themed heatmap
+    heatmap_img = create_heatmap_with_values(
+        matrix, valid_task_ids,
+        f'MOE {window_type} Task-Expert Selection Frequencies'
+    )
+
+    # Log heatmap to tensorboard
+    tb_logger.add_image(
+        f'MOE_Heatmap/{window_type}_TaskExpert_Heatmap',
+        heatmap_img,
+        global_step=train_iter,
+        dataformats='CHW'
+    )
+
+    # Log detailed and global statistics
+    log_expert_selection_details(tb_logger, merged_stats, valid_task_ids, matrix, window_type, train_iter)
+
+
+def convert_stats_to_serializable(moe_stats):
+    """
+    Overview:
+        Convert tensor data in MOE statistics to serializable numpy format.
+        Ensures compatibility with distributed communication protocols.
+    Arguments:
+        - moe_stats (:obj:`dict`): MOE statistics containing tensor data.
+    Returns:
+        - converted (:obj:`dict`): Converted statistics with numpy arrays.
+    Examples:
+        >>> tensor_stats = {'task_0': {'immediate': {'frequencies': torch.tensor([0.1, 0.9])}}}
+        >>> numpy_stats = convert_stats_to_serializable(tensor_stats)
+        >>> type(numpy_stats['task_0']['immediate']['frequencies'])  # <class 'list'>
+    """
+    if not moe_stats:
+        return {}
+
+    converted = {}
+    for task_id, task_stats in moe_stats.items():
+        converted[task_id] = {}
+        for window_type, stats in task_stats.items():
+            if stats and 'frequencies' in stats:
+                converted[task_id][window_type] = {
+                    'frequencies': stats['frequencies'].cpu().numpy().tolist(),
+                    'total_selections': stats['total_selections'],
+                    'data_points': stats['data_points']
+                }
+    return converted
+
+
+def gather_distributed_moe_stats(local_stats, world_size):
+    """
+    Overview:
+        Gather MOE statistics from all GPUs in distributed training environment.
+        Handles communication failures gracefully with fallback to local statistics.
+    Arguments:
+        - local_stats (:obj:`dict`): Local GPU's MOE statistics.
+        - world_size (:obj:`int`): Total number of distributed training processes.
+    Returns:
+        - all_stats (:obj:`list`): List of statistics from all ranks.
+    Examples:
+        >>> local_data = {'task_0': {'immediate': {'frequencies': [0.1, 0.9]}}}
+        >>> all_data = gather_distributed_moe_stats(local_data, 4)
+        >>> len(all_data)  # 4 (from 4 GPUs)
+    """
+    all_stats = [None for _ in range(world_size)]
+    try:
+        dist.all_gather_object(all_stats, local_stats)
+        return all_stats
+    except Exception as e:
+        print(f"Distributed MOE statistics gathering failed: {e}")
+        return [local_stats]  # fallback to local statistics
+
+
+def collect_and_log_moe_statistics(policy, tb_logger, train_iter, world_size, rank):
+    """
+    Overview:
+        Collect and log MoE expert selection statistics including heatmaps and distribution analysis.
+        Handles distributed data collection, merging, and TensorBoard visualization.
+
+    Arguments:
+        - policy (:obj:`Policy`): Training policy with world_model.transformer supporting get_expert_selection_stats.
+        - tb_logger (:obj:`SummaryWriter`): TensorBoard logger for metric recording.
+        - train_iter (:obj:`int`): Current training iteration number.
+        - world_size (:obj:`int`): Total number of GPUs in distributed training.
+        - rank (:obj:`int`): Current GPU rank identifier.
+
+    Returns:
+        - None: No return value; performs logging only.
+
+    Notes:
+        Logs heatmaps for immediate/short/medium/long windows and JS/Wasserstein divergence between task distributions.
+
+    Examples:
+        >>> collect_and_log_moe_statistics(policy, tb_logger, 1000, 8, 0)
+    """
+    try:
+        # Step 1: Get MOE statistics from policy's transformer model
+        moe_stats = None
+
+        transformer = policy._model.world_model.transformer
+        if hasattr(transformer, 'get_expert_selection_stats'):
+            moe_stats = transformer.get_expert_selection_stats()
+
+        if moe_stats is None:
+            print(f"Rank {rank}: Warning: Unable to get MOE statistics, train_iter={train_iter}")
+            return
+
+        # Step 2: Convert tensor data to serializable format
+        serializable_stats = convert_stats_to_serializable(moe_stats)
+
+        print(f"Rank {rank}: Local MOE statistics - tasks: {len(serializable_stats)}, train_iter={train_iter}")
+
+        # Step 3: Gather statistics from all GPUs in distributed setting
+        all_expert_stats = gather_distributed_moe_stats(serializable_stats, world_size)
+
+        # Step 4: Merge statistics data
+        merged_stats = merge_expert_stats_across_ranks(all_expert_stats)
+
+        if not merged_stats:
+            print(f"Rank {rank}: Warning: Merged MOE statistics empty, train_iter={train_iter}")
+            return
+
+        # Step 5: All GPUs log MOE statistics
+        print(f"Rank {rank}: Starting MOE statistics logging - merged tasks: {len(merged_stats)}, train_iter={train_iter}")
+
+        # Generate heatmaps and statistics for each time window
+        for window_type in ['immediate', 'short', 'medium', 'long']:
+            if any(window_type in task_stats for task_stats in merged_stats.values()):
+                process_and_log_moe_heatmaps_fast(tb_logger, merged_stats, window_type, train_iter)
+
+        # Log overall MOE usage
+        tb_logger.add_scalar('MOE_Global/ActiveTasks', len(merged_stats), global_step=train_iter)
+
+        # Step 6: Add distribution difference computation and logging
+        if any('immediate' in task_stats for task_stats in merged_stats.values()):
+            print(f"Rank {rank}: Starting inter-task distribution difference calculation...")
+            collect_and_log_divergences_with_heatmaps(tb_logger, merged_stats, train_iter)
+
+        print(f"Rank {rank}: MOE statistics logging completed, train_iter={train_iter}")
+
+    except Exception as e:
+        print(f"Rank {rank}: MOE statistics collection failed - {e}, train_iter={train_iter}")
+        import traceback
+        traceback.print_exc()
+
+
+# ====== GPU-Optimized Distribution Divergence Calculation and Visualization Functions ======
+def jensen_shannon_divergence_batch_gpu(distributions_tensor):
+    """
+    Overview:
+        GPU batch computation of JS divergence matrix - fully vectorized, no loops.
+        Efficiently computes Jensen-Shannon divergence between all pairs of distributions.
+
+    Arguments:
+        - distributions_tensor (:obj:`torch.Tensor`): Shape (n_tasks, n_experts), GPU tensor.
+
+    Returns:
+        - js_matrix (:obj:`torch.Tensor`): Shape (n_tasks, n_tasks), symmetric matrix.
+
+    Shapes:
+        - distributions_tensor: :math:`(N_{tasks}, N_{experts})`
+        - js_matrix: :math:`(N_{tasks}, N_{tasks})`
+
+    Notes:
+        Input is normalized to a probability distribution before computation.
+    Examples:
+        >>> dist_tensor = torch.rand(5, 8).cuda()
+        >>> js_matrix = jensen_shannon_divergence_batch_gpu(dist_tensor)
+        >>> print(js_matrix.shape)  # torch.Size([5, 5])
+    """
+    device = distributions_tensor.device
+    n_tasks, n_experts = distributions_tensor.shape
+
+    # 1. Normalize to probability distributions
+    eps = 1e-8
+    distributions_tensor = distributions_tensor / (distributions_tensor.sum(dim=1, keepdim=True) + eps)
+
+    # 2. Use broadcasting to compute average distributions for all task pairs
+    # P_i: (n_tasks, 1, n_experts), P_j: (1, n_tasks, n_experts)
+    P_i = distributions_tensor.unsqueeze(1)
+    P_j = distributions_tensor.unsqueeze(0)
+    M = 0.5 * (P_i + P_j)  # shape: (n_tasks, n_tasks, n_experts)
+
+    # 3. Batch compute KL divergences - fully vectorized
+    # KL(P_i || M) for all pairs
+    log_ratio_i = torch.log((P_i + eps) / (M + eps))
+    kl_i_m = torch.sum(P_i * log_ratio_i, dim=2)  # (n_tasks, n_tasks)
+
+    # KL(P_j || M) for all pairs
+    log_ratio_j = torch.log((P_j + eps) / (M + eps))
+    kl_j_m = torch.sum(P_j * log_ratio_j, dim=2)  # (n_tasks, n_tasks)
+
+    # 4. JS divergence matrix
+    js_matrix = 0.5 * (kl_i_m + kl_j_m)
+
+    return js_matrix
+
+
+def wasserstein_distance_batch_gpu(distributions_tensor):
+    """
+    Overview:
+        GPU batch computation of Wasserstein distance matrix - efficient 1D distribution implementation.
+        Computes Earth Mover's Distance between all pairs of discrete distributions.
+
+    Arguments:
+        - distributions_tensor (:obj:`torch.Tensor`): Shape (n_tasks, n_experts), GPU tensor.
+
+    Returns:
+        - wasserstein_matrix (:obj:`torch.Tensor`): Shape (n_tasks, n_tasks), symmetric matrix.
+
+    Shapes:
+        - distributions_tensor: :math:`(N_{tasks}, N_{experts})`
+        - wasserstein_matrix: :math:`(N_{tasks}, N_{tasks})`
+
+    Notes:
+        Uses L1 norm of CDF differences for discrete distributions.
+    Examples:
+        >>> dist_tensor = torch.rand(5, 8).cuda()
+        >>> wass_matrix = wasserstein_distance_batch_gpu(dist_tensor)
+        >>> print(wass_matrix.shape)  # torch.Size([5, 5])
+    """
+    device = distributions_tensor.device
+    n_tasks, n_experts = distributions_tensor.shape
+    eps = 1e-8
+
+    # 1. Normalize to probability distributions
+    distributions_tensor = distributions_tensor / (distributions_tensor.sum(dim=1, keepdim=True) + eps)
+
+    # 2. Compute cumulative distribution functions (CDF)
+    cdf_tensor = torch.cumsum(distributions_tensor, dim=1)  # (n_tasks, n_experts)
+
+    # 3. Use broadcasting to compute L1 distances between all CDF pairs
+    cdf_i = cdf_tensor.unsqueeze(1)  # (n_tasks, 1, n_experts)
+    cdf_j = cdf_tensor.unsqueeze(0)  # (1, n_tasks, n_experts)
+
+    # Wasserstein distance = L1 norm of cumulative distribution differences
+    wasserstein_matrix = torch.sum(torch.abs(cdf_i - cdf_j), dim=2)
+
+    return wasserstein_matrix
+
+
+def compute_distribution_divergences_optimized(merged_stats, window_type='immediate'):
+    """
+    Overview:
+        GPU-optimized version for efficient distribution divergence computation.
+        Leverages GPU acceleration for batch processing of divergence metrics.
+
+    Arguments:
+        - merged_stats (:obj:`dict`): Merged MoE statistics from all distributed ranks.
+        - window_type (:obj:`str`, optional): Time window type. Default is 'immediate'.
+
+    Returns:
+        - divergence_data (:obj:`dict`): Dictionary containing:
+            - task_ids: List of task IDs
+            - n_tasks, n_experts: Dimensions
+            - device, gpu_accelerated: Device info
+            - js_matrix, wasserstein_matrix: Divergence matrices (numpy)
+            - js_stats, wasserstein_stats: avg, max, min, std for each metric
+
+    Examples:
+        >>> stats = {'task_0': {'immediate': {'frequencies': [0.1, 0.9]}}}
+        >>> result = compute_distribution_divergences_optimized(stats)
+        >>> print(f"GPU accelerated: {result['gpu_accelerated']}")
+    """
+    # 1. Data preprocessing
+    valid_tasks = [(tid, stats[window_type]['frequencies'])
+                  for tid, stats in merged_stats.items()
+                  if window_type in stats]
+
+    if len(valid_tasks) < 2:
+        return {}
+
+    task_ids, frequencies_list = zip(*valid_tasks)
+
+    # 2. Efficient tensor conversion
+    try:
+        if isinstance(frequencies_list[0], torch.Tensor):
+            frequencies_tensor = torch.stack(frequencies_list)
+        else:
+            frequencies_tensor = torch.tensor(
+                np.array(frequencies_list),
+                dtype=torch.float32
+            )
+
+        # Automatic GPU acceleration
+        if torch.cuda.is_available():
+            frequencies_tensor = frequencies_tensor.cuda()
+
+    except Exception as e:
+        print(f"GPU conversion failed, using CPU: {e}")
+        frequencies_tensor = torch.tensor(np.array(frequencies_list), dtype=torch.float32)
+
+    device = frequencies_tensor.device
+    n_tasks, n_experts = frequencies_tensor.shape
+
+    # 3. GPU batch computation (no loops)
+    with torch.no_grad():
+        # Batch compute JS divergence and Wasserstein distance
+        js_matrix = jensen_shannon_divergence_batch_gpu(frequencies_tensor)
+        wasserstein_matrix = wasserstein_distance_batch_gpu(frequencies_tensor)
+
+        # Efficiently extract upper triangular values (avoid duplicate computation)
+        triu_indices = torch.triu_indices(n_tasks, n_tasks, offset=1, device=device)
+        js_values = js_matrix[triu_indices[0], triu_indices[1]]
+        wasserstein_values = wasserstein_matrix[triu_indices[0], triu_indices[1]]
+
+        # Statistical computation (vectorized)
+        js_stats = {
+            'avg': torch.mean(js_values).item(),
+            'max': torch.max(js_values).item(),
+            'min': torch.min(js_values).item(),
+            'std': torch.std(js_values).item()
+        }
+
+        wasserstein_stats = {
+            'avg': torch.mean(wasserstein_values).item(),
+            'max': torch.max(wasserstein_values).item(),
+            'min': torch.min(wasserstein_values).item(),
+            'std': torch.std(wasserstein_values).item()
+        }
+
+    return {
+        'task_ids': task_ids,
+        'n_tasks': n_tasks,
+        'n_experts': n_experts,
+        'device': str(device),
+        'gpu_accelerated': 'cuda' in str(device),
+
+        # Return CPU versions for logging
+        'js_matrix': js_matrix.cpu().numpy(),
+        'wasserstein_matrix': wasserstein_matrix.cpu().numpy(),
+        'js_stats': js_stats,
+        'wasserstein_stats': wasserstein_stats
+    }
+
+
+def create_similarity_heatmap_no_diagonal(similarity_matrix, task_ids, metric_name, title_suffix=""):
+    """
+    Overview:
+        Create task similarity heatmap with diagonal elements removed.
+        Provides clear visualization of inter-task relationships without self-similarity noise.
+    Arguments:
+        - similarity_matrix (:obj:`numpy.ndarray`): Similarity matrix (n_tasks, n_tasks).
+        - task_ids (:obj:`list`): Task identifier list for axis labels.
+        - metric_name (:obj:`str`): Metric name ('js_divergence', 'wasserstein_distance').
+        - title_suffix (:obj:`str`, optional): Additional title suffix. Default is "".
+    Returns:
+        - img_array (:obj:`numpy.ndarray`): Image array in CHW format for TensorBoard.
+    Shapes:
+        - similarity_matrix: :math:`(N_{tasks}, N_{tasks})`
+        - img_array: :math:`(3, H, W)` where H and W are image dimensions.
+    Examples:
+        >>> matrix = np.random.rand(5, 5)
+        >>> task_ids = [0, 1, 2, 3, 4]
+        >>> heatmap = create_similarity_heatmap_no_diagonal(matrix, task_ids, 'js_divergence')
+        >>> print(f"Output shape: {heatmap.shape}")  # (3, height, width)
+    """
+    try:
+        # Copy matrix to avoid modifying original data
+        matrix = similarity_matrix.copy()
+
+        # Set diagonal to NaN so matplotlib displays as blank
+        np.fill_diagonal(matrix, np.nan)
+
+        figsize = (max(6, len(task_ids)), max(4, len(task_ids)))
+        fig, ax = plt.subplots(figsize=figsize)  # Create new figure to avoid reuse issues
+
+        # Choose color mapping based on metric type
+        if 'js' in metric_name.lower():
+            cmap = 'Reds'
+            title_name = 'JS Divergence'
+            vmin, vmax = 0, 1.0
+        else:  # wasserstein
+            cmap = 'Blues'
+            title_name = 'Wasserstein Distance'
+            vmin, vmax = None, None  # Adaptive
+
+        # Use masked array to handle NaN values, diagonal displays as white
+        masked_matrix = np.ma.masked_invalid(matrix)
+        im = ax.imshow(masked_matrix, cmap=cmap, vmin=vmin, vmax=vmax, aspect='auto')
+
+        # Add value annotations (skip diagonal)
+        if len(task_ids) <= 15:  # Only add annotations for smaller task counts
+            for i in range(len(task_ids)):
+                for j in range(len(task_ids)):
+                    if i != j:  # Skip diagonal
+                        value = matrix[i, j]
+                        if not np.isnan(value):
+                            threshold = (vmax or np.nanmax(matrix)) * 0.5 if vmax else np.nanmax(matrix) * 0.5
+                            color = 'white' if value > threshold else 'black'
+                            ax.text(j, i, f'{value:.3f}', ha='center', va='center',
+                                   color=color, fontsize=8)
+
+        # Set labels
+        ax.set_xticks(range(len(task_ids)))
+        ax.set_yticks(range(len(task_ids)))
+        ax.set_xticklabels([f'T{tid}' for tid in task_ids], fontsize=9)
+        ax.set_yticklabels([f'T{tid}' for tid in task_ids], fontsize=9)
+        ax.set_title(f'Task {title_name} Matrix {title_suffix} (No Diagonal)', fontsize=12)
+        ax.set_xlabel('Tasks', fontsize=10)
+        ax.set_ylabel('Tasks', fontsize=10)
+
+        # Add colorbar
+        plt.colorbar(im, ax=ax, label=title_name, shrink=0.8)
+
+        # Convert to image array - fix matplotlib version compatibility
+        fig.canvas.draw()
+
+        try:
+            # New matplotlib uses buffer_rgba
+            if hasattr(fig.canvas, 'buffer_rgba'):
+                buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
+                h, w = fig.canvas.get_width_height()
+                img_array = buf.reshape(h, w, 4)[:, :, :3]  # Remove alpha channel
+            else:
+                # Old matplotlib fallback
+                buf = fig.canvas.print_to_string()
+                img_array = np.frombuffer(buf, dtype=np.uint8)
+                h, w = fig.canvas.get_width_height()
+                img_array = img_array.reshape(h, w, 3)
+        except Exception as conv_e:
+            print(f"Image conversion method failed: {conv_e}, trying PIL approach")
+            # Final fallback: convert through PIL
+            buf = BytesIO()
+            fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
+            buf.seek(0)
+            img = Image.open(buf)
+            img_array = np.array(img)[:, :, :3]  # Remove alpha channel
+            buf.close()
+
+        img_array = img_array.transpose(2, 0, 1)  # CHW format
+        plt.close(fig)  # Close figure to avoid memory leak
+
+        return img_array
+
+    except Exception as e:
+        print(f"Warning: No-diagonal heatmap generation failed: {e}")
+        return np.zeros((3, 100, 100), dtype=np.uint8)
+
+
+def log_pairwise_optimized(tb_logger, divergence_data, train_iter):
+    """
+    Overview:
+        Optimized task pair logging with batch processing.
+        Efficiently logs pairwise divergence metrics for all task combinations.
+    Arguments:
+        - tb_logger (:obj:`SummaryWriter`): TensorBoard logger for metric recording.
+        - divergence_data (:obj:`dict`): Divergence computation results.
+        - train_iter (:obj:`int`): Current training iteration for logging.
+    Examples:
+        >>> log_pairwise_optimized(tb_logger, divergence_data, 1000)
+    """
+    task_ids = divergence_data['task_ids']
+    js_matrix = divergence_data['js_matrix']
+    wasserstein_matrix = divergence_data['wasserstein_matrix']
+
+    # Batch construct task pair metric dictionary
+    pairwise_scalars = {}
+
+    for i, task_i in enumerate(task_ids):
+        for j, task_j in enumerate(task_ids):
+            if i < j:  # Only log upper triangle
+                # Construct metric names
+                js_key = f'TaskPairwise/Immediate_Task{task_i}_Task{task_j}_JS_Divergence'
+                wass_key = f'TaskPairwise/Immediate_Task{task_i}_Task{task_j}_Wasserstein_Distance'
+
+                pairwise_scalars[js_key] = js_matrix[i, j]
+                pairwise_scalars[wass_key] = wasserstein_matrix[i, j]
+
+    # Batch write to TensorBoard
+    for key, value in pairwise_scalars.items():
+        tb_logger.add_scalar(key, float(value), global_step=train_iter)
+
+
+def log_divergences_with_heatmaps(tb_logger, divergence_data, train_iter):
+    """
+    Overview:
+        Log distribution divergence metrics and heatmaps (with diagonal removed).
+        Comprehensive logging of inter-task distribution analysis results.
+    Arguments:
+        - tb_logger (:obj:`SummaryWriter`): TensorBoard logger for metric recording.
+        - divergence_data (:obj:`dict`): Divergence computation results.
+        - train_iter (:obj:`int`): Current training iteration for logging.
+    Examples:
+        >>> log_divergences_with_heatmaps(tb_logger, divergence_data, 1000)
+    """
+    if not divergence_data:
+        return
+
+    js_stats = divergence_data['js_stats']
+    wasserstein_stats = divergence_data['wasserstein_stats']
+    task_ids = divergence_data['task_ids']
+    n_tasks = divergence_data['n_tasks']
+
+    # Debug: Check matrix data
+    js_matrix = divergence_data['js_matrix']
+    wasserstein_matrix = divergence_data['wasserstein_matrix']
+    print(f"DEBUG: JS matrix shape={js_matrix.shape}, range=[{np.min(js_matrix):.6f}, {np.max(js_matrix):.6f}]")
+    print(f"DEBUG: Wasserstein matrix shape={wasserstein_matrix.shape}, range=[{np.min(wasserstein_matrix):.6f}, {np.max(wasserstein_matrix):.6f}]")
+
+    # 1. Log scalar metrics
+    scalar_dict = {
+        'MOE_Divergence/Immediate_AvgJS_Divergence': js_stats['avg'],
+        'MOE_Divergence/Immediate_MaxJS_Divergence': js_stats['max'],
+        'MOE_Divergence/Immediate_AvgWasserstein_Distance': wasserstein_stats['avg'],
+        'MOE_Divergence/Immediate_MaxWasserstein_Distance': wasserstein_stats['max'],
+    }
+
+    for key, value in scalar_dict.items():
+        tb_logger.add_scalar(key, value, global_step=train_iter)
+
+    # 1.1 Print core metrics to console
+    print("=" * 65)
+    print(f" Inter-Task Distribution Divergence Statistics (Iteration: {train_iter})")
+    print("=" * 65)
+    print(f"Participating tasks: {n_tasks} | Task IDs: {list(task_ids)}")
+    print(f"Computing device: {divergence_data.get('device', 'Unknown')} | GPU acceleration: {'Enabled' if divergence_data.get('gpu_accelerated', False) else 'Disabled'}")
+    print("-" * 65)
+    print("JS Divergence (Jensen-Shannon Divergence):")
+    print(f"  Average: {js_stats['avg']:.6f}  |  Maximum: {js_stats['max']:.6f}")
+    print(f"  Minimum: {js_stats['min']:.6f}  |  Std Dev: {js_stats['std']:.6f}")
+    print("-" * 65)
+    print("Wasserstein Distance:")
+    print(f"  Average: {wasserstein_stats['avg']:.6f}  |  Maximum: {wasserstein_stats['max']:.6f}")
+    print(f"  Minimum: {wasserstein_stats['min']:.6f}  |  Std Dev: {wasserstein_stats['std']:.6f}")
+    print("=" * 65)
+
+    # 2. Log similarity matrix heatmaps with diagonal removed
+    task_ids = divergence_data['task_ids']
+    n_tasks = divergence_data['n_tasks']
+
+    if n_tasks <= 25:  # Limit matrix size to avoid oversized heatmaps
+        try:
+            # JS divergence matrix heatmap (no diagonal)
+            js_heatmap = create_similarity_heatmap_no_diagonal(
+                divergence_data['js_matrix'],
+                task_ids,
+                'js_divergence',
+                f'(Immediate-{n_tasks} tasks)'
+            )
+            tb_logger.add_image(
+                'TaskSimilarity/Immediate_JS_Matrix_NoDiagonal',
+                js_heatmap,
+                global_step=train_iter,
+                dataformats='CHW'
+            )
+
+            # Wasserstein distance matrix heatmap (no diagonal)
+            wass_heatmap = create_similarity_heatmap_no_diagonal(
+                divergence_data['wasserstein_matrix'],
+                task_ids,
+                'wasserstein_distance',
+                f'(Immediate-{n_tasks} tasks)'
+            )
+            tb_logger.add_image(
+                'TaskSimilarity/Immediate_Wasserstein_Matrix_NoDiagonal',
+                wass_heatmap,
+                global_step=train_iter,
+                dataformats='CHW'
+            )
+
+        except Exception as e:
+            print(f"Warning: Similarity matrix heatmap generation failed: {e}")
+
+    # 3. Log task pair metrics (optional)
+    if n_tasks <= 20:
+        log_pairwise_optimized(tb_logger, divergence_data, train_iter)
+
+
+def collect_and_log_divergences_with_heatmaps(tb_logger, merged_stats, train_iter):
+    """
+    Overview:
+        Complete distribution divergence computation and logging (including no-diagonal heatmaps).
+        End-to-end pipeline for analyzing and visualizing inter-task distribution differences.
+    Arguments:
+        - tb_logger (:obj:`SummaryWriter`): TensorBoard logger for metric recording.
+        - merged_stats (:obj:`dict`): Merged MOE statistics from distributed training.
+        - train_iter (:obj:`int`): Current training iteration for logging.
+    Examples:
+        >>> collect_and_log_divergences_with_heatmaps(tb_logger, merged_stats, 1000)
+    """
+    try:
+        # GPU-optimized computation
+        divergence_data = compute_distribution_divergences_optimized(merged_stats, 'immediate')
+
+        if not divergence_data:
+            print(f"Skipping distribution divergence computation - insufficient tasks (need >=2 tasks)")
+            return
+
+        # Log metrics and heatmaps
+        log_divergences_with_heatmaps(tb_logger, divergence_data, train_iter)
+
+        # Summary print
+        print(f">> Distribution divergence statistics completed and logged to TensorBoard")
+        if divergence_data.get('n_tasks', 0) <= 25:
+            print(f">> Similarity matrix heatmaps generated (diagonal removed)")
+        if divergence_data.get('n_tasks', 0) <= 20:
+            print(f">> Task pair detailed metrics logged")
+        print()  # Blank line separator
+
+    except Exception as e:
+        print(f"ERROR: Distribution divergence computation failed - {e}")
+        import traceback
+        traceback.print_exc()
\ No newline at end of file
diff --git a/lzero/model/unizero_world_models/moe.py b/lzero/model/unizero_world_models/moe.py
index 17ce0605b..dcc774640 100644
--- a/lzero/model/unizero_world_models/moe.py
+++ b/lzero/model/unizero_world_models/moe.py
@@ -119,13 +119,30 @@ def __init__(self, config: Any, experts: List[nn.Module], gate: nn.Module, num_e
             )
         else:
             self.shared_expert = None
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        self.device = next(iter(experts)).w1.weight.device if experts else torch.device('cuda')
+        
+        # Sliding window configuration
+        self.window_sizes = {
+            'immediate': 100,    # Immediate statistics (last 100 steps)
+            'short': 1000,       # Short-term statistics (last 1000 steps)
+            'medium': 10000,     # Medium-term statistics (last 10000 steps)
+            'long': 100000       # Long-term statistics (last 100000 steps)
+        }
+        
+        # GPU statistics buffer: task_id -> {window_type -> [expert selection history]}
+        self.expert_stats_gpu = {}
+        self.step_count = 0
+
+    def forward(self, x: torch.Tensor, task_id: int = None) -> torch.Tensor:
         """
         Overview:
             Performs the forward pass for the MoE layer.
+
         Arguments:
             - x (:obj:`torch.Tensor`): The input tensor of shape [batch_size, seq_len, dim].
+            - task_id (:obj:`int`, optional): Current task ID. When not None and in training mode,
+                expert selection statistics are collected for gradient conflict analysis.
+
         Returns:
             - torch.Tensor: The output tensor with the same shape as the input.
         """
@@ -133,31 +150,27 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         original_shape = x.size()
         x = x.view(-1, self.dim)
 
-        # Compute gate logits, shape: [num_tokens, num_experts]
-        gate_logits = self.gate(x)
-        # Select top-k experts for each token.
-        weights, indices = torch.topk(gate_logits, self.num_experts_per_tok, dim=1)
-        # Normalize the weights of selected experts using softmax.
-        weights = F.softmax(weights, dim=1).to(x.dtype)
-
-        # Initialize the output tensor for expert computations.
-        expert_output = torch.zeros_like(x)
-
-        # Iterate over each expert to compute outputs for the tokens routed to it.
-        for expert_id in range(self.num_experts):
-            # Find the tokens that have this expert in their top-k list.
-            batch_idx, expert_tok_idx = torch.where(indices == expert_id)
-            if batch_idx.numel() == 0:
-                continue
-
-            # Select the subset of tokens for the current expert.
-            token_subset = x[batch_idx]  # Shape: [num_tokens_for_expert, dim]
-            # Compute the output from the current expert.
-            output_expert = self.experts[expert_id](token_subset)
-            # Get the corresponding weights for these tokens.
-            token_weights = weights[batch_idx, expert_tok_idx].unsqueeze(-1)
-            # Apply weights and accumulate the output.
-            expert_output[batch_idx] += output_expert * token_weights
+       
+        expert_output = x
+        if self.num_experts != 0:
+            # Gate logits: [N, num_experts], N = num tokens
+            gate_logits = self.gate(x)
+            # Top-k experts per token
+            weights, indices = torch.topk(gate_logits, self.num_experts_per_tok, dim=1)
+            weights = F.softmax(weights, dim=1).to(x.dtype)
+
+            if self.training and task_id is not None:
+                self._collect_expert_selection_stats(task_id, indices)
+
+            expert_output = torch.zeros_like(x)
+            for expert_id in range(self.num_experts):
+                batch_idx, expert_tok_idx = torch.where(indices == expert_id)
+                if batch_idx.numel() == 0:
+                    continue
+                token_subset = x[batch_idx]
+                output_expert = self.experts[expert_id](token_subset)
+                token_weights = weights[batch_idx, expert_tok_idx].unsqueeze(-1)
+                expert_output[batch_idx] += output_expert * token_weights
 
         # If a shared expert exists, add its output.
         if self.shared_expert is not None:
@@ -168,8 +181,132 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         # Restore the original tensor shape and return.
         return output.view(original_shape)
+    
+    def _collect_expert_selection_stats(self, task_id: int, indices: torch.Tensor) -> None:
+        """
+        Overview:
+            Collect expert selection statistics in GPU memory using multi-granularity sliding windows.
+            Maintains rolling buffers for immediate/short/medium/long windows to track expert usage.
+
+        Arguments:
+            - task_id (:obj:`int`): The identifier of the current task.
+            - indices (:obj:`torch.Tensor`): Expert indices selected by the router for the current batch.
+
+        Shapes:
+            - indices: :math:`(N, k)` where N is num tokens and k is num_experts_per_tok.
 
+        Examples:
+            >>> # Collect stats for task 0 with expert indices
+            >>> indices = torch.tensor([[0, 2], [1, 3]])  # batch_size=2, k=2
+            >>> moe_layer._collect_expert_selection_stats(task_id=0, indices=indices)
+        """
+        self.step_count += 1
+        
+        if task_id not in self.expert_stats_gpu:
+            self.expert_stats_gpu[task_id] = {}
+            for window_type in self.window_sizes.keys():
+                self.expert_stats_gpu[task_id][window_type] = torch.zeros(
+                    self.window_sizes[window_type], 
+                    self.num_experts, 
+                    dtype=torch.float32, 
+                    device=self.device
+                )
+        
+        # Calculate expert selection frequency for current batch
+        indices_flat = indices.flatten()  # [N*k]
+        expert_counts = torch.zeros(self.num_experts, device=self.device, dtype=torch.float32)
+        for expert_id in range(self.num_experts):
+            expert_counts[expert_id] = (indices_flat == expert_id).sum().float()
+        
+        # Update sliding windows for all granularities
+        for window_type, window_size in self.window_sizes.items():
+            buffer = self.expert_stats_gpu[task_id][window_type]
+            # Sliding window: new data goes to the end, old data moves forward
+            buffer[:-1] = buffer[1:].clone()
+            buffer[-1] = expert_counts
+    
+    def get_expert_selection_stats(self, task_id: int = None) -> dict:
+        """
+        Overview:
+            Get multi-granularity expert selection frequency statistics.
+
+        Arguments:
+            - task_id (:obj:`int`, optional): Specific task ID. If None, returns stats for all tasks.
+
+        Returns:
+            - stats (:obj:`dict`): {task_id: {window_type: {frequencies, total_counts, total_selections, data_points}}}.
+        Examples:
+            >>> # Get stats for all tasks
+            >>> all_stats = moe_layer.get_expert_selection_stats()
+            >>> # Get stats for specific task
+            >>> task_stats = moe_layer.get_expert_selection_stats(task_id=0)
+        """
+        if task_id is None:
+            # Return statistics for all tasks
+            all_stats = {}
+            for tid in self.expert_stats_gpu.keys():
+                all_stats[tid] = self._compute_task_stats(tid)
+            return all_stats
+        else:
+            # Return statistics for specified task
+            return self._compute_task_stats(task_id)
+    
+    def _compute_task_stats(self, task_id: int) -> dict:
+        """
+        Overview:
+            Compute multi-granularity statistics for a specified task.
+
+        Arguments:
+            - task_id (:obj:`int`): The task identifier.
+
+        Returns:
+            - stats (:obj:`dict`): {window_type: {frequencies, total_counts, total_selections, data_points}}.
+
+        Shapes:
+            - frequencies: :math:`(num\_experts,)` normalized selection frequencies per expert.
+            - total_counts: :math:`(num\_experts,)` absolute selection counts per expert.
+        Examples:
+            >>> # Compute stats for task 0
+            >>> task_stats = moe_layer._compute_task_stats(task_id=0)
+            >>> immediate_freq = task_stats['immediate']['frequencies']
+        """
+        if task_id not in self.expert_stats_gpu:
+            return {}
+        
+        stats = {}
+        for window_type, buffer in self.expert_stats_gpu[task_id].items():
+            # Simplified version: directly average all existing data, ignoring whether window is full
+            # buffer shape: [window_size, num_experts]
+            total_counts = buffer.sum(dim=0)  # [num_experts]
+            total_selections = total_counts.sum()
+            
+            if total_selections > 0:
+                frequencies = total_counts / total_selections
+            else:
+                frequencies = torch.zeros(self.num_experts, device=self.device)
+            
+            stats[window_type] = {
+                'frequencies': frequencies,  # Keep tensor format
+                'total_counts': total_counts,  # Keep tensor format  
+                'total_selections': total_selections.item(),
+                'data_points': min(self.step_count, self.window_sizes[window_type])
+            }
+        
+        return stats
+    
+    def reset_expert_selection_stats(self) -> None:
+        """
+        Overview:
+            Reset expert selection statistics and clear GPU buffers.
+
+        Examples:
+            >>> # Reset all expert selection statistics
+            >>> moe_layer.reset_expert_selection_stats()
+        """
+        self.expert_stats_gpu.clear()
+        self.step_count = 0
 
+class MoELayerOptimized(nn.Module):
     """
     Overview:
         An optimized implementation of the Mixture-of-Experts (MoE) layer that maintains the same API as `MoELayer`.
diff --git a/lzero/model/unizero_world_models/transformer.py b/lzero/model/unizero_world_models/transformer.py
index 399f98929..e3c309d36 100644
--- a/lzero/model/unizero_world_models/transformer.py
+++ b/lzero/model/unizero_world_models/transformer.py
@@ -12,7 +12,7 @@
 import math
 import logging
 from dataclasses import dataclass
-from typing import Optional
+from typing import Dict, List, Optional
 
 import torch
 import torch.nn as nn
@@ -349,6 +349,13 @@ def __init__(self, config: TransformerConfig, task_embed: Optional[nn.Module] =
         self.blocks = nn.ModuleList([Block(config) for _ in range(config.num_layers)])
         self.ln_f = nn.LayerNorm(config.embed_dim)
 
+        self.num_blocks=len(self.blocks)
+        self.num_experts=config.num_experts_of_moe_in_transformer
+
+        self.shared_expert = 0
+        if hasattr(config, "n_shared_experts") and config.n_shared_experts > 0:
+            self.shared_expert = config.n_shared_experts
+
         self.task_embed = task_embed
         self.task_embed_option = self.config.task_embed_option
         self.use_register_token = (self.task_embed_option == "register_task_embed")
@@ -443,10 +450,16 @@ def forward(
 
         x = self.drop(sequences)
 
+        # for i, block in enumerate(self.blocks):
+        #     kv_cache_layer = None if past_keys_values is None else past_keys_values[i]
+        #     x = block(x, kv_cache_layer, valid_context_lengths)
         for i, block in enumerate(self.blocks):
-            kv_cache_layer = None if past_keys_values is None else past_keys_values[i]
-            x = block(x, kv_cache_layer, valid_context_lengths)
+            is_last_block = (i == len(self.blocks) - 1)
+            x = block(x,
+                    None if past_keys_values is None else past_keys_values[i],
+                    valid_context_lengths, is_last_block=is_last_block, task_id=task_id)
 
+            
         x = self.ln_f(x)
 
         if self.use_register_token:
@@ -460,6 +473,147 @@ def forward(
 
         return x
 
+    def get_expert_selection_stats(self, task_id: int = None) -> dict:
+        """
+        Overview:
+            Retrieve MoE expert selection statistics from the last transformer block.
+        Arguments:
+            - task_id (:obj:`int`, optional): Task identifier for task-specific statistics. Default is None.
+        Returns:
+            - stats (:obj:`dict`): Dictionary containing expert selection statistics.
+        """
+        if len(self.blocks) == 0:
+            return {}
+        last_block = self.blocks[-1]
+        if not hasattr(last_block, 'feed_forward') or not hasattr(last_block.feed_forward, 'get_expert_selection_stats'):
+            return {}
+        return last_block.feed_forward.get_expert_selection_stats(task_id)
+
+    def reset_expert_selection_stats(self) -> None:
+        """
+        Overview:
+            Reset MoE expert selection statistics for the last transformer block.
+        """
+        if len(self.blocks) == 0:
+            return
+        last_block = self.blocks[-1]
+        if hasattr(last_block, 'feed_forward') and hasattr(last_block.feed_forward, 'reset_expert_selection_stats'):
+            last_block.feed_forward.reset_expert_selection_stats()
+
+    def get_shared_expert_gradients_by_block_id(self, block_id: int) -> Dict[str, torch.Tensor]:
+        """
+        Overview:
+            Retrieve parameter gradients of shared experts from a specified transformer block.
+        Arguments:
+            - block_id (:obj:`int`): Block identifier (0 to num_layers-1).
+        Returns:
+            - gradients (:obj:`Dict[str, torch.Tensor]`): Dictionary of parameter names and gradients.
+        """
+        if block_id < 0 or block_id >= len(self.blocks):
+            raise ValueError(f"Block ID {block_id} out of range. Available blocks: 0-{len(self.blocks)-1}")
+        block = self.blocks[block_id]
+        if not hasattr(block, 'feed_forward'):
+            raise ValueError(f"Block {block_id} doesn't have feed_forward layer")
+        if not hasattr(block.feed_forward, 'shared_expert') or block.feed_forward.shared_expert is None:
+            raise ValueError(f"Block {block_id} doesn't have shared expert")
+        gradients = {}
+        shared_expert = block.feed_forward.shared_expert
+        for name, param in shared_expert.named_parameters():
+            if param.grad is not None:
+                gradients[f"shared_expert.{name}"] = param.grad.clone()
+            else:
+                gradients[f"shared_expert.{name}"] = None
+        return gradients
+
+    def get_expert_gradients_for_last_block(self) -> List[torch.Tensor]:
+        """
+        Overview:
+            Retrieve flattened parameter gradients of all experts from the last transformer block.
+            Used for per-expert gradient conflict analysis.
+
+        Returns:
+            - gradients (:obj:`List[torch.Tensor]`): List of flattened gradient tensors, one per expert.
+        """
+        if len(self.blocks) == 0:
+            return []
+        last_block = self.blocks[-1]
+        gradients = []
+        if not hasattr(last_block, 'feed_forward'):
+            return gradients
+        feed_forward = last_block.feed_forward
+        if hasattr(feed_forward, 'experts') and feed_forward.experts is not None:
+            for expert_idx, expert in enumerate(feed_forward.experts):
+                expert_gradients = []
+                for name, param in expert.named_parameters():
+                    if param.grad is not None:
+                        expert_gradients.append(param.grad.clone().view(-1))
+                    else:
+                        expert_gradients.append(torch.zeros_like(param, device=param.device).view(-1))
+                gradients.append(torch.cat(expert_gradients, dim=0))
+        return gradients
+
+    def get_block_before_moe_gradients(self) -> Optional[torch.Tensor]:
+        """
+        Overview:
+            Retrieve gradients of the layer before MoE (ln2 output) in the last transformer block.
+            Used for gradient conflict analysis in multi-task learning.
+
+        Returns:
+            - gradients (:obj:`Optional[torch.Tensor]`): Gradient tensor, or None if hook not registered.
+        """
+        if len(self.blocks) == 0:
+            return None
+        return getattr(self.blocks[-1], 'block_before_moe_grad', None)
+
+    def get_last_shared_expert_gradients(self) -> torch.Tensor:
+        """
+        Overview:
+            Retrieve flattened parameter gradients from the shared expert in the last transformer block.
+
+        Returns:
+            - gradients (:obj:`torch.Tensor`): Flattened tensor of all shared expert gradients; empty tensor if none.
+        """
+        if len(self.blocks) == 0:
+            return []
+        last_block = self.blocks[-1]
+        if not hasattr(last_block, 'feed_forward') or not hasattr(last_block.feed_forward, 'shared_expert'):
+            return []
+        shared_expert = last_block.feed_forward.shared_expert
+        if shared_expert is None:
+            return []
+        shared_expert_gradients = []
+        for name, param in shared_expert.named_parameters():
+            if param.grad is not None:
+                shared_expert_gradients.append(param.grad.clone().view(-1))
+            else:
+                shared_expert_gradients.append(torch.zeros_like(param, device=param.device).view(-1))
+        return torch.cat(shared_expert_gradients, dim=0)
+
+    def get_last_block_expert_selection_stats(self) -> dict:
+        """
+        Overview:
+            Retrieve MoE expert selection statistics from the last transformer block.
+        Returns:
+            - stats (:obj:`dict`): Dictionary containing expert selection statistics.
+        """
+        if len(self.blocks) == 0:
+            return {}
+        last_block = self.blocks[-1]
+        if hasattr(last_block, 'feed_forward') and hasattr(last_block.feed_forward, 'get_expert_selection_stats'):
+            return last_block.feed_forward.get_expert_selection_stats()
+        return {}
+
+    def reset_last_block_expert_selection_stats(self) -> None:
+        """
+        Overview:
+            Reset MoE expert selection statistics for the last transformer block.
+        """
+        if len(self.blocks) == 0:
+            return
+        last_block = self.blocks[-1]
+        if hasattr(last_block, 'feed_forward') and hasattr(last_block.feed_forward, 'reset_expert_selection_stats'):
+            last_block.feed_forward.reset_expert_selection_stats()
+
 
 class Block(nn.Module):
     """
@@ -529,16 +683,24 @@ def __init__(self, config: TransformerConfig) -> None:
                 _maybe_wrap_linear(nn.Linear(4 * config.embed_dim, config.embed_dim), config, "feed_forward"),
                 nn.Dropout(config.resid_pdrop),
             )
+        self.config = config
+        self.block_before_moe_grad = None
 
     def forward(self, x: torch.Tensor, past_keys_values: Optional[KeysValues] = None,
-                valid_context_lengths: Optional[torch.Tensor] = None) -> torch.Tensor:
+                valid_context_lengths: Optional[torch.Tensor] = None,
+                is_last_block: bool = False, task_id: int = 0) -> torch.Tensor:
         """
         Overview:
             Performs the forward pass of the Transformer block.
+
         Arguments:
             - x (:obj:`torch.Tensor`): Input tensor of shape (batch_size, seq_length, embed_dim).
             - past_keys_values (:obj:`Optional[KeysValues]`): Precomputed keys and values for faster generation.
             - valid_context_lengths (:obj:`Optional[torch.Tensor]`): Valid lengths of context for masking.
+            - is_last_block (:obj:`bool`): Whether this is the last block; only the last block registers
+                gradient hook for block-before-MoE and passes task_id to MoE for expert selection stats.
+            - task_id (:obj:`int`): Current task ID, passed to MoE forward for expert selection statistics.
+
         Returns:
             - torch.Tensor: Output tensor of shape (batch_size, seq_length, embed_dim).
         """
@@ -548,8 +710,27 @@ def forward(self, x: torch.Tensor, past_keys_values: Optional[KeysValues] = None
             ff_output = self.feed_forward(self.ln2(x))
             x = self.gate2(x, ff_output)
         else:
+            # x = x + attn_output
+            # x = x + self.feed_forward(self.ln2(x))
+            
+            
             x = x + attn_output
-            x = x + self.feed_forward(self.ln2(x))
+            block_before_moe = self.ln2(x)
+            # Register gradient hook on last block for gradient conflict analysis
+            if self.training and is_last_block:
+                self.block_before_moe_grad = None
+
+                def grad_hook(grad):
+                    self.block_before_moe_grad = grad.clone()
+                    return None
+
+                block_before_moe.register_hook(grad_hook)
+
+            if is_last_block and getattr(self.config, 'multiplication_moe_in_transformer', False) and hasattr(self.feed_forward, 'forward'):
+                x = x + self.feed_forward(block_before_moe, task_id=task_id)
+            else:
+                x = x + self.feed_forward(block_before_moe)
+            
         return x
 
 
diff --git a/lzero/model/unizero_world_models/world_model_multitask.py b/lzero/model/unizero_world_models/world_model_multitask.py
index 836a463c7..bf756b696 100644
--- a/lzero/model/unizero_world_models/world_model_multitask.py
+++ b/lzero/model/unizero_world_models/world_model_multitask.py
@@ -255,6 +255,9 @@ def __init__(self, config: TransformerConfig, tokenizer: Tokenizer) -> None:
 
         self.reanalyze_phase = False
         self._rank = get_rank()
+        
+        # For gradient conflict analysis: populated by obs_embeddings register_hook during backward
+        self.obs_embeddings_grad = None
 
     def _scale_grad(self, grad: torch.Tensor) -> torch.Tensor:
         """
@@ -799,7 +802,7 @@ def _transformer_pass(
             ]
             return torch.cat(x, dim=0)
         else:
-            return self.transformer(sequences, past_keys_values, valid_context_lengths=valid_context_lengths)
+            return self.transformer(sequences, past_keys_values, valid_context_lengths=valid_context_lengths,task_id=task_id)
 
     @torch.no_grad()
     def reset_for_initial_inference(self, obs_act_dict: dict, task_id: int = 0) -> Tuple[WorldModelOutput, torch.Tensor]:
@@ -1590,7 +1593,10 @@ def gather_and_plot(
     def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar_transform_handle=None, task_id = 0, **kwargs: Any) -> LossWithIntermediateLosses:
         # Encode observations into latent state representations
         obs_embeddings = self.tokenizer.encode_to_obs_embeddings(batch['observations'], task_id=task_id)
-
+        
+        # Register hook to capture obs_embeddings gradients for gradient conflict analysis
+        obs_embeddings.register_hook(lambda grad: setattr(self, 'obs_embeddings_grad', grad))
+   
         if self.analysis_tsne:
             # =========== tsne analysis ===========
             if not obs_embeddings.is_cuda:
@@ -1605,7 +1611,7 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
         if self.analysis_dormant_ratio_weight_rank:
             self._analysis_step_counter += 1
             self.do_analysis = (
-                self.analysis_dormant_ratio_weight_rank          # 总开关
+                self.analysis_dormant_ratio_weight_rank
                 and self._analysis_step_counter % self.analysis_dormant_ratio_interval == 0
             )
 
diff --git a/lzero/policy/scaling_transform.py b/lzero/policy/scaling_transform.py
index ac3305332..fc8093e1f 100644
--- a/lzero/policy/scaling_transform.py
+++ b/lzero/policy/scaling_transform.py
@@ -167,8 +167,11 @@ def phi_transform(
 
     # --- 4. Stack indices / probs and scatter ------------------------------
     # Clamp high_idx to handle the edge case where x is exactly max_bound
-    idx   = torch.stack([low_idx_long,
-                         torch.clamp(high_idx, max=size - 1)], dim=-1)  # (*x, 2)
+   
+    low_idx_clamped  = torch.clamp(low_idx_long, min=0, max=size - 1)
+    high_idx_clamped = torch.clamp(high_idx, min=0, max=size - 1)
+    idx   = torch.stack([low_idx_clamped, high_idx_clamped], dim=-1)  # (*x, 2)    
+    
     prob  = torch.stack([p_low, p_high], dim=-1)                        # (*x, 2)
 
     target = torch.zeros(*x.shape, size,
diff --git a/lzero/policy/unizero_multitask.py b/lzero/policy/unizero_multitask.py
index 64a52d8af..baf9366da 100644
--- a/lzero/policy/unizero_multitask.py
+++ b/lzero/policy/unizero_multitask.py
@@ -15,7 +15,9 @@
                           select_action, to_torch_float_tensor)
 from lzero.policy.unizero import UniZeroPolicy, scale_module_weights_vectorized
 
-from .utils import configure_optimizers_nanogpt, initialize_zeros_batch
+from .utils import configure_optimizers_nanogpt, initialize_zeros_batch, compute_gradient_conflict_distributed, log_gradient_conflict_heatmaps_distributed_fast
+# Gradient conflict monitoring: analyze encoder/MoE gradient conflicts in multi-task learning.
+# Imports: compute_gradient_conflict_distributed, log_gradient_conflict_heatmaps_distributed_fast
 
 # Please replace the path with the actual location of your LibMTL library.
 sys.path.append('/path/to/your/LibMTL')
@@ -25,7 +27,6 @@
 from LibMTL.weighting.moco_fast_mem_eff import MoCoCfg
 from LibMTL.weighting.MoCo_unizero import MoCo as GradCorrect
 
-
 def generate_task_loss_dict(multi_task_losses: List[Union[torch.Tensor, float]], task_name_template: str, task_id: int) -> Dict[str, float]:
     """
     Overview:
@@ -160,7 +161,9 @@ class UniZeroMTPolicy(UniZeroPolicy):
         by addressing the limitations of MuZero-style algorithms, particularly in environments requiring the
         capture of long-term dependencies. More details can be found at: https://arxiv.org/abs/2406.10667.
     """
-
+    
+    
+       
     # The default_config for UniZero multi-task policy.
     config = dict(
         type='unizero_multitask',
@@ -470,7 +473,15 @@ class UniZeroMTPolicy(UniZeroPolicy):
             decay=int(1e5),
         ),
     )
-
+    def __init__(self, cfg, model = None, enable_field = None):
+        super().__init__(cfg, model, enable_field)
+        # Gradient conflict monitoring: global step counter for TensorBoard logging
+        self.step = 0
+        # Frequency (in train iters) to force log gradient conflict scalars
+        self.save_freq = 200
+        # Whether MoE is enabled in the model (affects gradient conflict collection)
+        self.use_moe = False
+        
     def default_model(self) -> Tuple[str, List[str]]:
         """
         Overview:
@@ -1086,6 +1097,123 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
             e_rank_sim_norm_multi_task.append(e_rank_sim_norm)
 
 
+        
+        # ==================== Gradient Conflict Monitoring ====================
+        # Before main backward: run per-task backward to collect gradients from encoder,
+        # block-before-MoE, shared expert, and individual experts. Compute cosine-similarity-based
+        # conflict scores and log to TensorBoard. NOTE: obs_embeddings_grad, before_moe_grad may be None.
+        self._optimizer_world_model.zero_grad()
+        self._learn_model.world_model.tokenizer.encoder[0].grad = None
+        multi_gpu = dist.is_initialized() and self._cfg.multi_gpu
+        rank = dist.get_rank() if multi_gpu else 0
+
+        self.log_conflict_var = True
+        self.log_conflict_matrix = True
+        if self.step % self.save_freq == 0:
+            self.log_conflict_var = True
+
+        if self.log_conflict_var:
+            matrix_dict = {}
+            num_experts = self._learn_model.world_model.transformer.num_experts
+
+
+            local_task_num = len(losses_list)
+            local_encoder_grad_list = []
+            local_before_moe_grad_list = []
+            local_shared_expert_grad_list = []
+            local_last_block_expert_grad_list = [[] for _ in range(num_experts)] 
+            
+            print(f'Rank {rank} collecting gradients') 
+            gradient_conflict_log_dict = {}
+
+            for i in range(local_task_num):
+                # Clear gradients before each computation to ensure independence
+                self._optimizer_world_model.zero_grad()
+                # Per-task backward to collect gradients (retain_graph for later main backward)
+                losses_list[i].backward(retain_graph=True)
+                # obs_embeddings_grad is set by world_model register_hook; may be None before first backward
+                local_encoder_grad_list.append(self._learn_model.world_model.obs_embeddings_grad.view(-1).detach().clone())
+
+                # Gradients of ln2(x) before MoE in last transformer block (set by Block grad_hook)
+                before_moe_grad = self._learn_model.world_model.transformer.get_block_before_moe_gradients()
+                local_before_moe_grad_list.append(before_moe_grad.view(-1).detach().clone())
+
+                # Get gradients of the shared expert in the last block
+                if self._learn_model.world_model.transformer.shared_expert > 0:
+                    shared_expert_grad_for_last_task = self._learn_model.world_model.transformer.get_last_shared_expert_gradients()
+                    local_shared_expert_grad_list.append(shared_expert_grad_for_last_task)
+                    
+                # Compute gradient conflicts of experts in the last block
+                if num_experts>0:
+                    last_block_expert_grad_list = self._learn_model.world_model.transformer.get_expert_gradients_for_last_block()
+                    for j in range(num_experts):
+                        local_last_block_expert_grad_list[j].append(last_block_expert_grad_list[j])
+
+
+            
+            print(f'Rank {rank} computing gradient conflicts')
+                
+            # Clear shared parameter gradients to avoid accumulation
+            self._optimizer_world_model.zero_grad()
+            
+            print(f'Rank {rank} computing attention gradient conflicts')
+            # 1. Compute gradient conflicts after attention and before MOE
+            local_before_moe_grad_list = torch.stack(local_before_moe_grad_list, dim=0)  # (local_task_num, grad_dim)
+            before_moe_grad_conflict_ddp=compute_gradient_conflict_distributed(local_before_moe_grad_list, device=self._cfg.device)
+            gradient_conflict_log_dict['avg_before_moe_grad_conflict'] = before_moe_grad_conflict_ddp.avg_conflict_score if before_moe_grad_conflict_ddp is not None else 0
+            gradient_conflict_log_dict['max_before_moe_grad_conflict'] = before_moe_grad_conflict_ddp.max_conflict_score if before_moe_grad_conflict_ddp is not None else 0
+            if self.log_conflict_matrix and before_moe_grad_conflict_ddp is not None :
+                matrix_dict['before_moe_grad_conflict_matrix'] = before_moe_grad_conflict_ddp.cosine_similarity_matrix
+
+            print(f'Rank {rank} computing encoder gradient conflicts')
+            # 2. Compute gradient conflicts of encoder
+            local_encoder_grad_list = torch.stack(local_encoder_grad_list, dim=0)  # (local_task_num, grad_dim)
+            encoder_grad_conflict_ddp=compute_gradient_conflict_distributed(local_encoder_grad_list, device=self._cfg.device)
+            gradient_conflict_log_dict['avg_encoder_grad_conflict'] = encoder_grad_conflict_ddp.avg_conflict_score if encoder_grad_conflict_ddp is not None else 0
+            gradient_conflict_log_dict['max_encoder_grad_conflict'] = encoder_grad_conflict_ddp.max_conflict_score if encoder_grad_conflict_ddp is not None else 0
+            if self.log_conflict_matrix and encoder_grad_conflict_ddp is not None:
+                matrix_dict['encoder_grad_conflict_matrix']=encoder_grad_conflict_ddp.cosine_similarity_matrix 
+
+
+            print(f'Rank {rank} computing shared expert gradient conflicts')
+            # 3. If shared expert exists, compute gradient conflicts on shared expert
+            if self._learn_model.world_model.transformer.shared_expert > 0:
+                local_shared_expert_grad_list = torch.stack(local_shared_expert_grad_list, dim=0)
+                shared_expert_grad_conflict = compute_gradient_conflict_distributed(local_shared_expert_grad_list, device=self._cfg.device) if len(local_shared_expert_grad_list) > 0 else None
+                gradient_conflict_log_dict['avg_shared_expert_grad_conflict'] = shared_expert_grad_conflict.avg_conflict_score if shared_expert_grad_conflict is not None else 0
+                gradient_conflict_log_dict['max_shared_expert_grad_conflict'] = shared_expert_grad_conflict.max_conflict_score if shared_expert_grad_conflict is not None else 0
+
+               
+                if self.log_conflict_matrix and shared_expert_grad_conflict is not None:
+                    matrix_dict['shared_expert_grad_conflict_matrix'] = shared_expert_grad_conflict.cosine_similarity_matrix
+
+            # 4. Gradient conflicts of experts in the last block
+            last_block_expert_grad_conflict_ddp_list = []
+            if num_experts > 0:
+                for i in range(num_experts):
+                    # Stack gradients of the last block experts across tasks
+                    local_last_block_expert_grad_list[i] = torch.stack(local_last_block_expert_grad_list[i], dim=0)
+                    # Compute gradient conflicts of each expert
+                    expert_conflict = compute_gradient_conflict_distributed(local_last_block_expert_grad_list[i], device=self._cfg.device)
+                    last_block_expert_grad_conflict_ddp_list.append(expert_conflict)
+                    gradient_conflict_log_dict[f'avg_expert_{i}_grad_conflict'] = expert_conflict.avg_conflict_score if expert_conflict is not None else 0
+                    gradient_conflict_log_dict[f'max_expert_{i}_grad_conflict'] = expert_conflict.max_conflict_score if expert_conflict is not None else 0
+
+                    if self.log_conflict_matrix and expert_conflict is not None:
+                        matrix_dict[f'expert_{i}_grad_conflict_matrix'] = expert_conflict.cosine_similarity_matrix
+
+                all_moe_gradient = torch.cat(local_last_block_expert_grad_list, dim=1)
+                if self._learn_model.world_model.transformer.shared_expert > 0:
+                    all_moe_gradient = torch.cat((local_shared_expert_grad_list, all_moe_gradient), dim=1)
+                all_moe_gradient_ddp = compute_gradient_conflict_distributed(all_moe_gradient, device=self._cfg.device)
+
+                gradient_conflict_log_dict['avg_moe_layer_grad_conflict'] = all_moe_gradient_ddp.avg_conflict_score if all_moe_gradient_ddp is not None else 0
+                gradient_conflict_log_dict['max_moe_layer_grad_conflict'] = all_moe_gradient_ddp.max_conflict_score if all_moe_gradient_ddp is not None else 0
+                if self.log_conflict_matrix and all_moe_gradient_ddp is not None:
+                    matrix_dict['max_moe_layer_grad_conflict_matrix'] = all_moe_gradient_ddp.cosine_similarity_matrix
+
+        self._optimizer_world_model.zero_grad()
+
         # ==================== Integrate norm monitoring logic ====================
         norm_log_dict = {}
         # Check if monitoring frequency is reached
@@ -1282,6 +1410,20 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
             'total_grad_norm_before_clip_wm': total_grad_norm_before_clip_wm.item(),
         }
 
+        # Get tb_logger from self.logger or config (cfg.policy.logger)
+        tb_logger = getattr(self, 'logger', None) or getattr(self._cfg, 'logger', None)
+        if self.log_conflict_matrix and tb_logger is not None:
+            # Log gradient conflict heatmaps to TensorBoard (converted to list for distributed processing)
+            matrix_list = list(matrix_dict.items())
+            log_gradient_conflict_heatmaps_distributed_fast(tb_logger, matrix_list, self.step)
+            
+        if self.log_conflict_var and tb_logger is not None:
+            # Log gradient conflict scalar metrics to TensorBoard
+            for key, value in gradient_conflict_log_dict.items():
+                tb_logger.add_scalar(f'gradient_conflict/{key}', value, self.step) 
+            
+        
+        
         # ==================== START: Add new log items ====================
         if self.use_adaptive_entropy_weight:
             return_log_dict['adaptive_alpha'] = current_alpha.item()
@@ -1333,7 +1475,8 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
         if norm_log_dict:
             return_log_dict.update(norm_log_dict)
 
-        # Return the final loss dictionary.
+        # Increment step counter for gradient conflict logging
+        self.step += 1
         return return_log_dict
 
     def monitor_weights_and_grads(self, model: torch.nn.Module) -> None:
@@ -1399,7 +1542,9 @@ def _monitor_vars_learn(self, num_tasks: int = 2) -> List[str]:
             'cur_lr_world_model',
             'weighted_total_loss',
             'total_grad_norm_before_clip_wm',
-
+            'avg_encoder_grad_conflict',
+            'avg_before_moe_grad_conflict',
+            'avg_shared_expert_grad_conflict',
             'adaptive_alpha',
             "adaptive_target_entropy_ratio",
             'final_alpha_loss',
diff --git a/lzero/policy/utils.py b/lzero/policy/utils.py
index d2bdaa6c8..d2b68360f 100644
--- a/lzero/policy/utils.py
+++ b/lzero/policy/utils.py
@@ -800,3 +800,531 @@ def mz_network_output_unpack(network_output: Dict) -> Tuple:
     value = network_output.value  # shape: (batch_size, support_support_size)
     policy_logits = network_output.policy_logits  # shape: (batch_size, action_space_size)
     return latent_state, reward, value, policy_logits
+
+
+# ==================== #=============================
+import torch.distributed as dist
+
+# ==================== Gradient Conflict Matrix Visualization Module =============================
+"""
+Overview:
+    Gradient conflict matrix visualization and computation for multi-task learning analysis.
+    Provides heatmap generation and distributed logging for gradient conflict analysis.
+
+Interfaces:
+    - _get_or_create_figure: Get or create reusable matplotlib figure
+    - _fast_tensor_heatmap: Generate optimized heatmap tensor from conflict matrix
+    - log_gradient_conflict_heatmaps_distributed_fast: High-performance distributed heatmap logging
+    - compute_gradient_conflict_distributed: Distributed gradient conflict computation (cosine similarity)
+"""
+
+# Pre-import matplotlib module to avoid repeated import overhead
+import matplotlib
+matplotlib.use('Agg')
+
+# Global figure cache
+_GLOBAL_FIG_CACHE = None
+_GLOBAL_AX_CACHE = None
+
+def _get_or_create_figure(figsize=(8, 6)):
+    """
+    Overview:
+        Get or create reusable matplotlib figure for memory efficiency.
+    Arguments:
+        - figsize (:obj:`tuple`): Figure size as (width, height), default is (8, 6).
+    Returns:
+        - fig (:obj:`matplotlib.figure.Figure`): Matplotlib figure object.
+        - ax (:obj:`matplotlib.axes.Axes`): Matplotlib axes object.
+    Examples:
+        >>> fig, ax = _get_or_create_figure((10, 8))
+        >>> ax.plot([1, 2, 3], [4, 5, 6])
+    """
+    global _GLOBAL_FIG_CACHE, _GLOBAL_AX_CACHE
+    if _GLOBAL_FIG_CACHE is None:
+        _GLOBAL_FIG_CACHE, _GLOBAL_AX_CACHE = plt.subplots(figsize=figsize)
+    return _GLOBAL_FIG_CACHE, _GLOBAL_AX_CACHE
+
+def _fast_tensor_heatmap(matrix_np, tag):
+    """
+    Overview:
+        Generate optimized heatmap tensor with performance enhancements by skipping text annotations
+        and removing diagonal elements for better visualization.
+    Arguments:
+        - matrix_np (:obj:`numpy.ndarray`): Input matrix for heatmap generation.
+        - tag (:obj:`str`): Tag label for the heatmap title.
+    Returns:
+        - img_tensor (:obj:`torch.Tensor`): RGB image tensor with shape :math:`(3, H, W)`.
+    Shapes:
+        - matrix_np: :math:`(N, M)` where N and M are matrix dimensions.
+        - img_tensor: :math:`(3, H, W)` where H and W are image dimensions.
+    Examples:
+        >>> matrix = np.random.randn(5, 5)
+        >>> heatmap_tensor = _fast_tensor_heatmap(matrix, "conflict_matrix")
+        >>> print(heatmap_tensor.shape)  # torch.Size([3, height, width])
+    """
+    # Copy matrix to avoid modifying original data
+    matrix_no_diag = matrix_np.copy()
+
+    # Set diagonal to 0 for better conflict visualization (self-similarity is always 1)
+    if matrix_no_diag.shape[0] == matrix_no_diag.shape[1]:
+        np.fill_diagonal(matrix_no_diag, 0)
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+    # Blues colormap with range [-0.2, 0.2] for conflict scores
+    im = ax.imshow(matrix_no_diag, cmap='Blues', vmin=-0.2, vmax=0.2)
+    ax.set_title(f'{tag}', fontsize=12)
+    
+    # Add value annotations only for small matrices (avoid O(n^2) overhead)
+    if matrix_no_diag.size <= 64:
+        for row in range(matrix_no_diag.shape[0]):
+            for col in range(matrix_no_diag.shape[1]):
+                if row != col:
+                    value = matrix_no_diag[row, col]
+                    text_color = "white" if value > 0.5 else "black"
+                    ax.text(col, row, f'{value:.2f}',
+                           ha="center", va="center", color=text_color, fontsize=8)
+    
+    fig.canvas.draw()
+    try:
+        if hasattr(fig.canvas, 'buffer_rgba'):
+            buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
+            buf = buf.reshape(fig.canvas.get_width_height()[::-1] + (4,))
+            img_tensor = torch.from_numpy(buf[:, :, :3]).permute(2, 0, 1).float() / 255.0
+        elif hasattr(fig.canvas, 'tostring_rgb'):
+            buf = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+            buf = buf.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+            img_tensor = torch.from_numpy(buf).permute(2, 0, 1).float() / 255.0
+        else:
+            try:
+                from PIL import Image
+                import io
+                buf = io.BytesIO()
+                fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+                buf.seek(0)
+                pil_img = Image.open(buf).convert('RGB')
+                img_array = np.array(pil_img)
+                img_tensor = torch.from_numpy(img_array).permute(2, 0, 1).float() / 255.0
+            except Exception:
+                h, w = matrix_no_diag.shape
+                # Fallback: simple upsampling for visualization
+                img_tensor = torch.zeros(3, h*50, w*50)
+                img_tensor[2] = torch.from_numpy(matrix_no_diag).repeat_interleave(50, 0).repeat_interleave(50, 1)
+    except Exception:
+        h, w = matrix_no_diag.shape
+        # Fallback: simple upsampling for visualization
+        img_tensor = torch.zeros(3, h*50, w*50)
+        img_tensor[2] = torch.from_numpy(matrix_no_diag).repeat_interleave(50, 0).repeat_interleave(50, 1)
+    finally:
+        plt.close(fig)
+    
+    return img_tensor
+
+
+def log_gradient_conflict_heatmaps_distributed_fast(tb_logger, matrix_list, step):
+    """
+    Overview:
+        High-performance distributed heatmap processing with optimizations for reduced latency.
+        Key optimizations: figure reuse, text annotation skipping for large matrices,
+        conditional barriers, and robust error recovery.
+
+    Arguments:
+        - tb_logger (:obj:`SummaryWriter`): TensorBoard logger for heatmap logging.
+        - matrix_list (:obj:`list`): List of (tag, matrix) tuples; tag is string, matrix is tensor.
+        - step (:obj:`int`): Global training step for logging.
+
+    Returns:
+        - None: No return value; performs logging only.
+
+    Shapes:
+        - matrix_list[i][1]: :math:`(N, N)` cosine similarity matrix, N = number of tasks.
+    Examples:
+        >>> import torch
+        >>> from torch.utils.tensorboard import SummaryWriter
+        >>> tb_logger = SummaryWriter()
+        >>> matrices = [("task1", torch.randn(5, 5)), ("task2", torch.randn(3, 3))]
+        >>> log_gradient_conflict_heatmaps_distributed_fast(tb_logger, matrices, 100)
+    """
+    if not matrix_list:
+        return
+    
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    
+    try:
+        # Each GPU processes its own subset of matrices
+        processed_any = False
+        for i in range(rank, len(matrix_list), world_size):
+            tag, matrix = matrix_list[i]
+            if matrix is not None and matrix.numel() > 0:
+                matrix_np = matrix.detach().cpu().numpy()
+                
+                img_tensor = _fast_tensor_heatmap(matrix_np, tag)
+                tb_logger.add_image(f'gradient_conflict_matrix/{tag}', img_tensor, global_step=step)
+                processed_any = True
+        
+        # Barrier only when needed; rank 0 always participates to avoid deadlock
+        if processed_any or rank == 0:
+            dist.barrier()
+        
+    except Exception as e:
+        print(f"Rank {rank}: Error in optimized heatmap logging: {e}")
+        try:
+            dist.barrier()
+        except:
+            pass
+
+# ==================== Gradient Conflict Computation Module =============================
+
+
+
+def example_usage():
+    """
+    Overview:
+        Example usage demonstration for gradient conflict analysis computation.
+        Generates sample gradients and computes conflict analysis results including average conflict score,
+        maximum conflict score, number of conflicting gradient pairs, average conflict intensity,
+        gradient norms, and cosine similarity matrix.
+    Arguments:
+        - None: Function generates sample gradients internally for demonstration.
+    Returns:
+        - None: Function prints results to console without return values.
+    Examples:
+        >>> example_usage()
+        # Output:
+        # Gradient Conflict Analysis Results:
+        # Average conflict score: 0.1234
+        # Maximum conflict score: 0.5678
+        # Number of conflicting pairs: 3
+        # Average conflict intensity: 0.2345
+        # Gradient norms: [tensor1, tensor2, tensor3]
+        # Cosine similarity matrix:
+        # tensor([[1.0000, -0.1234, 0.5678],
+        #         [-0.1234, 1.0000, -0.3456],
+        #         [0.5678, -0.3456, 1.0000]])
+    """
+    torch.manual_seed(42)
+    gradients = [
+        torch.randn(100),
+        torch.randn(100),
+        torch.randn(100),
+    ]
+    conflicts = compute_gradient_conflicts(gradients)
+    print("Gradient Conflict Analysis Results:")
+    print(f"Average conflict score: {conflicts['avg_conflict_score']:.4f}")
+    print(f"Max conflict score: {conflicts['max_conflict_score']:.4f}")
+    print(f"Num conflicting pairs: {conflicts['num_conflicting_pairs']}")
+    print(f"Avg conflict intensity: {conflicts['avg_conflict_intensity']:.4f}")
+    print(f"Gradient norms: {conflicts['gradient_norms']}")
+    print("\nCosine similarity matrix:")
+    print(conflicts['cosine_similarity_matrix'])
+
+
+
+def compute_gradient_conflicts(gradients: List[torch.Tensor]) -> dict:
+    """
+    Overview:
+        Compute conflicts between multiple gradients using CUDA-optimized vectorized operations.
+        Calculates cosine similarity matrix and derives conflict scores for gradient analysis.
+    Arguments:
+        - gradients (:obj:`List[torch.Tensor]`): List of gradient tensors with identical shapes.
+    Returns:
+        - result (:obj:`dict`): Dictionary containing conflict analysis results with keys:
+                               'avg_conflict_score', 'max_conflict_score', 'min_conflict_score',
+                               and 'cosine_similarity_matrix'.
+    Shapes:
+        - gradients[i]: :math:`(D_1, D_2, ..., D_n)` where all gradients have identical dimensions.
+        - cosine_similarity_matrix: :math:`(N, N)` where N is the number of gradients.
+    Examples:
+        >>> import torch
+        >>> gradients = [torch.randn(100), torch.randn(100), torch.randn(100)]
+        >>> conflicts = compute_gradient_conflicts(gradients)
+        >>> print(f"Average conflict: {conflicts['avg_conflict_score']:.4f}")
+        >>> print(f"Similarity matrix shape: {conflicts['cosine_similarity_matrix'].shape}")
+    """
+    n_gradients = len(gradients)
+    
+    # No conflict with single gradient
+    if n_gradients <= 1:
+        device = gradients[0].device if gradients else torch.device('cuda')
+        return EasyDict({
+            'avg_conflict_score': 0.0, 
+            'max_conflict_score': 0.0, 
+            'min_conflict_score': 0.0,
+            'cosine_similarity_matrix': torch.zeros(1, 1, device=device)
+        })
+    
+    assert all(g.shape == gradients[0].shape for g in gradients), "All gradients must have same shape"
+    
+    device = gradients[0].device
+    
+    # Vectorized: stack and normalize all gradients
+    stacked_grads = torch.stack([g.flatten() for g in gradients])
+    normalized_grads = F.normalize(stacked_grads, p=2, dim=1)
+    
+    # Compute full pairwise cosine similarity matrix in one step
+    cosine_sim_matrix = torch.mm(normalized_grads, normalized_grads.t())
+    
+    # Exclude diagonal elements
+    mask = ~torch.eye(n_gradients, device=device, dtype=torch.bool)
+    conflict_scores = -cosine_sim_matrix[mask]
+    
+    return EasyDict({
+        'avg_conflict_score': conflict_scores.mean().item(),
+        'max_conflict_score': conflict_scores.max().item(),
+        'min_conflict_score': conflict_scores.min().item(),
+        'cosine_similarity_matrix': cosine_sim_matrix
+    })
+ 
+ 
+def compute_gradient_conflict_distributed(local_grads, multi_gpu=True, device=0):
+    """
+    Overview:
+        Distributed gradient conflict computation with hierarchical aggregation optimization.
+        Uses layered preprocessing, NCCL communication, and vectorized cosine-similarity computation.
+
+    Arguments:
+        - local_grads (:obj:`torch.Tensor`): Local gradient tensor, shape (L, D) with L tasks, D dim.
+        - multi_gpu (:obj:`bool`, optional): Whether to use multi-GPU distributed mode. Default True.
+        - device (:obj:`Union[int, str, torch.device]`, optional): Device for computation. Default 0.
+    Returns:
+        - gradient_conflict (:obj:`dict`): Dictionary containing conflict analysis results identical
+                                          across all ranks, including 'avg_conflict_score',
+                                          'max_conflict_score', 'min_conflict_score', and
+                                          'cosine_similarity_matrix'.
+    Shapes:
+        - local_grads: :math:`(L, D)` where L is local task number and D is encoder gradient dimension.
+        - cosine_similarity_matrix: :math:`(N, N)` where N is total number of valid gradients across all ranks.
+    Examples:
+        >>> import torch
+        >>> import torch.distributed as dist
+        >>> local_grads = torch.randn(5, 128)  # 5 local tasks, 128-dim gradients
+        >>> conflicts = compute_gradient_conflict_distributed(local_grads, multi_gpu=True, device=0)
+        >>> print(f"Average conflict: {conflicts['avg_conflict_score']:.4f}")
+    """
+    if not multi_gpu:
+        # Single-GPU mode: use optimized single-node version
+        norms = torch.norm(local_grads, dim=1)
+        valid_grads = local_grads[norms > 1e-8]
+        if valid_grads.shape[0] <= 1:
+            device = valid_grads.device
+            return EasyDict({
+                'avg_conflict_score': 0.0, 
+                'max_conflict_score': 0.0, 
+                'min_conflict_score': 0.0,
+                'cosine_similarity_matrix': torch.zeros(1, 1, device=device)
+            })
+        
+        device = valid_grads.device
+        normalized = F.normalize(valid_grads, p=2, dim=1)
+        similarity = torch.mm(normalized, normalized.t())
+        mask = ~torch.eye(valid_grads.shape[0], device=device, dtype=torch.bool)
+        conflicts = -similarity[mask]
+        return EasyDict({
+            'avg_conflict_score': conflicts.mean().item(),
+            'max_conflict_score': conflicts.max().item(),
+            'min_conflict_score': conflicts.min().item(),
+            'cosine_similarity_matrix': similarity
+        })
+    
+    # Multi-GPU distributed: layered aggregation
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    device = torch.device(f'{device}')
+    
+    # Layer 1: local preprocessing
+    norms = torch.norm(local_grads, dim=1)
+    valid_grads = local_grads[norms > 1e-8]
+    local_normalized = F.normalize(valid_grads, p=2, dim=1)
+
+    # Gather valid gradient counts from all ranks
+    valid_count = torch.tensor(valid_grads.shape[0], device=device)
+    valid_counts = [torch.tensor(0, device=device) for _ in range(world_size)]
+    dist.all_gather(valid_counts, valid_count)
+    
+    total_valid = sum(v.item() for v in valid_counts)
+    if total_valid <= 1:
+        return EasyDict({
+            'avg_conflict_score': 0.0, 
+            'max_conflict_score': 0.0, 
+            'min_conflict_score': 0.0,
+            'cosine_similarity_matrix': torch.zeros(1, 1, device=device)
+        })
+    
+    # Pad to same size for all_gather
+    max_valid = max(v.item() for v in valid_counts)
+    if valid_grads.shape[0] < max_valid:
+        pad_size = max_valid - valid_grads.shape[0]
+        pad_tensor = torch.zeros(pad_size, valid_grads.shape[1], device=device, dtype=valid_grads.dtype)
+        local_normalized = torch.cat([local_normalized, pad_tensor], dim=0)
+    
+    # Layer 2: NCCL all_gather
+    gathered_normalized = [torch.empty_like(local_normalized) for _ in range(world_size)]
+    dist.all_gather(gathered_normalized, local_normalized)
+    
+    # if rank == 0:
+        # Layer 3: rebuild valid normalized gradients and compute conflicts
+    all_valid_normalized = []
+    for i, count in enumerate(valid_counts):
+        if count > 0:
+            all_valid_normalized.append(gathered_normalized[i][:count.item()])
+    
+    if len(all_valid_normalized) == 0:
+        return EasyDict({
+            'avg_conflict_score': 0.0, 
+            'max_conflict_score': 0.0, 
+            'min_conflict_score': 0.0,
+            'cosine_similarity_matrix': torch.zeros(1, 1, device=device)
+        })
+    
+    all_normalized = torch.cat(all_valid_normalized, dim=0)
+    
+    # Vectorized: single matrix multiply instead of O(n^2) loops
+    similarity = torch.mm(all_normalized, all_normalized.t())
+    mask = ~torch.eye(similarity.shape[0], device=device, dtype=torch.bool)
+    conflicts = -similarity[mask]
+    
+    return EasyDict({
+        'avg_conflict_score': conflicts.mean().item(),
+        'max_conflict_score': conflicts.max().item(),
+        'min_conflict_score': conflicts.min().item(),
+        'cosine_similarity_matrix': similarity
+    })
+
+def compute_gradient_conflicts_batch(gradient_groups: Dict[str, torch.Tensor], device=0) -> Dict[str, dict]:
+    """
+    Overview:
+        Batch computation of gradient conflicts for multiple gradient groups to reduce
+        distributed communication overhead through optimized data aggregation.
+    Arguments:
+        - gradient_groups (:obj:`Dict[str, torch.Tensor]`): Dictionary mapping group names to
+                                                           local gradient tensors.
+        - device (:obj:`int`, optional): Device index for tensor operations. Default is 0.
+    Returns:
+        - results (:obj:`Dict[str, dict]`): Dictionary mapping group names to conflict analysis
+                                          results, each containing 'avg_conflict_score',
+                                          'max_conflict_score', 'min_conflict_score', and
+                                          'cosine_similarity_matrix'.
+    Shapes:
+        - gradient_groups[group_name]: :math:`(L, D)` where L is local task number and D is gradient dimension.
+        - results[group_name]['cosine_similarity_matrix']: :math:`(N, N)` where N is total valid gradients for the group.
+    Examples:
+        >>> import torch
+        >>> gradient_groups = {
+        ...     "encoder": torch.randn(5, 128),
+        ...     "decoder": torch.randn(3, 64)
+        ... }
+        >>> results = compute_gradient_conflicts_batch(gradient_groups, device=0)
+        >>> print(f"Encoder conflicts: {results['encoder']['avg_conflict_score']:.4f}")
+        >>> print(f"Decoder conflicts: {results['decoder']['avg_conflict_score']:.4f}")
+    """
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    world_size = dist.get_world_size() if dist.is_initialized() else 1
+    
+    results = {}
+    
+    if world_size == 1:
+        # Single-GPU mode
+        for group_name, local_grads in gradient_groups.items():
+            if local_grads.numel() == 0:
+                results[group_name] = EasyDict({'avg_conflict_score': 0.0})
+                continue
+                
+            # Filter zero gradients
+            norms = torch.norm(local_grads, dim=1)
+            valid_mask = norms > 1e-8
+            local_grads_filtered = local_grads[valid_mask]
+            
+            if local_grads_filtered.shape[0] <= 1:
+                results[group_name] = EasyDict({
+                    'avg_conflict_score': 0.0, 
+                    'max_conflict_score': 0.0, 
+                    'min_conflict_score': 0.0,
+                    'cosine_similarity_matrix': torch.zeros(1, 1, device=device)
+                })
+            else:
+                grad_list = [local_grads_filtered[i] for i in range(local_grads_filtered.shape[0])]
+                results[group_name] = compute_gradient_conflicts(grad_list)
+        return results
+    
+    # Multi-GPU: collect all gradient groups at once
+    local_filtered_groups = {}
+    local_valid_counts = {}
+    
+    for group_name, local_grads in gradient_groups.items():
+        if local_grads.numel() == 0:
+            local_filtered_groups[group_name] = torch.empty(0, 0, device=device)
+            local_valid_counts[group_name] = 0
+            continue
+            
+        norms = torch.norm(local_grads, dim=1)
+        valid_mask = norms > 1e-8
+        filtered = local_grads[valid_mask]
+        local_filtered_groups[group_name] = filtered
+        local_valid_counts[group_name] = filtered.shape[0]
+    
+    # Collect valid sample counts from all ranks
+    all_valid_counts = [None for _ in range(world_size)]
+    dist.all_gather_object(all_valid_counts, local_valid_counts)
+    
+    # Compute per-group maximum task counts for padding
+    max_counts = {}
+    for group_name in gradient_groups.keys():
+        counts = [counts_dict.get(group_name, 0) for counts_dict in all_valid_counts]
+        max_counts[group_name] = max(counts) if counts else 0
+    
+    # Pad local groups to the maximum count and prepare for communication
+    local_padded_groups = {}
+    for group_name, filtered_grads in local_filtered_groups.items():
+        max_count = max_counts[group_name]
+        if max_count == 0:
+            local_padded_groups[group_name] = torch.empty(0, 0)
+            continue
+            
+        if filtered_grads.shape[0] < max_count:
+            if filtered_grads.numel() > 0:
+                pad_size = max_count - filtered_grads.shape[0]
+                grad_dim = filtered_grads.shape[1]
+                pad_tensor = torch.zeros(pad_size, grad_dim, device=device)
+                padded = torch.cat([filtered_grads, pad_tensor], dim=0)
+            else:
+                grad_dim = gradient_groups[group_name].shape[1] if gradient_groups[group_name].numel() > 0 else 1
+                padded = torch.zeros(max_count, grad_dim, device=device)
+        else:
+            padded = filtered_grads
+            
+        local_padded_groups[group_name] = padded.cpu()
+    
+    # Gather all padded gradient groups from all ranks in one shot
+    all_gradient_groups = [None for _ in range(world_size)]
+    dist.all_gather_object(all_gradient_groups, local_padded_groups)
+    
+    if rank == 0:
+        for group_name in gradient_groups.keys():
+            valid_grad_list = []
+            for rank_idx, rank_data in enumerate(all_gradient_groups):
+                if group_name in rank_data:
+                    valid_count = all_valid_counts[rank_idx].get(group_name, 0)
+                    if valid_count > 0:
+                        tensor_valid = rank_data[group_name][:valid_count, :].to(device)
+                        valid_grad_list.append(tensor_valid)
+            
+            if len(valid_grad_list) == 0:
+                results[group_name] = EasyDict({'avg_conflict_score': 0.0})
+            else:
+                all_grads = torch.cat(valid_grad_list, dim=0)
+                if all_grads.shape[0] <= 1:
+                    results[group_name] = EasyDict({'avg_conflict_score': 0.0})
+                else:
+                    grad_list = [all_grads[i] for i in range(all_grads.shape[0])]
+                    results[group_name] = compute_gradient_conflicts(grad_list)
+    else:
+        results = None
+    
+    # Broadcast final results to all ranks
+    results_list = [results]
+    dist.broadcast_object_list(results_list, src=0)
+    return results_list[0]
+
+
+if __name__ == "__main__":
+    example_usage()
diff --git a/zoo/atari/config/README.md b/zoo/atari/config/README.md
new file mode 100644
index 000000000..23b4f669a
--- /dev/null
+++ b/zoo/atari/config/README.md
@@ -0,0 +1,49 @@
+## One Model for All Tasks: Leveraging Efficient World Models in Multi-Task Planning — MoE Gradient Conflict Experiment Reproduction
+
+This README describes how to reproduce the **MoE gradient conflict experiments** from the paper **"One Model for All Tasks: Leveraging Efficient World Models in Multi-Task Planning"**.  
+Paper link: `https://arxiv.org/pdf/2509.07945`
+
+The two main configs in this folder are used to reproduce **Figures 16–20**:
+
+| Config file | Model type | Role in the paper |
+| --- | --- | --- |
+| `atari_unizero_nomoe_multitask_segment_ddp_config.py` | Multi-task UniZero, **no MoE** | **"MLP / dense Transformer"** baseline |
+| `atari_unizero_moe_multitask_segment_ddp_config.py` | ScaleZero-style model with **MoE backbone** | **"MoE Transformer"** variant |
+
+Both scripts train on a multi-task Atari benchmark and log the statistics needed to rebuild  
+**Figure 16, 17, 18, 19, 20** (gradient conflicts, expert-selection entropy, expert usage heatmaps, task-wise Wasserstein distances, etc.).
+
+> **Figure 18 note**: the y-axis uses **log-scaled gradient conflict values**.  
+> When plotting, first apply a log transform (e.g., log10) to the raw conflict metrics.
+
+### How to run
+
+From the LightZero project root:
+
+```bash
+# Non-MoE baseline
+torchrun --nproc_per_node=4 zoo/atari/config/atari_unizero_nomoe_multitask_segment_ddp_config.py
+
+# MoE version
+torchrun --nproc_per_node=4 zoo/atari/config/atari_unizero_moe_multitask_segment_ddp_config.py
+```
+
+After running both scripts, you can follow the descriptions in the paper appendix (especially Sections E.1–E.2) plus the logged statistics to reproduce Figures 16–20.
+
+### Reference Figures
+
+**moe_expert_selection_wasserstein_distance**
+
+<img src="../../../assets/moe_expert_selection_wasserstein_distance.png" alt="moe_expert_selection_wasserstein_distance" width="800"/>
+
+**moe_expert_selection_js_divergence**
+
+<img src="../../../assets/moe_expert_selection_js_divergence.png" alt="moe_expert_selection_js_divergence" width="800"/>
+
+**gradient_conflict_comparison_moe_vs_nomoe**
+
+<img src="../../../assets/gradient_conflict_comparison_moe_vs_nomoe.png" alt="gradient_conflict_comparison_moe_vs_nomoe" width="800"/>
+
+**expert_selection_heatmaps**
+
+<img src="../../../assets/expert_selection_heatmaps.jpg" alt="expert_selection_heatmaps" width="800"/>
diff --git a/zoo/atari/config/README_zh.md b/zoo/atari/config/README_zh.md
new file mode 100644
index 000000000..b4cf29b8b
--- /dev/null
+++ b/zoo/atari/config/README_zh.md
@@ -0,0 +1,48 @@
+## One Model for All Tasks: Leveraging Efficient World Models in Multi-Task Planning — MoE 梯度冲突实验复现
+
+本文档为论文 **《One Model for All Tasks: Leveraging Efficient World Models in Multi-Task Planning》** 中 **MoE 梯度冲突相关实验** 的复现说明。  
+论文链接：`https://arxiv.org/pdf/2509.07945`
+
+本目录下两个主要脚本用于复现 **Figure 16–20** 相关实验：
+
+| 配置文件 | 模型类型 | 在论文中的角色 |
+| --- | --- | --- |
+| `atari_unizero_nomoe_multitask_segment_ddp_config.py` | 多任务 UniZero，**无 MoE** | 作为 **“MLP / Dense Transformer” 基线** |
+| `atari_unizero_moe_multitask_segment_ddp_config.py` | 使用 **MoE backbone** 的 ScaleZero 风格模型 | 作为 **“MoE Transformer” 版本** |
+
+两份脚本都会在 Atari 多任务基准上训练，并记录复现以下图所需的统计量：  
+**Figure 16, 17, 18, 19, 20**（梯度冲突、专家选择熵、专家利用热力图、任务间 Wasserstein 距离等）。
+
+> **Figure 18 提醒**：图中纵轴为 **梯度冲突的 log 值**，画图时请先对原始冲突数值取对数（如 log10）。
+
+### 运行示例
+
+在 LightZero 根目录下执行：
+
+```bash
+# 非 MoE 基线
+torchrun --nproc_per_node=4 zoo/atari/config/atari_unizero_nomoe_multitask_segment_ddp_config.py
+
+# MoE 版本
+torchrun --nproc_per_node=4 zoo/atari/config/atari_unizero_moe_multitask_segment_ddp_config.py
+```
+
+跑完这两个脚本后，结合日志与论文附录（尤其 E.1–E.2）的说明，即可复现 Figure 16–20。
+
+### 参考图片
+
+**moe_expert_selection_wasserstein_distance**
+
+<img src="../../../assets/moe_expert_selection_wasserstein_distance.png" alt="moe_expert_selection_wasserstein_distance" width="800"/>
+
+**moe_expert_selection_js_divergence**
+
+<img src="../../../assets/moe_expert_selection_js_divergence.png" alt="moe_expert_selection_js_divergence" width="800"/>
+
+**gradient_conflict_comparison_moe_vs_nomoe**
+
+<img src="../../../assets/gradient_conflict_comparison_moe_vs_nomoe.png" alt="gradient_conflict_comparison_moe_vs_nomoe" width="800"/>
+
+**expert_selection_heatmaps**
+
+<img src="../../../assets/expert_selection_heatmaps.jpg" alt="expert_selection_heatmaps" width="800"/>
diff --git a/zoo/atari/config/atari_unizero_moe_multitask_segment_ddp_config.py b/zoo/atari/config/atari_unizero_moe_multitask_segment_ddp_config.py
new file mode 100644
index 000000000..af8464d38
--- /dev/null
+++ b/zoo/atari/config/atari_unizero_moe_multitask_segment_ddp_config.py
@@ -0,0 +1,326 @@
+from easydict import EasyDict
+import math
+from typing import List, Tuple, Any, Dict, Union
+
+# -------------------------------------------------
+# 1. Refactored compute_batch_config
+# -------------------------------------------------
+def compute_batch_config(
+    env_id_list: List[str],
+    effective_batch_size: int,
+    gpu_num: int = 8,
+    max_micro_batch_one_gpu: int = 400,
+) -> Tuple[List[int], int]:
+    """
+    Overview:
+        Calculate the micro-batch size for each environment and the number of gradient accumulation steps
+        to approach a target effective batch size across multiple GPUs and environments.
+
+    Arguments:
+        - env_id_list (:obj:`List[str]`): A list of environment IDs for all tasks.
+        - effective_batch_size (:obj:`int`): The target global batch size for one backward pass.
+        - gpu_num (:obj:`int`): The number of GPUs actually used. Defaults to 8.
+        - max_micro_batch_one_gpu (:obj:`int`): The maximum micro-batch size a single GPU can handle. Defaults to 400.
+
+    Returns:
+        - batch_sizes (:obj:`List[int]`): A list of micro-batch sizes for each environment.
+        - grad_acc_steps (:obj:`int`): The number of gradient accumulation steps.
+    """
+    n_env = len(env_id_list)
+    # Number of environments that each GPU needs to handle simultaneously.
+    envs_per_gpu = max(1, math.ceil(n_env / gpu_num))
+    # Reduce the micro-batch limit if multiple environments share one GPU.
+    max_micro_batch = max(1, max_micro_batch_one_gpu // envs_per_gpu)
+
+    # First, calculate a candidate micro-batch by distributing the effective batch size evenly.
+    candidate = max(1, effective_batch_size // n_env)
+    micro_batch = min(candidate, max_micro_batch)
+
+    # Gradient accumulation steps = ceil(global_batch / (micro_batch * n_env)).
+    grad_acc_steps = max(1, math.ceil(effective_batch_size / (micro_batch * n_env)))
+
+    # Fine-tune the micro-batch downwards to ensure:
+    # micro_batch * n_env * grad_acc_steps <= effective_batch_size
+    # This aims to get as close as possible to the target without exceeding it.
+    while micro_batch * n_env * grad_acc_steps > effective_batch_size:
+        micro_batch -= 1
+        if micro_batch == 0:  # Defensive check, should not happen in theory.
+            micro_batch = 1
+            break
+
+    batch_sizes = [micro_batch] * n_env
+
+    # --- Debug Information --- #
+    real_total_batch_size = micro_batch * n_env * grad_acc_steps
+    print(
+        f"[BatchConfig] Envs={n_env}, TargetTotalBS={effective_batch_size}, "
+        f"MicroBS={micro_batch}, GradAccSteps={grad_acc_steps}, RealTotalBS={real_total_batch_size}"
+    )
+
+    return batch_sizes, grad_acc_steps
+
+def create_config(
+        env_id: str, action_space_size: int, collector_env_num: int, evaluator_env_num: int, n_episode: int,
+        num_simulations: int, reanalyze_ratio: float, batch_size: int, num_unroll_steps: int,
+        infer_context_length: int, norm_type: str, buffer_reanalyze_freq: float, reanalyze_batch_size: int,
+        reanalyze_partition: float, num_segments: int, total_batch_size: int, num_layers: int
+) -> EasyDict:
+    """
+    Overview:
+        Creates the main configuration structure for a single training task.
+
+    Arguments:
+        - env_id (:obj:`str`): The environment ID.
+        - action_space_size (:obj:`int`): The size of the action space.
+        - collector_env_num (:obj:`int`): Number of environments for data collection.
+        - evaluator_env_num (:obj:`int`): Number of environments for evaluation.
+        - n_episode (:obj:`int`): Number of episodes to run for evaluation.
+        - num_simulations (:obj:`int`): Number of simulations in MCTS.
+        - reanalyze_ratio (:obj:`float`): The ratio of reanalyzed samples in a batch.
+        - batch_size (:obj:`int`): The batch size for training.
+        - num_unroll_steps (:obj:`int`): The number of steps to unroll the model dynamics.
+        - infer_context_length (:obj:`int`): The context length for inference.
+        - norm_type (:obj:`str`): The type of normalization layer to use (e.g., 'LN').
+        - buffer_reanalyze_freq (:obj:`float`): Frequency of reanalyzing the replay buffer.
+        - reanalyze_batch_size (:obj:`int`): Batch size for reanalysis.
+        - reanalyze_partition (:obj:`float`): Partition ratio for reanalysis.
+        - num_segments (:obj:`int`): Number of segments for data collection.
+        - total_batch_size (:obj:`int`): The total effective batch size.
+        - num_layers (:obj:`int`): Number of layers in the transformer model.
+
+    Returns:
+        - (:obj:`EasyDict`): A configuration object.
+    """
+    return EasyDict(dict(
+        env=dict(
+            stop_value=int(1e6),
+            env_id=env_id,
+            observation_shape=(3, 64, 64),
+            gray_scale=False,
+            collector_env_num=collector_env_num,
+            evaluator_env_num=evaluator_env_num,
+            n_evaluator_episode=evaluator_env_num,
+            manager=dict(shared_memory=False),
+            full_action_space=True,
+        ),
+        policy=dict(
+            multi_gpu=True,
+            only_use_moco_stats=False,
+            use_moco=False,
+            moco_version="v1",
+            total_task_num=len(env_id_list),
+            task_num=len(env_id_list),
+            task_id=0,  # This will be overridden for each task
+            model=dict(
+                observation_shape=(3, 64, 64),
+                action_space_size=action_space_size,
+                norm_type=norm_type,
+                num_res_blocks=2,
+                num_channels=256,
+                num_layers=num_layers,
+                world_model_cfg=dict(
+                    norm_type=norm_type,
+                    action_space_size=action_space_size,
+                    num_layers=num_layers,
+                    num_heads=8,
+                    embed_dim=768,
+                    env_num=len(env_id_list),
+                    task_num=len(env_id_list),
+                    max_blocks=num_unroll_steps,
+                    max_tokens=2 * num_unroll_steps,
+                    context_length=2 * infer_context_length,
+                    final_norm_option_in_obs_head='LayerNorm',
+                    final_norm_option_in_encoder='LayerNorm',
+                    predict_latent_loss_type='mse',
+                    encoder_type='vit',
+                    device='cuda',
+                    game_segment_length=20,
+                    # MoE: multiplication-based MoE in transformer, 8 experts, 2 per token, 1 shared expert
+                    use_normal_head=True,
+                    use_softmoe_head=False,
+                    use_moe_head=False,
+                    num_experts_in_moe_head=1,
+                    moe_in_transformer=False,
+                    multiplication_moe_in_transformer=True,
+                    n_shared_experts=1,
+                    num_experts_per_tok=2,
+                    num_experts_of_moe_in_transformer=8,
+                    moe_use_lora=False,
+                ),
+            ),
+            device='cuda',
+            game_segment_length=20,
+            update_per_collect=80,  # Corresponds to replay_ratio=0.5 for 8 games (20*8*0.5=80)
+            learning_rate=0.0001,
+            weight_decay=1e-2,
+            batch_size=batch_size,
+            num_unroll_steps=num_unroll_steps,
+            num_segments=num_segments,
+            num_simulations=num_simulations,
+            reanalyze_ratio=reanalyze_ratio,
+            n_episode=n_episode,
+            total_batch_size=total_batch_size,
+            collector_env_num=collector_env_num,
+            evaluator_env_num=evaluator_env_num,
+            buffer_reanalyze_freq=buffer_reanalyze_freq,
+            reanalyze_batch_size=reanalyze_batch_size,
+            reanalyze_partition=reanalyze_partition,
+            replay_buffer_size=int(5e5),
+            eval_freq=int(1e4),
+        ),
+    ))
+
+def generate_configs(
+        env_id_list: List[str], action_space_size: int, collector_env_num: int, n_episode: int,
+        evaluator_env_num: int, num_simulations: int, reanalyze_ratio: float, batch_size: List[int],
+        num_unroll_steps: int, infer_context_length: int, norm_type: str, seed: int,
+        buffer_reanalyze_freq: float, reanalyze_batch_size: int, reanalyze_partition: float,
+        num_segments: int, total_batch_size: int, num_layers: int
+) -> List[List[Union[int, List[EasyDict]]]]:
+    """
+    Overview:
+        Generates a list of configurations for all specified tasks.
+
+    Arguments:
+        (See arguments for `create_config` function)
+        - seed (:obj:`int`): The random seed for the experiment.
+
+    Returns:
+        - (:obj:`List[List[Union[int, List[EasyDict]]]]`): A list where each element contains a task_id
+          and its corresponding configuration objects.
+    """
+    configs = []
+
+    # --- Experiment Name Template ---
+    benchmark_tag = "data_unizero_mt"
+    model_tag = f"vit_nlayer{num_layers}_tbs{total_batch_size}"
+    exp_name_prefix = f'{benchmark_tag}/atari_{len(env_id_list)}games_{model_tag}_seed{seed}/'
+
+    for task_id, env_id in enumerate(env_id_list):
+        config = create_config(
+            env_id, action_space_size, collector_env_num, evaluator_env_num, n_episode, num_simulations,
+            reanalyze_ratio, batch_size, num_unroll_steps, infer_context_length, norm_type,
+            buffer_reanalyze_freq, reanalyze_batch_size, reanalyze_partition, num_segments, total_batch_size, num_layers
+        )
+        config.policy.task_id = task_id
+        # Correctly extract the game name from 'ALE/GameName-v5' format.
+        game_name = env_id.split('/')[1].split('-')[0]
+        config.exp_name = exp_name_prefix + f"{game_name}_seed{seed}"
+        configs.append([task_id, [config, create_env_manager()]])
+    return configs
+
+def create_env_manager() -> EasyDict:
+    """
+    Overview:
+        Creates the environment manager configuration, specifying the types of environment,
+        policy, and their import paths.
+
+    Returns:
+        - (:obj:`EasyDict`): A configuration object for the environment manager.
+    """
+    return EasyDict(dict(
+        env=dict(
+            type='atari_lightzero',
+            import_names=['zoo.atari.envs.atari_lightzero_env'],
+        ),
+        env_manager=dict(type='subprocess'),
+        policy=dict(
+            type='unizero_multitask',
+            import_names=['lzero.policy.unizero_multitask'],
+        ),
+    ))
+
+if __name__ == "__main__":
+    """
+    Overview:
+        This script should be executed with <nproc_per_node> GPUs for distributed training.
+
+        Example launch commands:
+
+        export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+        cd /path/to/your/project/
+
+        torchrun --nproc_per_node=4 /mnt/shared-storage-user/puyuan/code/LightZero/zoo/atari/config/atari_unizero_moe_multitask_segment_ddp_config.py
+    """
+    from lzero.entry import train_unizero_multitask_segment_ddp
+    from ding.utils import DDPContext
+    import torch.distributed as dist
+    import os
+
+    # ==================== Main Experiment Settings ====================
+    num_games = 8  # Options: 3, 8, 26
+    num_layers = 1  # Transformer depth (reduced for faster iteration)
+    action_space_size = 18
+    collector_env_num = 8
+    num_segments = 8
+    n_episode = 8
+    evaluator_env_num = 3
+    num_simulations = 25  # MCTS simulations per step
+    max_env_step = int(5e6)
+    reanalyze_ratio = 0.0
+
+    # ==================== Environment Configuration ====================
+    if num_games == 3:
+        env_id_list = ['ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5']
+    elif num_games == 8:
+        env_id_list = [
+            'ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5', 'ALE/Boxing-v5',
+            'ALE/Alien-v5', 'ALE/ChopperCommand-v5', 'ALE/Hero-v5', 'ALE/RoadRunner-v5',
+        ]
+    elif num_games == 26:
+        env_id_list = [
+            'ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5', 'ALE/Boxing-v5',
+            'ALE/Alien-v5', 'ALE/ChopperCommand-v5', 'ALE/Hero-v5', 'ALE/RoadRunner-v5',
+            'ALE/Amidar-v5', 'ALE/Assault-v5', 'ALE/Asterix-v5', 'ALE/BankHeist-v5',
+            'ALE/BattleZone-v5', 'ALE/CrazyClimber-v5', 'ALE/DemonAttack-v5', 'ALE/Freeway-v5',
+            'ALE/Frostbite-v5', 'ALE/Gopher-v5', 'ALE/Jamesbond-v5', 'ALE/Kangaroo-v5',
+            'ALE/Krull-v5', 'ALE/KungFuMaster-v5', 'ALE/PrivateEye-v5', 'ALE/UpNDown-v5',
+            'ALE/Qbert-v5', 'ALE/Breakout-v5',
+        ]
+    else:
+        raise ValueError(f"Unsupported number of environments: {num_games}")
+
+    # ==================== Batch Size Calculation ====================
+    if len(env_id_list) == 8:
+        if num_layers in [1, 4]:
+            effective_batch_size = 10
+        elif num_layers == 8:
+            effective_batch_size = 10
+        else:
+            effective_batch_size = 10
+    elif len(env_id_list) == 26:
+        effective_batch_size = 512
+    elif len(env_id_list) == 3:
+        effective_batch_size = 10  # Reduced for debugging; use 512 for full training
+    else:
+        raise ValueError(f"Batch size not configured for {len(env_id_list)} environments.")
+
+    batch_sizes, grad_acc_steps = compute_batch_config(env_id_list, effective_batch_size, gpu_num=4)
+    total_batch_size = effective_batch_size
+
+    # ==================== Model and Training Settings ====================
+    num_unroll_steps = 10
+    infer_context_length = 4
+    norm_type = 'LN'
+    buffer_reanalyze_freq = 1 / 100000000  # Effectively disable buffer reanalyze
+    reanalyze_batch_size = 160
+    reanalyze_partition = 0.75
+
+    # ==================== Training Loop ====================
+    # Set NCCL timeout to prevent watchdog hang due to unbalanced data collection speeds
+    os.environ.setdefault('NCCL_TIMEOUT', '3600')  # 60 minutes in seconds
+    os.environ.setdefault('NCCL_BLOCKING_WAIT', '1')
+
+    for seed in [0]:
+        configs = generate_configs(
+            env_id_list, action_space_size, collector_env_num, n_episode, evaluator_env_num,
+            num_simulations, reanalyze_ratio, batch_sizes, num_unroll_steps, infer_context_length,
+            norm_type, seed, buffer_reanalyze_freq, reanalyze_batch_size, reanalyze_partition,
+            num_segments, total_batch_size, num_layers
+        )
+
+        with DDPContext():
+            train_unizero_multitask_segment_ddp(configs, seed=seed, max_env_step=max_env_step, benchmark_name="atari")
+            print(f"Seed: {seed} training finished!")
+            if dist.is_initialized():
+                dist.destroy_process_group()
diff --git a/zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py b/zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py
index 08636286c..57def1014 100644
--- a/zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py
+++ b/zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py
@@ -311,3 +311,330 @@ def create_env_manager() -> EasyDict:
             print(f"Seed: {seed} training finished!")
             if dist.is_initialized():
                 dist.destroy_process_group()
+
+from easydict import EasyDict
+import math
+from typing import List, Tuple, Any, Dict, Union
+
+# -------------------------------------------------
+# 1. Refactored compute_batch_config
+# -------------------------------------------------
+def compute_batch_config(
+    env_id_list: List[str],
+    effective_batch_size: int,
+    gpu_num: int = 8,
+    max_micro_batch_one_gpu: int = 400,
+) -> Tuple[List[int], int]:
+    """
+    Overview:
+        Calculate the micro-batch size for each environment and the number of gradient accumulation steps
+        to approach a target effective batch size across multiple GPUs and environments.
+
+    Arguments:
+        - env_id_list (:obj:`List[str]`): A list of environment IDs for all tasks.
+        - effective_batch_size (:obj:`int`): The target global batch size for one backward pass.
+        - gpu_num (:obj:`int`): The number of GPUs actually used. Defaults to 8.
+        - max_micro_batch_one_gpu (:obj:`int`): The maximum micro-batch size a single GPU can handle. Defaults to 400.
+
+    Returns:
+        - batch_sizes (:obj:`List[int]`): A list of micro-batch sizes for each environment.
+        - grad_acc_steps (:obj:`int`): The number of gradient accumulation steps.
+    """
+    n_env = len(env_id_list)
+    # Number of environments that each GPU needs to handle simultaneously.
+    envs_per_gpu = max(1, math.ceil(n_env / gpu_num))
+    # Reduce the micro-batch limit if multiple environments share one GPU.
+    max_micro_batch = max(1, max_micro_batch_one_gpu // envs_per_gpu)
+
+    # First, calculate a candidate micro-batch by distributing the effective batch size evenly.
+    candidate = max(1, effective_batch_size // n_env)
+    micro_batch = min(candidate, max_micro_batch)
+
+    # Gradient accumulation steps = ceil(global_batch / (micro_batch * n_env)).
+    grad_acc_steps = max(1, math.ceil(effective_batch_size / (micro_batch * n_env)))
+
+    # Fine-tune the micro-batch downwards to ensure:
+    # micro_batch * n_env * grad_acc_steps <= effective_batch_size
+    # This aims to get as close as possible to the target without exceeding it.
+    while micro_batch * n_env * grad_acc_steps > effective_batch_size:
+        micro_batch -= 1
+        if micro_batch == 0:  # Defensive check, should not happen in theory.
+            micro_batch = 1
+            break
+
+    batch_sizes = [micro_batch] * n_env
+
+    # --- Debug Information --- #
+    real_total_batch_size = micro_batch * n_env * grad_acc_steps
+    print(
+        f"[BatchConfig] Envs={n_env}, TargetTotalBS={effective_batch_size}, "
+        f"MicroBS={micro_batch}, GradAccSteps={grad_acc_steps}, RealTotalBS={real_total_batch_size}"
+    )
+
+    return batch_sizes, grad_acc_steps
+
+def create_config(
+        env_id: str, action_space_size: int, collector_env_num: int, evaluator_env_num: int, n_episode: int,
+        num_simulations: int, reanalyze_ratio: float, batch_size: int, num_unroll_steps: int,
+        infer_context_length: int, norm_type: str, buffer_reanalyze_freq: float, reanalyze_batch_size: int,
+        reanalyze_partition: float, num_segments: int, total_batch_size: int, num_layers: int
+) -> EasyDict:
+    """
+    Overview:
+        Creates the main configuration structure for a single training task.
+
+    Arguments:
+        - env_id (:obj:`str`): The environment ID.
+        - action_space_size (:obj:`int`): The size of the action space.
+        - collector_env_num (:obj:`int`): Number of environments for data collection.
+        - evaluator_env_num (:obj:`int`): Number of environments for evaluation.
+        - n_episode (:obj:`int`): Number of episodes to run for evaluation.
+        - num_simulations (:obj:`int`): Number of simulations in MCTS.
+        - reanalyze_ratio (:obj:`float`): The ratio of reanalyzed samples in a batch.
+        - batch_size (:obj:`int`): The batch size for training.
+        - num_unroll_steps (:obj:`int`): The number of steps to unroll the model dynamics.
+        - infer_context_length (:obj:`int`): The context length for inference.
+        - norm_type (:obj:`str`): The type of normalization layer to use (e.g., 'LN').
+        - buffer_reanalyze_freq (:obj:`float`): Frequency of reanalyzing the replay buffer.
+        - reanalyze_batch_size (:obj:`int`): Batch size for reanalysis.
+        - reanalyze_partition (:obj:`float`): Partition ratio for reanalysis.
+        - num_segments (:obj:`int`): Number of segments for data collection.
+        - total_batch_size (:obj:`int`): The total effective batch size.
+        - num_layers (:obj:`int`): Number of layers in the transformer model.
+
+    Returns:
+        - (:obj:`EasyDict`): A configuration object.
+    """
+    return EasyDict(dict(
+        env=dict(
+            stop_value=int(1e6),
+            env_id=env_id,
+            observation_shape=(3, 64, 64),
+            gray_scale=False,
+            collector_env_num=collector_env_num,
+            evaluator_env_num=evaluator_env_num,
+            n_evaluator_episode=evaluator_env_num,
+            manager=dict(shared_memory=False),
+            full_action_space=True,
+        ),
+        policy=dict(
+            multi_gpu=True,
+            only_use_moco_stats=False,
+            use_moco=False,
+            moco_version="v1",
+            total_task_num=len(env_id_list),
+            task_num=len(env_id_list),
+            task_id=0,  # This will be overridden for each task
+            model=dict(
+                observation_shape=(3, 64, 64),
+                action_space_size=action_space_size,
+                norm_type=norm_type,
+                num_res_blocks=2,
+                num_channels=256,
+                num_layers=num_layers,
+                world_model_cfg=dict(
+                    norm_type=norm_type,
+                    action_space_size=action_space_size,
+                    num_layers=num_layers,
+                    num_heads=8,
+                    embed_dim=768,
+                    env_num=len(env_id_list),
+                    task_num=len(env_id_list),
+                    max_blocks=num_unroll_steps,
+                    max_tokens=2 * num_unroll_steps,
+                    context_length=2 * infer_context_length,
+                    final_norm_option_in_obs_head='LayerNorm',
+                    final_norm_option_in_encoder='LayerNorm',
+                    predict_latent_loss_type='mse',
+                    encoder_type='vit',
+                    device='cuda',
+                    game_segment_length=20,
+                    # MoE: multiplication-based MoE in transformer, 8 experts, 2 per token, 1 shared expert
+                    use_normal_head=True,
+                    use_softmoe_head=False,
+                    use_moe_head=False,
+                    num_experts_in_moe_head=1,
+                    moe_in_transformer=False,
+                    multiplication_moe_in_transformer=True,
+                    n_shared_experts=1,
+                    num_experts_per_tok=2,
+                    num_experts_of_moe_in_transformer=8,
+                    moe_use_lora=False,
+                ),
+            ),
+            device='cuda',
+            game_segment_length=20,
+            update_per_collect=80,  # Corresponds to replay_ratio=0.5 for 8 games (20*8*0.5=80)
+            learning_rate=0.0001,
+            weight_decay=1e-2,
+            batch_size=batch_size,
+            num_unroll_steps=num_unroll_steps,
+            num_segments=num_segments,
+            num_simulations=num_simulations,
+            reanalyze_ratio=reanalyze_ratio,
+            n_episode=n_episode,
+            total_batch_size=total_batch_size,
+            collector_env_num=collector_env_num,
+            evaluator_env_num=evaluator_env_num,
+            buffer_reanalyze_freq=buffer_reanalyze_freq,
+            reanalyze_batch_size=reanalyze_batch_size,
+            reanalyze_partition=reanalyze_partition,
+            replay_buffer_size=int(5e5),
+            eval_freq=int(1e4),
+        ),
+    ))
+
+def generate_configs(
+        env_id_list: List[str], action_space_size: int, collector_env_num: int, n_episode: int,
+        evaluator_env_num: int, num_simulations: int, reanalyze_ratio: float, batch_size: List[int],
+        num_unroll_steps: int, infer_context_length: int, norm_type: str, seed: int,
+        buffer_reanalyze_freq: float, reanalyze_batch_size: int, reanalyze_partition: float,
+        num_segments: int, total_batch_size: int, num_layers: int
+) -> List[List[Union[int, List[EasyDict]]]]:
+    """
+    Overview:
+        Generates a list of configurations for all specified tasks.
+
+    Arguments:
+        (See arguments for `create_config` function)
+        - seed (:obj:`int`): The random seed for the experiment.
+
+    Returns:
+        - (:obj:`List[List[Union[int, List[EasyDict]]]]`): A list where each element contains a task_id
+          and its corresponding configuration objects.
+    """
+    configs = []
+
+    # --- Experiment Name Template ---
+    benchmark_tag = "data_unizero_mt"
+    model_tag = f"vit_nlayer{num_layers}_tbs{total_batch_size}"
+    exp_name_prefix = f'{benchmark_tag}/atari_{len(env_id_list)}games_{model_tag}_seed{seed}/'
+
+    for task_id, env_id in enumerate(env_id_list):
+        config = create_config(
+            env_id, action_space_size, collector_env_num, evaluator_env_num, n_episode, num_simulations,
+            reanalyze_ratio, batch_size, num_unroll_steps, infer_context_length, norm_type,
+            buffer_reanalyze_freq, reanalyze_batch_size, reanalyze_partition, num_segments, total_batch_size, num_layers
+        )
+        config.policy.task_id = task_id
+        # Correctly extract the game name from 'ALE/GameName-v5' format.
+        game_name = env_id.split('/')[1].split('-')[0]
+        config.exp_name = exp_name_prefix + f"{game_name}_seed{seed}"
+        configs.append([task_id, [config, create_env_manager()]])
+    return configs
+
+def create_env_manager() -> EasyDict:
+    """
+    Overview:
+        Creates the environment manager configuration, specifying the types of environment,
+        policy, and their import paths.
+
+    Returns:
+        - (:obj:`EasyDict`): A configuration object for the environment manager.
+    """
+    return EasyDict(dict(
+        env=dict(
+            type='atari_lightzero',
+            import_names=['zoo.atari.envs.atari_lightzero_env'],
+        ),
+        env_manager=dict(type='subprocess'),
+        policy=dict(
+            type='unizero_multitask',
+            import_names=['lzero.policy.unizero_multitask'],
+        ),
+    ))
+
+if __name__ == "__main__":
+    """
+    Overview:
+        This script should be executed with <nproc_per_node> GPUs for distributed training.
+
+        Example launch commands:
+
+        export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+        cd /path/to/your/project/
+
+        torchrun --nproc_per_node=4 /mnt/shared-storage-user/puyuan/code/LightZero/zoo/atari/config/atari_unizero_moe_multitask_segment_ddp_config.py
+    """
+    from lzero.entry import train_unizero_multitask_segment_ddp
+    from ding.utils import DDPContext
+    import torch.distributed as dist
+    import os
+
+    # ==================== Main Experiment Settings ====================
+    num_games = 8  # Options: 3, 8, 26
+    num_layers = 1  # Transformer depth (reduced for faster iteration)
+    action_space_size = 18
+    collector_env_num = 8
+    num_segments = 8
+    n_episode = 8
+    evaluator_env_num = 3
+    num_simulations = 25  # MCTS simulations per step
+    max_env_step = int(5e6)
+    reanalyze_ratio = 0.0
+
+    # ==================== Environment Configuration ====================
+    if num_games == 3:
+        env_id_list = ['ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5']
+    elif num_games == 8:
+        env_id_list = [
+            'ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5', 'ALE/Boxing-v5',
+            'ALE/Alien-v5', 'ALE/ChopperCommand-v5', 'ALE/Hero-v5', 'ALE/RoadRunner-v5',
+        ]
+    elif num_games == 26:
+        env_id_list = [
+            'ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5', 'ALE/Boxing-v5',
+            'ALE/Alien-v5', 'ALE/ChopperCommand-v5', 'ALE/Hero-v5', 'ALE/RoadRunner-v5',
+            'ALE/Amidar-v5', 'ALE/Assault-v5', 'ALE/Asterix-v5', 'ALE/BankHeist-v5',
+            'ALE/BattleZone-v5', 'ALE/CrazyClimber-v5', 'ALE/DemonAttack-v5', 'ALE/Freeway-v5',
+            'ALE/Frostbite-v5', 'ALE/Gopher-v5', 'ALE/Jamesbond-v5', 'ALE/Kangaroo-v5',
+            'ALE/Krull-v5', 'ALE/KungFuMaster-v5', 'ALE/PrivateEye-v5', 'ALE/UpNDown-v5',
+            'ALE/Qbert-v5', 'ALE/Breakout-v5',
+        ]
+    else:
+        raise ValueError(f"Unsupported number of environments: {num_games}")
+
+    # ==================== Batch Size Calculation ====================
+    if len(env_id_list) == 8:
+        if num_layers in [1, 4]:
+            effective_batch_size = 10
+        elif num_layers == 8:
+            effective_batch_size = 10
+        else:
+            effective_batch_size = 10
+    elif len(env_id_list) == 26:
+        effective_batch_size = 512
+    elif len(env_id_list) == 3:
+        effective_batch_size = 10  # Reduced for debugging; use 512 for full training
+    else:
+        raise ValueError(f"Batch size not configured for {len(env_id_list)} environments.")
+
+    batch_sizes, grad_acc_steps = compute_batch_config(env_id_list, effective_batch_size, gpu_num=4)
+    total_batch_size = effective_batch_size
+
+    # ==================== Model and Training Settings ====================
+    num_unroll_steps = 10
+    infer_context_length = 4
+    norm_type = 'LN'
+    buffer_reanalyze_freq = 1 / 100000000  # Effectively disable buffer reanalyze
+    reanalyze_batch_size = 160
+    reanalyze_partition = 0.75
+
+    # ==================== Training Loop ====================
+    # Set NCCL timeout to prevent watchdog hang due to unbalanced data collection speeds
+    os.environ.setdefault('NCCL_TIMEOUT', '3600')  # 60 minutes in seconds
+    os.environ.setdefault('NCCL_BLOCKING_WAIT', '1')
+
+    for seed in [0]:
+        configs = generate_configs(
+            env_id_list, action_space_size, collector_env_num, n_episode, evaluator_env_num,
+            num_simulations, reanalyze_ratio, batch_sizes, num_unroll_steps, infer_context_length,
+            norm_type, seed, buffer_reanalyze_freq, reanalyze_batch_size, reanalyze_partition,
+            num_segments, total_batch_size, num_layers
+        )
+
+        with DDPContext():
+            train_unizero_multitask_segment_ddp(configs, seed=seed, max_env_step=max_env_step, benchmark_name="atari")
+            print(f"Seed: {seed} training finished!")
+            if dist.is_initialized():
+                dist.destroy_process_group()
diff --git a/zoo/atari/config/atari_unizero_nomoe_multitask_segment_ddp_config.py b/zoo/atari/config/atari_unizero_nomoe_multitask_segment_ddp_config.py
new file mode 100644
index 000000000..29762cf01
--- /dev/null
+++ b/zoo/atari/config/atari_unizero_nomoe_multitask_segment_ddp_config.py
@@ -0,0 +1,298 @@
+from easydict import EasyDict
+import math
+from typing import List, Tuple, Any, Dict, Union
+
+# -------------------------------------------------
+# 1. Refactored compute_batch_config (same as MoE config)
+# -------------------------------------------------
+def compute_batch_config(
+    env_id_list: List[str],
+    effective_batch_size: int,
+    gpu_num: int = 8,
+    max_micro_batch_one_gpu: int = 400,
+) -> Tuple[List[int], int]:
+    """
+    Overview:
+        Calculate the micro-batch size for each environment and the number of gradient accumulation steps
+        to approach a target effective batch size across multiple GPUs and environments.
+
+    Arguments:
+        - env_id_list (:obj:`List[str]`): A list of environment IDs for all tasks.
+        - effective_batch_size (:obj:`int`): The target global batch size for one backward pass.
+        - gpu_num (:obj:`int`): The number of GPUs actually used. Defaults to 8.
+        - max_micro_batch_one_gpu (:obj:`int`): The maximum micro-batch size a single GPU can handle. Defaults to 400.
+
+    Returns:
+        - batch_sizes (:obj:`List[int]`): A list of micro-batch sizes for each environment.
+        - grad_acc_steps (:obj:`int`): The number of gradient accumulation steps.
+    """
+    n_env = len(env_id_list)
+    # Number of environments that each GPU needs to handle simultaneously.
+    envs_per_gpu = max(1, math.ceil(n_env / gpu_num))
+    # Reduce the micro-batch limit if multiple environments share one GPU.
+    max_micro_batch = max(1, max_micro_batch_one_gpu // envs_per_gpu)
+
+    # First, calculate a candidate micro-batch by distributing the effective batch size evenly.
+    candidate = max(1, effective_batch_size // n_env)
+    micro_batch = min(candidate, max_micro_batch)
+
+    # Gradient accumulation steps = ceil(global_batch / (micro_batch * n_env)).
+    grad_acc_steps = max(1, math.ceil(effective_batch_size / (micro_batch * n_env)))
+
+    # Fine-tune the micro-batch downwards to ensure:
+    # micro_batch * n_env * grad_acc_steps <= effective_batch_size
+    # This aims to get as close as possible to the target without exceeding it.
+    while micro_batch * n_env * grad_acc_steps > effective_batch_size:
+        micro_batch -= 1
+        if micro_batch == 0:  # Defensive check, should not happen in theory.
+            micro_batch = 1
+            break
+
+    batch_sizes = [micro_batch] * n_env
+
+    # --- Debug Information --- #
+    real_total_batch_size = micro_batch * n_env * grad_acc_steps
+    print(
+        f"[BatchConfig] Envs={n_env}, TargetTotalBS={effective_batch_size}, "
+        f"MicroBS={micro_batch}, GradAccSteps={grad_acc_steps}, RealTotalBS={real_total_batch_size}"
+    )
+
+    return batch_sizes, grad_acc_steps
+
+
+def create_config(
+        env_id: str, action_space_size: int, collector_env_num: int, evaluator_env_num: int, n_episode: int,
+        num_simulations: int, reanalyze_ratio: float, batch_size: int, num_unroll_steps: int,
+        infer_context_length: int, norm_type: str, buffer_reanalyze_freq: float, reanalyze_batch_size: int,
+        reanalyze_partition: float, num_segments: int, total_batch_size: int, num_layers: int
+) -> EasyDict:
+    """
+    Overview:
+        Creates the main configuration structure for a single training task (MoE disabled).
+    """
+    return EasyDict(dict(
+        env=dict(
+            stop_value=int(1e6),
+            env_id=env_id,
+            observation_shape=(3, 64, 64),
+            gray_scale=False,
+            collector_env_num=collector_env_num,
+            evaluator_env_num=evaluator_env_num,
+            n_evaluator_episode=evaluator_env_num,
+            manager=dict(shared_memory=False),
+            full_action_space=True,
+        ),
+        policy=dict(
+            multi_gpu=True,
+            only_use_moco_stats=False,
+            use_moco=False,
+            moco_version="v1",
+            total_task_num=len(env_id_list),
+            task_num=len(env_id_list),
+            task_id=0,  # This will be overridden for each task
+            model=dict(
+                observation_shape=(3, 64, 64),
+                action_space_size=action_space_size,
+                norm_type=norm_type,
+                num_res_blocks=2,
+                num_channels=256,
+                num_layers=num_layers,
+                world_model_cfg=dict(
+                    norm_type=norm_type,
+                    action_space_size=action_space_size,
+                    num_layers=num_layers,
+                    num_heads=8,
+                    embed_dim=768,
+                    env_num=len(env_id_list),
+                    task_num=len(env_id_list),
+                    max_blocks=num_unroll_steps,
+                    max_tokens=2 * num_unroll_steps,
+                    context_length=2 * infer_context_length,
+                    final_norm_option_in_obs_head='LayerNorm',
+                    final_norm_option_in_encoder='LayerNorm',
+                    predict_latent_loss_type='mse',
+                    encoder_type='vit',
+                    device='cuda',
+                    game_segment_length=20,
+                    # MoE disabled: run plain transformer encoder
+                    use_normal_head=True,
+                    use_softmoe_head=False,
+                    use_moe_head=False,
+                    num_experts_in_moe_head=1,
+                    moe_in_transformer=False,
+                    multiplication_moe_in_transformer=False,
+                    n_shared_experts=1,
+                    num_experts_per_tok=2,
+                    num_experts_of_moe_in_transformer=8,
+                    moe_use_lora=False,
+                ),
+            ),
+            device='cuda',
+            game_segment_length=20,
+            update_per_collect=80,  # Corresponds to replay_ratio=0.5 for 8 games (20*8*0.5=80)
+            learning_rate=0.0001,
+            weight_decay=1e-2,
+            batch_size=batch_size,
+            num_unroll_steps=num_unroll_steps,
+            num_segments=num_segments,
+            num_simulations=num_simulations,
+            reanalyze_ratio=reanalyze_ratio,
+            n_episode=n_episode,
+            total_batch_size=total_batch_size,
+            collector_env_num=collector_env_num,
+            evaluator_env_num=evaluator_env_num,
+            buffer_reanalyze_freq=buffer_reanalyze_freq,
+            reanalyze_batch_size=reanalyze_batch_size,
+            reanalyze_partition=reanalyze_partition,
+            replay_buffer_size=int(5e5),
+            eval_freq=int(1e4),
+        ),
+    ))
+
+
+def generate_configs(
+        env_id_list: List[str], action_space_size: int, collector_env_num: int, n_episode: int,
+        evaluator_env_num: int, num_simulations: int, reanalyze_ratio: float, batch_size: List[int],
+        num_unroll_steps: int, infer_context_length: int, norm_type: str, seed: int,
+        buffer_reanalyze_freq: float, reanalyze_batch_size: int, reanalyze_partition: float,
+        num_segments: int, total_batch_size: int, num_layers: int
+) -> List[List[Union[int, List[EasyDict]]]]:
+    """
+    Overview:
+        Generates a list of configurations for all specified tasks (MoE disabled).
+    """
+    configs = []
+
+    # --- Experiment Name Template ---
+    benchmark_tag = "data_unizero_mt"
+    model_tag = f"vit_nlayer{num_layers}_tbs{total_batch_size}_nomoe"
+    exp_name_prefix = f'{benchmark_tag}/atari_{len(env_id_list)}games_{model_tag}_seed{seed}/'
+
+    for task_id, env_id in enumerate(env_id_list):
+        config = create_config(
+            env_id, action_space_size, collector_env_num, evaluator_env_num, n_episode, num_simulations,
+            reanalyze_ratio, batch_size, num_unroll_steps, infer_context_length, norm_type,
+            buffer_reanalyze_freq, reanalyze_batch_size, reanalyze_partition, num_segments, total_batch_size, num_layers
+        )
+        config.policy.task_id = task_id
+        # Correctly extract the game name from 'ALE/GameName-v5' format.
+        game_name = env_id.split('/')[1].split('-')[0]
+        config.exp_name = exp_name_prefix + f"{game_name}_seed{seed}"
+        configs.append([task_id, [config, create_env_manager()]])
+    return configs
+
+
+def create_env_manager() -> EasyDict:
+    """
+    Overview:
+        Creates the environment manager configuration, specifying the types of environment,
+        policy, and their import paths.
+    """
+    return EasyDict(dict(
+        env=dict(
+            type='atari_lightzero',
+            import_names=['zoo.atari.envs.atari_lightzero_env'],
+        ),
+        env_manager=dict(type='subprocess'),
+        policy=dict(
+            type='unizero_multitask',
+            import_names=['lzero.policy.unizero_multitask'],
+        ),
+    ))
+
+
+if __name__ == "__main__":
+    """
+    Overview:
+        This script should be executed with <nproc_per_node> GPUs for distributed training.
+
+        Example launch commands:
+
+        export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+        cd /path/to/your/project/
+
+        torchrun --nproc_per_node=4 /mnt/shared-storage-user/puyuan/code/LightZero/zoo/atari/config/atari_unizero_nomoe_multitask_segment_ddp_config.py
+    """
+    from lzero.entry import train_unizero_multitask_segment_ddp
+    from ding.utils import DDPContext
+    import torch.distributed as dist
+    import os
+
+    # ==================== Main Experiment Settings ====================
+    num_games = 8  # Options: 3, 8, 26
+    num_layers = 1  # Transformer depth (reduced for faster iteration)
+    action_space_size = 18
+    collector_env_num = 8
+    num_segments = 8
+    n_episode = 8
+    evaluator_env_num = 3
+    num_simulations = 25  # MCTS simulations per step
+    max_env_step = int(5e6)
+    reanalyze_ratio = 0.0
+
+    # ==================== Environment Configuration ====================
+    if num_games == 3:
+        env_id_list = ['ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5']
+    elif num_games == 8:
+        env_id_list = [
+            'ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5', 'ALE/Boxing-v5',
+            'ALE/Alien-v5', 'ALE/ChopperCommand-v5', 'ALE/Hero-v5', 'ALE/RoadRunner-v5',
+        ]
+    elif num_games == 26:
+        env_id_list = [
+            'ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5', 'ALE/Boxing-v5',
+            'ALE/Alien-v5', 'ALE/ChopperCommand-v5', 'ALE/Hero-v5', 'ALE/RoadRunner-v5',
+            'ALE/Amidar-v5', 'ALE/Assault-v5', 'ALE/Asterix-v5', 'ALE/BankHeist-v5',
+            'ALE/BattleZone-v5', 'ALE/CrazyClimber-v5', 'ALE/DemonAttack-v5', 'ALE/Freeway-v5',
+            'ALE/Frostbite-v5', 'ALE/Gopher-v5', 'ALE/Jamesbond-v5', 'ALE/Kangaroo-v5',
+            'ALE/Krull-v5', 'ALE/KungFuMaster-v5', 'ALE/PrivateEye-v5', 'ALE/UpNDown-v5',
+            'ALE/Qbert-v5', 'ALE/Breakout-v5',
+        ]
+    else:
+        raise ValueError(f"Unsupported number of environments: {num_games}")
+
+    # ==================== Batch Size Calculation ====================
+    if len(env_id_list) == 8:
+        if num_layers in [1, 4]:
+            effective_batch_size = 10
+        elif num_layers == 8:
+            effective_batch_size = 10
+        else:
+            effective_batch_size = 10
+    elif len(env_id_list) == 26:
+        effective_batch_size = 512
+    elif len(env_id_list) == 3:
+        effective_batch_size = 10  # Reduced for debugging; use 512 for full training
+    else:
+        raise ValueError(f"Batch size not configured for {len(env_id_list)} environments.")
+
+    batch_sizes, grad_acc_steps = compute_batch_config(env_id_list, effective_batch_size, gpu_num=4)
+    total_batch_size = effective_batch_size
+
+    # ==================== Model and Training Settings ====================
+    num_unroll_steps = 10
+    infer_context_length = 4
+    norm_type = 'LN'
+    buffer_reanalyze_freq = 1 / 100000000  # Effectively disable buffer reanalyze
+    reanalyze_batch_size = 160
+    reanalyze_partition = 0.75
+
+    # ==================== Training Loop ====================
+    # Set NCCL timeout to prevent watchdog hang due to unbalanced data collection speeds
+    os.environ.setdefault('NCCL_TIMEOUT', '3600')  # 60 minutes in seconds
+    os.environ.setdefault('NCCL_BLOCKING_WAIT', '1')
+
+    for seed in [0]:
+        configs = generate_configs(
+            env_id_list, action_space_size, collector_env_num, n_episode, evaluator_env_num,
+            num_simulations, reanalyze_ratio, batch_sizes, num_unroll_steps, infer_context_length,
+            norm_type, seed, buffer_reanalyze_freq, reanalyze_batch_size, reanalyze_partition,
+            num_segments, total_batch_size, num_layers
+        )
+
+        with DDPContext():
+            train_unizero_multitask_segment_ddp(configs, seed=seed, max_env_step=max_env_step, benchmark_name="atari")
+            print(f"Seed: {seed} training finished!")
+            if dist.is_initialized():
+                dist.destroy_process_group()
+