update

tr2d2-pep/peptide_mcts.py

@@ -345,7 +345,7 @@ class MCTS:
         print(f"[BUFFER] reason={reason} sv={np.round(sv,4)} "
             f"buf_len={len(self.buffer)} extra={extra}")
 
-    def updateBuffer(self, x_final, log_rnd, score_vectors, childSequences):
+    def updateBuffer(self, x_final, log_rnd, score_vectors, validSequences):
         B = x_final.shape[0]
         traj_log_rnds, scalar_rewards = [], []
 
@@ -367,7 +367,7 @@ class MCTS:
                 "log_rnd": traj_log_rnd.clone(),
                 "final_reward": scalar_reward,
                 "score_vector": sv.copy(),
-                "seq": childSequences[i],
+                "seq": validSequences[i],
             }
 
             # Drop if dominated by any existing
@@ -601,7 +601,7 @@ class MCTS:
         valid_traj_log_rnd = torch.stack(valid_traj_log_rnd, dim=0)
         # update buffer and get rewards
         with self.timer.section("expand.update_buffer"):
-            traj_log_rnds, scalar_rewards = self.updateBuffer(valid_x_final, valid_traj_log_rnd, score_vectors, childSequences)
+            traj_log_rnds, scalar_rewards = self.updateBuffer(valid_x_final, valid_traj_log_rnd, score_vectors, validSequences)
         
         allChildReward = np.zeros_like(score_vectors[0])