feature(xjy): Refine PriorZero Implementation

lzero/model/unizero_world_models/world_model.py

@@ -2064,7 +2064,7 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
                 value_priority=value_priority,
                 intermediate_tensor_x=intermediate_tensor_x,
                 obs_embeddings=detached_obs_embeddings, # <-- 新增
-            )
+            ), inverse_scalar_transform_handle(outputs.logits_value.reshape(-1, outputs.logits_value.shape[-1])).detach()
 
 
     # TODO: test correctness

lzero/worker/muzero_segment_collector.py

@@ -477,16 +477,6 @@ def collect(
                     if self.policy_config.use_ture_chance_label_in_chance_encoder:
                         append_kwargs['chance'] = self.chance_dict_tmp[env_id]
 
-                    # [PRIORZERO-NEW] Add raw_obs_text if available in obs (not info!)
-                    # Jericho env puts raw_obs_text in the obs dictionary
-                    if env_id == 0 and collected_step < 5:  # Debug first few steps
-                        print(f"[OBS_DEBUG] Step {collected_step} env {env_id}: obs keys = {list(obs.keys())}")
-                        print(f"[OBS_DEBUG] obs type = {type(obs)}")
-                        if 'raw_obs_text' in obs:
-                            print(f"[OBS_DEBUG] Found raw_obs_text: {str(obs['raw_obs_text'])[:100]}...")
-                        else:
-                            print(f"[OBS_DEBUG] NO raw_obs_text in obs!")
-
                     if 'raw_obs_text' in obs:
                         append_kwargs['raw_obs_text'] = obs['raw_obs_text']
                     elif 'raw_obs_text' in info:

zoo/jericho/envs/jericho_env.py

@@ -4,6 +4,7 @@
 import json
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Union
+from collections import OrderedDict
 
 import gym
 import numpy as np
@@ -49,12 +50,13 @@ class JerichoEnv(BaseEnv):
         'max_seq_len': 512,
         'remove_stuck_actions': False,
         'add_location_and_inventory': False,
-        # 'for_unizero': False,
         'for_unizero': True,
         'save_replay': False,
         'save_replay_path': None,
         'env_type': "zork1",
-        'collect_policy_mode': "agent"
+        'collect_policy_mode': "agent",
+        'use_cache': True,
+        'cache_size': 100000,
     }
 
     def __init__(self, cfg: Dict[str, Any]) -> None:
@@ -93,6 +95,12 @@ def __init__(self, cfg: Dict[str, Any]) -> None:
         self.add_location_and_inventory: bool = self.cfg['add_location_and_inventory']
         self.for_unizero: bool = self.cfg['for_unizero']
 
+        self.use_cache = self.cfg['use_cache']
+        if self.use_cache:
+            self.cache_size = self.cfg['cache_size']
+            self.cache_buffer = OrderedDict()
+            print(f'[jericho]: use_cache: {self.use_cache}, cache_size={self.cache_size}')
+
         # Initialize the tokenizer once (only in rank 0 process if distributed)
         if JerichoEnv.tokenizer is None:
             if self.rank == 0:
@@ -138,7 +146,18 @@ def prepare_obs(self, obs: str, return_str: bool = False) -> Dict[str, Any]:
         raw_obs_text = obs  # Save original text BEFORE any modification
 
         if self._action_list is None:
-            self._action_list = self._env.get_valid_actions()
+            if self.use_cache:
+                cache_key = self._env.get_world_state_hash()
+                if cache_key in self.cache_buffer:
+                    self.cache_buffer.move_to_end(cache_key)
+                    self._action_list = self.cache_buffer[cache_key]
+                else:
+                    self._action_list = self._env.get_valid_actions()
+                    self.cache_buffer[cache_key] = self._action_list
+                    if len(self.cache_buffer) > self.cache_size:
+                        self.cache_buffer.popitem(last=False)
+            else:
+                self._action_list = self._env.get_valid_actions()
 
         # Filter available actions based on whether stuck actions are removed.
         if self.remove_stuck_actions:
@@ -344,6 +363,7 @@ def step(self, action: Union[int, np.ndarray, str], return_str: bool = False) ->
         previous_obs: Optional[str] = self.last_observation if (self.remove_stuck_actions and self.last_observation is not None) else None
 
         observation, reward, done, info = self._env.step(action_str)
+        info['action_str'] = action_str
 
         self._timestep += 1
         if not self.for_unizero: