diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/ckp_select.py b/rdagent/scenarios/data_science/proposal/exp_gen/ckp_select.py index 4169bd37..d84c0f49 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/ckp_select.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/ckp_select.py @@ -70,7 +70,7 @@ class LimitTimeCKPSelector(CheckpointSelector): 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,7 +81,7 @@ class LimitTimeCKPSelector(CheckpointSelector): 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" ) @@ -89,11 +89,11 @@ class LimitTimeCKPSelector(CheckpointSelector): 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 @@ class SOTAJumpCKPSelector(CheckpointSelector): 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 @@ class SOTAJumpCKPSelector(CheckpointSelector): 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 @@ class BackJumpCKPSelector(CheckpointSelector): 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 @@ class BackJumpCKPSelector(CheckpointSelector): 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 @@ class BackJumpCKPSelector(CheckpointSelector): 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: