polish(pu): polish unizero config

Author: puyuan1996

lzero/mcts/buffer/game_buffer.py

@@ -697,16 +697,16 @@ def _remove(self, excess_game_segment_index: int) -> None:
             f"  - base_idx: {base_idx_after}\n"
             f"------------------------------\n\n"
         )
-
+        # TODO
         # 5. Print to console and write to file
-        print(log_message)
+        # print(log_message)
 
-        log_filename = f"game_buffer_remove_log_{timestamp.strftime('%Y%m%d_%H%M%S')}.txt"
-        try:
-            with open(log_filename, 'a', encoding='utf-8') as f:
-                f.write(log_message)
-        except Exception as e:
-            print(f"[ERROR] Failed to write to log file {log_filename}: {e}")
+        # log_filename = f"game_buffer_remove_log_{timestamp.strftime('%Y%m%d_%H%M%S')}.txt"
+        # try:
+        #     with open(log_filename, 'a', encoding='utf-8') as f:
+        #         f.write(log_message)
+        # except Exception as e:
+        #     print(f"[ERROR] Failed to write to log file {log_filename}: {e}")
 
         # --- End of logging modification ---
 

lzero/model/unizero_world_models/utils.py

@@ -263,18 +263,18 @@ def __init__(self, latent_recon_loss_weight=0, perceptual_loss_weight=0, continu
         # NOTE: Define the weights for each loss type
         if not continuous_action_space:
             # like EZV2, for atari and memory
-            self.obs_loss_weight = 10
-            self.value_loss_weight = 0.5
-            self.reward_loss_weight = 1.
-            self.policy_loss_weight = 1.
-            self.ends_loss_weight = 0.
+            # self.obs_loss_weight = 10
+            # self.value_loss_weight = 0.5
+            # self.reward_loss_weight = 1.
+            # self.policy_loss_weight = 1.
+            # self.ends_loss_weight = 0.
 
             # muzero loss weight
-            # self.obs_loss_weight = 2
-            # self.value_loss_weight = 0.25
-            # self.reward_loss_weight = 1
-            # self.policy_loss_weight = 1
-            # self.ends_loss_weight = 0.
+            self.obs_loss_weight = 2
+            self.value_loss_weight = 0.25
+            self.reward_loss_weight = 1
+            self.policy_loss_weight = 1
+            self.ends_loss_weight = 0.
 
             # like TD-MPC2 for DMC
             # self.obs_loss_weight = 10