feature(pu): add entropy_loss_coef and prior_mixing_cfg

Author: puyuan1996

zoo/jericho/priorzero/models/actor.py

@@ -224,7 +224,8 @@ def train_batch(self, batch_data: Dict[str, torch.Tensor], kl_ctl: float, step_i
                 micro_batch['action_mask'],
                 attention_mask=micro_batch['attention_mask'],
                 return_output=True,
-                logits_to_keep=logits_to_keep, 
+                return_entropy=self.args.entropy_loss_coef is not None,
+                logits_to_keep=logits_to_keep,
             )
             actor_loss, clipfrac, clip_ratio, approx_kl, vllm_kl = self.policy_loss(
                 action_log_probs,
@@ -242,8 +243,20 @@ def train_batch(self, batch_data: Dict[str, torch.Tensor], kl_ctl: float, step_i
                 kl_loss = masked_mean(kl, micro_batch["action_mask"])
             else:
                 kl_loss = 0.0
-            
+
+            # Entropy loss for exploration bonus
+            if self.args.entropy_loss_coef is not None:
+                # Extract entropy for action tokens only
+                # Note: output.entropy is already [:, :-1] from Actor.forward (line 89)
+                # So we extract the last action_mask.shape[1] tokens
+                entropy = output.entropy[:, -micro_batch['action_mask'].shape[1]:]
+                entropy_loss = masked_mean(entropy, micro_batch['action_mask'])
+            else:
+                entropy_loss = 0.0
+
             loss = actor_loss + kl_loss * float(kl_ctl.value)
+            if self.args.entropy_loss_coef is not None and self.args.entropy_loss_coef != 0:
+                loss -= entropy_loss * self.args.entropy_loss_coef
 
             self.strategy.backward(loss, self.actor, self.actor_optim)
 
@@ -326,6 +339,13 @@ def train_batch(self, batch_data: Dict[str, torch.Tensor], kl_ctl: float, step_i
                 status["cur_refer_kl"] = kl_loss.detach().float().mean().item()
             else:
                 status["cur_refer_kl"] = float(kl_loss)
+
+            # Add entropy loss logging
+            if self.args.entropy_loss_coef is not None:
+                if isinstance(entropy_loss, torch.Tensor):
+                    status["entropy_loss"] = entropy_loss.detach().float().mean().item()
+                else:
+                    status["entropy_loss"] = float(entropy_loss)
 
             status = self.strategy.all_reduce(status)
             status_list.append(status)

zoo/jericho/priorzero/priorzero_config.py

@@ -181,7 +181,25 @@ class PriorZeroLLMConfig:
     eps_clip_low_high: Tuple[float, float] = (0.2, 0.2)
     rft_kl_coef: float = 0.01
     kl_estimator: str = "k3"
-    
+
+    # Entropy loss for exploration bonus
+    entropy_loss_coef: Optional[float] = 0.01  # None = disabled, typical values: 0.001-0.01
+    # entropy_loss_coef: Optional[float] = None  # None = disabled, typical values: 0.001-0.01  
+
+    # LLM Prior Mixing Configuration
+    prior_mixing_cfg: Optional[EasyDict] = field(default_factory=lambda: EasyDict({
+        'enable_soft_mixing': True,              # Enable soft mixing instead of hard override
+        # 'mixing_alpha': 0.5,                     # Weight for LLM prior (0=network only, 1=LLM only)
+        'mixing_alpha': 0.,                     # Weight for LLM prior (0=network only, 1=LLM only)
+        # 'alpha_schedule': None,                  # 'linear', 'cosine', 'exponential', or None (fixed)
+        # 'alpha_schedule': 'cosine',  # Smooth decay          
+        'alpha_init': 0.8,                       # Initial alpha (high LLM influence)
+        'alpha_final': 0.2,                      # Final alpha (low LLM influence)
+        'alpha_decay_steps': 10000,              # Steps to decay from init to final
+        'enable_clip_prior': True,               # Enable clipping of LLM prior probabilities
+        'clip_prior_epsilon': 0.01,              # Minimum probability for each action (exploration)
+    }))
+
     train_llm_after_wm_warm_step: int = int(1e2) # TODO
     value_norm_cfg: Optional[EasyDict] = field(default_factory=lambda: EasyDict({
         'enable_stability_optimizer': True,
@@ -327,7 +345,9 @@ def get_priorzero_config(
         n_episode=n_episode,
         train_start_after_envsteps=0,
         replay_buffer_size=replay_buffer_size,
-        eval_freq=int(3e4),
+        # eval_freq=int(3e4),
+        eval_freq=int(1e3), # TODO
+        # eval_freq=int(2), # TODO
         collector_env_num=collector_env_num,
         evaluator_env_num=evaluator_env_num,
         buffer_reanalyze_freq=1 / 1000000,
@@ -435,11 +455,16 @@ def get_priorzero_config(
 
         # Format reward info
         fmt_rew_str = "fmt" if llm_config.reward_func.format_reward else "nofmt"
+        # entropy_loss_coef = 
 
         # Build exp_name
+        # exp_name = (
+        #     f"data_priorzero/pz_{env_id}_{model_key}_"
+        #     f"{cot_str}_{adv_type_short}_{prior_temp_str}_{fmt_rew_str}_pel{entropy_loss_coef}_llm-mix-0-true_seed{seed}" # TODO
+        # )
         exp_name = (
             f"data_priorzero/pz_{env_id}_{model_key}_"
-            f"{cot_str}_{adv_type_short}_{prior_temp_str}_{fmt_rew_str}_seed{seed}"
+            f"{cot_str}_{adv_type_short}_{prior_temp_str}_{fmt_rew_str}_pel001_llm-mix-0-true_seed{seed}" # TODO
         )
 
         # Update config with generated exp_name

zoo/jericho/priorzero/priorzero_policy.py

@@ -26,8 +26,9 @@
 
 @POLICY_REGISTRY.register('priorzero', force_overwrite=True)
 class PriorZeroPolicy(OriginalUniZeroPolicy):
-    def __init__(self, cfg: Dict, model: torch.nn.Module = None, enable_field: List[str] = None, **kwargs):   
+    def __init__(self, cfg: Dict, model: torch.nn.Module = None, enable_field: List[str] = None, **kwargs):
         super().__init__(cfg, model, enable_field)
+        self._mixing_step = 0  # Track steps for alpha scheduling
 
     def _init_learn(self) -> None:
         super()._init_learn()
@@ -283,7 +284,76 @@ def pad_to_fixed_length(self, data, target_len=55, pad_val=-1e9, dtype=torch.flo
             if L > 0:
                 out[i, :L] = torch.tensor(seq[:L], dtype=dtype)
         return out
-    
+
+    def _clip_prior_probabilities(self, policy_logits: torch.Tensor, epsilon: float, action_mask: List[np.ndarray]) -> torch.Tensor:
+        """
+        Clip LLM prior probabilities to ensure minimum exploration.
+
+        Args:
+            policy_logits: Log probabilities from LLM [B, A]
+            epsilon: Minimum probability for each legal action
+            action_mask: List of action masks for each environment
+
+        Returns:
+            Clipped policy logits
+        """
+        # Convert logits to probabilities
+        policy_probs = F.softmax(policy_logits, dim=-1)
+
+        # Clip probabilities for legal actions
+        batch_size = policy_probs.shape[0]
+        for i in range(batch_size):
+            legal_actions = action_mask[i] == 1
+            num_legal = legal_actions.sum()
+
+            if num_legal > 0:
+                # Clip legal action probabilities to be at least epsilon
+                policy_probs[i, legal_actions] = torch.clamp(
+                    policy_probs[i, legal_actions],
+                    min=epsilon
+                )
+
+                # Renormalize to sum to 1
+                policy_probs[i, legal_actions] = policy_probs[i, legal_actions] / policy_probs[i, legal_actions].sum()
+
+        # Convert back to log probabilities
+        clipped_logits = torch.log(policy_probs + 1e-10)
+        return clipped_logits
+
+    def _compute_mixing_alpha(self, cfg: Dict) -> float:
+        """
+        Compute the mixing alpha based on schedule configuration.
+
+        Args:
+            cfg: Prior mixing configuration
+
+        Returns:
+            Current alpha value
+        """
+        if not cfg.get('alpha_schedule'):
+            # Fixed alpha
+            return cfg.get('mixing_alpha', 0.5)
+
+        schedule_type = cfg['alpha_schedule']
+        init_alpha = cfg.get('alpha_init', 0.8)
+        final_alpha = cfg.get('alpha_final', 0.2)
+        decay_steps = cfg.get('alpha_decay_steps', 10000)
+
+        # Compute progress
+        progress = min(self._mixing_step / decay_steps, 1.0)
+
+        if schedule_type == 'linear':
+            alpha = init_alpha + (final_alpha - init_alpha) * progress
+        elif schedule_type == 'cosine':
+            alpha = final_alpha + (init_alpha - final_alpha) * 0.5 * (1 + np.cos(np.pi * progress))
+        elif schedule_type == 'exponential':
+            decay_rate = cfg.get('alpha_decay_rate', 0.95)
+            alpha = final_alpha + (init_alpha - final_alpha) * (decay_rate ** self._mixing_step)
+        else:
+            alpha = cfg.get('mixing_alpha', 0.5)
+
+        return alpha
+
     def _forward_collect(
         self,
         data: torch.Tensor,
@@ -324,12 +394,41 @@ def _forward_collect(
                     prior.append(llm_prior_logprob[env_id][action])
             policy_priors.append(prior)
         policy_priors = self.pad_to_fixed_length(data=policy_priors, target_len=self.cfg.model.action_space_size, pad_val=-1e9)
-        
+
         with torch.no_grad():
             network_output = self._collect_model.initial_inference(self.last_batch_obs, self.last_batch_action, data, timestep)
             latent_state_roots, reward_roots, pred_values, policy_logits = mz_network_output_unpack(network_output)
 
-            network_output.policy_logits = policy_priors
+            # ======================================================================
+            # LLM Prior Mixing: Soft Mixing + Clip Prior
+            # ======================================================================
+            # Get mixing configuration from llm_config if available
+            from zoo.jericho.priorzero.priorzero_entry_sync_ddp import llm_config
+            mixing_cfg = llm_config.prior_mixing_cfg if hasattr(llm_config, 'prior_mixing_cfg') else {}
+
+            # Store original network policy for logging
+            network_policy_logits = policy_logits.clone()
+
+            # Apply clip prior if enabled
+            if mixing_cfg.get('enable_clip_prior', False):
+                epsilon = mixing_cfg.get('clip_prior_epsilon', 0.01)
+                policy_priors = self._clip_prior_probabilities(policy_priors, epsilon, action_mask)
+
+            # Apply soft mixing if enabled
+            if mixing_cfg.get('enable_soft_mixing', False):
+                alpha = self._compute_mixing_alpha(mixing_cfg)
+                # Soft mixing: (1 - alpha) * network + alpha * LLM
+                mixed_policy_logits = (1 - alpha) * policy_logits + alpha * policy_priors
+                final_policy_logits = mixed_policy_logits
+                self._mixing_step += 1  # Increment step for alpha scheduling
+            else:
+                # Hard override (original behavior)
+                final_policy_logits = policy_priors
+                alpha = 1.0  # For logging
+
+            # Update network output with final policy
+            network_output.policy_logits = final_policy_logits
+
             if not self._cfg.mcts_ctree:
                 raise NotImplementedError("Python MCTS not supported for PriorZero")
 
@@ -338,7 +437,7 @@ def _forward_collect(
             # ======================================================================
             pred_values_np = self.value_inverse_scalar_transform_handle(pred_values).detach().cpu().numpy()
             latent_state_roots_np = latent_state_roots.detach().cpu().numpy()
-            policy_logits = policy_priors.detach().cpu().numpy().tolist()
+            policy_logits_for_mcts = final_policy_logits.detach().cpu().numpy().tolist()
 
 
             legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_collect_env_num)]
@@ -347,7 +446,7 @@ def _forward_collect(
                                     ).astype(np.float32).tolist() for j in range(active_collect_env_num)
             ]
             roots = MCTSCtree.roots(active_collect_env_num, legal_actions)
-            roots.prepare(self._cfg.root_noise_weight, noises, reward_roots, policy_logits, to_play)
+            roots.prepare(self._cfg.root_noise_weight, noises, reward_roots, policy_logits_for_mcts, to_play)
             self._mcts_collect.search(roots, self._collect_model, latent_state_roots_np, to_play, timestep=timestep)
 
             roots_visit_count = roots.get_distributions()
@@ -373,8 +472,19 @@ def _forward_collect(
                     'visit_count_distribution_entropy': visit_count_distribution_entropy,
                     'searched_value': value,
                     'predicted_value': pred_values_np[i],
-                    'predicted_policy_logits': policy_logits[i],
+                    'predicted_policy_logits': policy_logits_for_mcts[i],
                     'timestep': timestep[i],
+                    # Add mixing metrics for logging
+                    'mixing_alpha': alpha,
+                    'network_policy_entropy': -torch.sum(
+                        F.softmax(network_policy_logits[i], dim=-1) * F.log_softmax(network_policy_logits[i], dim=-1)
+                    ).item(),
+                    'llm_policy_entropy': -torch.sum(
+                        F.softmax(policy_priors[i], dim=-1) * F.log_softmax(policy_priors[i], dim=-1)
+                    ).item(),
+                    'mixed_policy_entropy': -torch.sum(
+                        F.softmax(final_policy_logits[i], dim=-1) * F.log_softmax(final_policy_logits[i], dim=-1)
+                    ).item(),
                 }
                 batch_action.append(action)
             self.last_batch_obs = data