fix: use trace count as index (#909)

Author: qew21

rdagent/scenarios/data_science/proposal/exp_gen/ckp_select.py

@@ -70,7 +70,7 @@ def get_selection(self, trace: Trace) -> tuple[int, ...]:
             return (-1,)  # Continue with latest trial for new sub-trace
 
         # Calculate elapsed time for current sub-trace
-        elapsed_time = current_time - self.sub_trace_start_times[trace.sub_trace_count]
+        elapsed_time = current_time - self.sub_trace_start_times[trace.sub_trace_count - 1]
 
         if elapsed_time < self.time_limit_pre_trace:
             # Continue with current sub-trace
@@ -81,19 +81,19 @@ def get_selection(self, trace: Trace) -> tuple[int, ...]:
             return (-1,)
         else:
             # Check if we've reached the maximum number of traces
-            if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
+            if trace.sub_trace_count >= self.MAX_TRACE_NUM:
                 logger.info(
                     f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
                 )
                 logger.info(f"current sub-trace count: {trace.sub_trace_count}")
                 return (-1,)
 
             # Time limit exceeded, start a new sub-trace
-            self.sub_trace_start_times[trace.sub_trace_count + 1] = current_time
+            self.sub_trace_start_times[trace.sub_trace_count] = current_time
             logger.info(
                 f"Elapsed time {elapsed_time} exceeds time limit {self.time_limit_pre_trace}, jump to a new sub-trace"
             )
-            logger.info(f"current sub-trace count: {trace.sub_trace_count + 1}")
+            logger.info(f"current sub-trace count: {trace.sub_trace_count}")
             return tuple()  # Empty tuple signals starting a new sub-trace
 
 
@@ -129,7 +129,7 @@ def get_selection(self, trace: Trace) -> tuple[int, ...]:
                     sota_count += 1
             if sota_count < self.SOTA_COUNT_THRESHOLD:
                 # Check if we've reached the maximum number of traces
-                if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
+                if trace.sub_trace_count >= self.MAX_TRACE_NUM:
                     logger.info(
                         f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
                     )
@@ -139,7 +139,7 @@ def get_selection(self, trace: Trace) -> tuple[int, ...]:
                 logger.info(
                     f"SOTA count {sota_count} is below threshold {self.SOTA_COUNT_THRESHOLD}, jump to a new sub-trace"
                 )
-                logger.info(f"current sub-trace count: {trace.sub_trace_count + 1}")
+                logger.info(f"current sub-trace count: {trace.sub_trace_count}")
                 return ()
             else:
                 logger.info(
@@ -189,7 +189,7 @@ def get_selection(self, trace: Trace) -> tuple[int, ...]:
 
             if sota_count < self.SOTA_COUNT_THRESHOLD:
                 # Check if we've reached the maximum number of traces before creating a new one
-                if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
+                if trace.sub_trace_count >= self.MAX_TRACE_NUM:
                     logger.info(
                         f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
                     )
@@ -216,7 +216,7 @@ def get_selection(self, trace: Trace) -> tuple[int, ...]:
                         return (last_second_sota_idx,)
                     else:
                         # Check max trace limit again before creating a new trace
-                        if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
+                        if trace.sub_trace_count >= self.MAX_TRACE_NUM:
                             logger.info(
                                 f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
                             )
@@ -226,7 +226,7 @@ def get_selection(self, trace: Trace) -> tuple[int, ...]:
                         logger.info(
                             f"SOTA count {sota_count} is below threshold {self.SOTA_COUNT_THRESHOLD}, jump a new sub-trace"
                         )
-                        logger.info(f"current sub-trace count: {trace.sub_trace_count + 1}")
+                        logger.info(f"current sub-trace count: {trace.sub_trace_count}")
                         return ()  # reboot a new sub-trace
 
             else: