Method _train_step

agentlightning/verl/trainer.py:239–436

(self, batch_dict: dict)

        return ref_worker.compute_ref_log_prob(batch)

    def _train_step(self, batch_dict: dict) -> dict:
        # Isolate in a separate method to automatically recycle the variables before validation.
        batch: DataProto = DataProto.from_single_dict(batch_dict)
        metrics = {}
        timing_raw = {}

        with _timer("step", timing_raw):

            # When agent mode is enabled, we read the batch as it is.
            gen_batch = batch

            # generate a batch
            with _timer("gen", timing_raw):
                self.async_rollout_manager.wake_up()
                self.agent_mode_daemon.set_up_data_and_server(
                    gen_batch.non_tensor_batch, self.async_rollout_manager.server_addresses
                )
                self.agent_mode_daemon.run_until_all_finished()
                batch, agent_metrics = self.agent_mode_daemon.get_train_data_batch(
                    max_prompt_length=(
                        self.config.agentlightning.trace_aggregator.trajectory_max_prompt_length
                        if self.config.agentlightning.trace_aggregator.level.startswith("trajectory")
                        else self.config.data.max_prompt_length
                    ),
                    max_response_length=(
                        self.config.agentlightning.trace_aggregator.trajectory_max_response_length
                        if self.config.agentlightning.trace_aggregator.level.startswith("trajectory")
                        else self.config.data.max_response_length
                    ),
                    device=gen_batch.batch["fake_ids"].device,
                    global_steps=self.global_steps,
                )
                metrics.update(agent_metrics)
                self.agent_mode_daemon.clear_data_and_server()
                self.async_rollout_manager.sleep()

            if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:
                with _timer("gen_max", timing_raw):
                    gen_baseline_batch = deepcopy(gen_batch)
                    gen_baseline_batch.meta_info["do_sample"] = False
                    gen_baseline_output = self.async_rollout_manager.generate_sequences(gen_baseline_batch)

                    batch = batch.union(gen_baseline_output)
                    reward_baseline_tensor = self.reward_fn(batch)
                    reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1)

                    batch.pop(batch_keys=list(gen_baseline_output.batch.keys()))

                    batch.batch["reward_baselines"] = reward_baseline_tensor

                    del gen_baseline_batch, gen_baseline_output

            # uid is used for algorithm like GRPO, should be aligned to data id
            batch.non_tensor_batch["uid"] = batch.non_tensor_batch["data_id_list"]

            if "response_mask" not in batch.batch:
                batch.batch["response_mask"] = compute_response_mask(batch)