Remove dead code with direct config access

src/reasoning/step_level_cot.py

@@ -636,79 +636,6 @@ class StepLevelCoTTrainer:
         
         progress_bar.close()
         logger.info(f"PRM training completed - Avg loss: {total_loss / max(global_step, 1):.4f}")
-        
-        prm_trainer = PRMTrainer(
-            model=self.prm,
-            learning_rate=learning_rate,  # Use the passed parameter
-            weight_decay=self.config.weight_decay,
-            warmup_steps=self.config.warmup_steps,
-        )
-        
-        train_dataloader = torch.utils.data.DataLoader(
-            train_dataset,
-            batch_size=self.config.train_batch_size,
-            shuffle=True,
-            num_workers=self.config.dataloader_num_workers,
-        )
-        
-        # Prepare with accelerator
-        prm_trainer.model, prm_trainer.optimizer, train_dataloader = self.accelerator.prepare(
-            prm_trainer.model, prm_trainer.optimizer, train_dataloader
-        )
-        
-        best_val_loss = float('inf')
-        global_step = 0
-        
-        # Training loop using steps instead of epochs
-        prm_trainer.model.train()
-        
-        progress_bar = tqdm(
-            total=max_steps,
-            desc=f"PRM Training",
-            disable=not self.accelerator.is_local_main_process,
-        )
-        
-        while global_step < max_steps:
-            for batch in train_dataloader:
-                if global_step >= max_steps:
-                    break
-                    
-                # Extract batch data
-                vision_features = batch['visual_features']
-                step_embeddings = batch['step_descriptions']  # Need to encode these
-                target_rewards = batch['step_rewards']
-                
-                # Encode step descriptions
-                # (In practice, you'd encode text properly)
-                
-                metrics = prm_trainer.train_step(
-                    vision_features,
-                    step_embeddings,
-                    target_rewards,
-                )
-                
-                global_step += 1
-                progress_bar.update(1)
-                progress_bar.set_postfix({'loss': metrics['loss'], 'step': global_step})
-                
-                # Save checkpoint
-                if global_step % save_steps == 0:
-                    logger.info(f"Saving checkpoint at step {global_step}")
-                    self._save_prm(global_step)
-                
-                # Validation
-                if eval_steps > 0 and global_step % eval_steps == 0 and val_dataset is not None:
-                    val_metrics = self._evaluate_prm(prm_trainer, val_dataset)
-                    logger.info(f"Step {global_step} - Validation: {val_metrics}")
-                    
-                    # Save best model
-                    if val_metrics['mse'] < best_val_loss:
-                        best_val_loss = val_metrics['mse']
-                        self._save_prm(f"best_step_{global_step}")
-        
-        progress_bar.close()
-        
-        logger.info("PRM training completed")
     
     def train_rl(
         self,
@@ -811,14 +738,6 @@ class StepLevelCoTTrainer:
         
         progress_bar.close()
         logger.info(f"RL training completed - Avg reward: {total_reward / max(global_step, 1):.4f}")
-        
-        self.rl_trainer.train(
-            train_dataset=train_dataset,
-            num_iterations=num_iterations,
-            log_interval=self.config.logging_steps,
-        )
-        
-        logger.info("RL training completed")
     
     def evaluate_inference_scaling(
         self,