def run(self):
    start_step = self.global_step
    start_time = time.time()
    for _step in range(start_step, self.max_steps):
        # set global step
        self.actor.set_global_step(self.global_step)
        self.rollout.set_global_step(self.global_step)

        with self.timer("step"):
            with self.timer("sync_weights"):
                if _step % self.weight_sync_interval == 0:
                    self.update_rollout_weights()
            with self.timer("generate_rollouts"):
                env_handle: Handle = self.env.interact(
                    input_channel=self.env_channel,
                    rollout_channel=self.rollout_channel,
                    reward_channel=self.reward_channel,
                    actor_channel=self.actor_channel,
                )
                rollout_handle: Handle = self.rollout.generate(
                    input_channel=self.rollout_channel,
                    output_channel=self.env_channel,
                )
                if self.reward is not None:
                    reward_handle: Handle = self.reward.compute_rewards(
                        input_channel=self.reward_channel,
                        output_channel=self.env_channel,
                    )
                self.actor.recv_rollout_trajectories(
                    input_channel=self.actor_channel
                ).wait()
                rollout_handle.wait()
                if self.reward is not None:
                    reward_handle.wait()

            # compute advantages and returns.
            with self.timer("cal_adv_and_returns"):
                actor_rollout_metrics = (
                    self.actor.compute_advantages_and_returns().wait()
                )

            # actor training.
            actor_training_handle: Handle = self.actor.run_training()
            env_bootstrap_handle: Handle | None = None
            if self.overlap_env_bootstrap and _step + 1 < self.max_steps:
                env_bootstrap_handle = self.env.prefetch_train_bootstrap(
                    rollout_channel=self.rollout_channel
                )

         # Some metrics code
         # .......

    self.metric_logger.finish()

    # Stop logging thread
    self.stop_logging = True
    self.log_queue.join()  # Wait for all queued logs to be processed
    self.log_thread.join(timeout=1.0)

class CommMapper:
  # 计算目的 worker rank 的分配逻辑，用户数据发送
  def get_dst_ranks(
          batch_size: int, src_world_size: int, dst_world_size: int, src_rank: int
      ) -> list[tuple[int, int]]:
          """Compute destination ranks and transfer sizes for one source rank."""
          # 要求数据的batch_size必须被 两组 worker平均分配
          assert batch_size % src_world_size == 0, (
              f"batch_size ({batch_size}) must be divisible by src_world_size ({src_world_size})."
          )
          assert batch_size % dst_world_size == 0, (
              f"batch_size ({batch_size}) must be divisible by dst_world_size ({dst_world_size})."
          )
          assert 0 <= src_rank < src_world_size, (
              f"src_rank ({src_rank}) must be in [0, {src_world_size})."
          )
  
          batch_size_per_src_rank = batch_size // src_world_size
          batch_size_per_dst_rank = batch_size // dst_world_size
  
          dst_ranks_and_sizes: list[tuple[int, int]] = []
          # 根据src_rank的batch大小计算分割区间
          batch_begin = src_rank * batch_size_per_src_rank
          batch_end = (src_rank + 1) * batch_size_per_src_rank
          while batch_begin < batch_end:
              # 将数据切分至目标 des_rank 中，并截取数据大小，存入dst_ranks_and_sizes列表
              dst_rank = batch_begin // batch_size_per_dst_rank
              dst_batch_begin = dst_rank * batch_size_per_dst_rank
              dst_remaining = batch_size_per_dst_rank - (batch_begin - dst_batch_begin)
              src_remaining = batch_end - batch_begin
              dst_size = min(dst_remaining, src_remaining)
              dst_ranks_and_sizes.append((dst_rank, dst_size))
              batch_begin += dst_size
          return dst_ranks_and_sizes
  
  # 计算源 worker rank 的分配逻辑，用于数据返回
  def get_src_ranks(
          batch_size: int, src_world_size: int, dst_world_size: int, dst_rank: int
      ) -> list[tuple[int, int]]:
          """Compute source ranks/sizes for one destination rank."""
          assert batch_size % src_world_size == 0, (
              f"batch_size ({batch_size}) must be divisible by src_world_size ({src_world_size})."
          )
          assert batch_size % dst_world_size == 0, (
              f"batch_size ({batch_size}) must be divisible by dst_world_size ({dst_world_size})."
          )
          assert 0 <= dst_rank < dst_world_size, (
              f"dst_rank ({dst_rank}) must be in [0, {dst_world_size})."
          )
  
          src_ranks_and_sizes: list[tuple[int, int]] = []
          for src_rank in range(src_world_size):
              dst_ranks_and_sizes = CommMapper.get_dst_ranks(
                  batch_size=batch_size,
                  src_world_size=src_world_size,
                  dst_world_size=dst_world_size,
                  src_rank=src_rank,
              )
              for mapped_dst_rank, size in dst_ranks_and_sizes:
                  if mapped_dst_rank == dst_rank:
                      src_ranks_and_sizes.append((src_rank, size))
  
          expected_size = batch_size // dst_world_size
          actual_size = sum(size for _, size in src_ranks_and_sizes)
          assert actual_size == expected_size, (
              f"Expected receive size {expected_size} for destination rank {dst_rank}, "
              f"got {actual_size} from mappings {src_ranks_and_sizes}."
          )
          return src_ranks_and_sizes

@dataclass(kw_only=True)
class EnvOutput:
    """Environment output for a single chunk step."""
    # 仿真实例，是一个[str, tensor]的 map
    obs: dict[str, Any]
    # 当仿真结束(正常、异常)时，所观察到的 obs
    final_obs: Optional[dict[str, Any]] = None
    # dones 代表仿真正常结束的状态矩阵：terminations OR truncations
    # terminations 代表 自然终止 的状态矩阵：比如任务成功/失败、掉落、碰撞等环境定义的终止条件触发
    # truncations 代表 截断终止 的状态矩阵：比如到达 time limit、外部强制中止、达到最大步数等“非任务本身终止”的原因
    dones: Optional[torch.Tensor] = None  # [B, T]
    terminations: Optional[torch.Tensor] = None  # [B, T]
    truncations: Optional[torch.Tensor] = None  # [B, T]
    rewards: Optional[torch.Tensor] = None  # [B, T]

# truncations: shape [B, T] = [3, 4]
truncations = torch.tensor([
    [False, False, False, False],  # env0 这 4 步都没截断
    [False, True,  False, False],  # env1 在第 2 步发生过截断（中途）
    [False, False, False, True ],  # env2 在最后一步发生截断
], dtype=torch.bool)
# 然而，我们在env worker只取T最后一步的结果作为该step的状态，
# 因此env worker的step代码中会看到下面的写法
last_step_truncations = truncations[:, -1]
# 结果: tensor([False, False, True])

async def recv_env_output(
    self, input_channel: Channel, mode: Literal["train", "eval"] = "train"
) -> dict[str, Any]:
    """Receive env outputs from mapped env ranks and merge if needed.

    Args:
        input_channel: Channel carrying env->rollout outputs.
        mode: Rollout mode, either ``"train"`` or ``"eval"``.

    Returns:
        A single env output dict. When multiple env ranks are mapped to this
        rollout worker, outputs are merged on batch dimension.
    """
    assert mode in ["train", "eval"], f"{mode=} is not supported"
    src_ranks_and_sizes = self.src_ranks[mode]
    obs_batches = []
    # 通过上文计算的对应关系，从目标 env worker 中获取数据
    for src_rank, expected_size in src_ranks_and_sizes:
        obs_batch = await input_channel.get(
            key=CommMapper.build_channel_key(
                src_rank, self._rank, extra=f"{mode}_obs"
            ),
            async_op=True,
        ).async_wait()
        actual_size = self._infer_env_batch_size(obs_batch)
        assert actual_size == expected_size, (
            f"Expected env output batch size {expected_size} from env rank {src_rank}, "
            f"got {actual_size}."
        )
        obs_batches.append(obs_batch)
    # 然后将获取的目标数据数据中的 tensor 进行merge，组成一个大的结构体，方便一次性推理
    # {"obs": merged_obs, "final_obs": merged_final_obs}
    return self._merge_obs_batches(obs_batches)

def send_chunk_actions(
    self,
    output_channel: Channel,
    chunk_actions: torch.Tensor | np.ndarray,
    mode: Literal["train", "eval"] = "train",
):
    """Send action shards to mapped env ranks.

    Args:
        output_channel: Channel carrying rollout->env action chunks.
        chunk_actions: Predicted action chunk batch (tensor or ndarray).
        mode: Rollout mode, either ``"train"`` or ``"eval"``.
    """
    assert mode in ["train", "eval"], f"{mode=} is not supported"
    dst_ranks_and_sizes = self.dst_ranks[mode]
    # 从 dst_ranks中获取，每个 env worker 需要发送多大 batch size 是数据，并进行分割发送
    split_sizes = [size for _, size in dst_ranks_and_sizes]
    chunk_actions_split = self._split_actions(chunk_actions, split_sizes)
    for (dst_rank, _), chunk_action_i in zip(
        dst_ranks_and_sizes, chunk_actions_split
    ):
        if isinstance(chunk_action_i, torch.Tensor):
            chunk_action_i = (
                chunk_action_i.detach().cpu().contiguous()
            )  # for evaluation
        output_channel.put(
            chunk_action_i,
            key=CommMapper.build_channel_key(
                self._rank, dst_rank, extra=f"{mode}_actions"
            ),
            async_op=True,
        )

sync def recv_rollout_trajectories(self, input_channel: Channel) -> None:
    """
    Receive rollout trajectories from rollout workers.

    Args:
        input_channel: The input channel to read from.
    """
    clear_memory(sync=False)
    # 计算这个 actor rank worker 的 input_channel 有多少批数据存在
    send_num = self._component_placement.get_world_size("env") * self.stage_num
    recv_num = self._component_placement.get_world_size("actor")
    split_num = compute_split_num(send_num, recv_num)

    recv_list = []
    # 每批数据取一次
    for _ in range(split_num):
        trajectory: Trajectory = await input_channel.get(async_op=True).async_wait()
        recv_list.append(trajectory)
    # 把trajectories list 合并成一个大的 dictionary，其 value 为合并的 tensor
    self.rollout_batch = convert_trajectories_to_batch(recv_list)
    # 如果是 not auto_reset模式，这里需要对 reward 进行过滤，并改变 reward 的 shape，
    # 方便后面的 adv 的计算，同时也生成loss_mask矩阵，在 adv 计算屏蔽需要忽略的 values
    self.rollout_batch = self._process_received_rollout_batch(self.rollout_batch)

def compute_advantages_and_returns(self) -> dict[str, torch.Tensor]:
    """
    Compute the advantages and returns.
    """
    # 这里转换为标准的输入参数
    kwargs = {
        "task_type": self.cfg.runner.task_type,
        "adv_type": self.cfg.algorithm.adv_type,
        "rewards": self.rollout_batch["rewards"],
        "dones": self.rollout_batch["dones"],
        "values": self.rollout_batch.get("prev_values", None),
        "gamma": self.cfg.algorithm.get("gamma", 1),
        "gae_lambda": self.cfg.algorithm.get("gae_lambda", 1),
        "group_size": self.cfg.algorithm.get("group_size", 8),
        "reward_type": self.cfg.algorithm.reward_type,
        "loss_mask": self.rollout_batch.get("loss_mask", None),
        "loss_mask_sum": self.rollout_batch.get("loss_mask_sum", None),
    }
    # adv的核心计算函数
    advantages_and_returns = calculate_adv_and_returns(**kwargs)

    # 统计rollout的metrics数据
    self.rollout_batch.update(advantages_and_returns)
    if kwargs["loss_mask"] is not None:
        self.rollout_batch.update({"loss_mask": kwargs["loss_mask"]})
    if kwargs["loss_mask_sum"] is not None:
        self.rollout_batch.update({"loss_mask_sum": kwargs["loss_mask_sum"]})
    # 获得rewards的统计信息，并打印
    rollout_metrics = compute_rollout_metrics(self.rollout_batch)
    return rollout_metrics

def calculate_adv_and_returns(**kwargs) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
    """
    Unified entry for advantage + return computation.
    Accepts variable keyword arguments, preprocesses them, then dispatches
    to specific algorithm via registry.
    """
    adv_type = kwargs["adv_type"]
    # 查找gae的计算方法
    fn = get_adv_and_returns(adv_type)

    task_type = kwargs["task_type"]
    # 我们主要看embodied类型的计算，LLM的计算和 VLA 差别较大，暂时忽略
    if task_type == "embodied":
        # 1. 如果计算 rewards 是通过chunk_level方式，即我们不需要统计每一步的action的结果
        #    那么我们需要把rewards、dones、loss_mask、loss_mask_sum等值压缩成一维
        # 2. 获取最新的 reward 的shape[num_chunk，batch_size, chunk_size]
        #    n_steps = num_chunk * chunk_size
        kwargs = preprocess_embodied_advantages_inputs(**kwargs)
        # 除了 gae(PPO) 外，还有其他算法: grpo、grpo_dynamic、reinpp、raw
        if adv_type != "gae":
            kwargs = calculate_scores(**kwargs)
        advantages, returns = fn(**kwargs)
        # 将结果展开维一维数据，进行 batch training
        res = postprocess_embodied_advantages_outputs(
            advantages=advantages, returns=returns, **kwargs
        )
    else:
        # reasoning tasks
        kwargs = preprocess_reasoning_advantages_inputs(**kwargs)
        advantages, returns = fn(**kwargs)
        res = postprocess_reasoning_advantages_outputs(advantages, returns)
    return res

@register_advantage("gae")
def compute_gae_advantages_and_returns(
    rewards: torch.Tensor,
    gamma: float = 1.0,
    gae_lambda: float = 1.0,
    values: Optional[torch.Tensor] = None,
    normalize_advantages: bool = True,
    normalize_returns: bool = False,
    loss_mask: Optional[torch.Tensor] = None,
    dones: Optional[torch.Tensor] = None,
    **kwargs,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Calculate advantages and returns for Proximal Policy Optimization (PPO).
    NOTE: currently this function does not support auto-reset.

    This function implements Generalized Advantage Estimation (GAE) to compute
    advantages and returns for PPO training. The advantages are normalized
    using mean and standard deviation for stable training.

    Args:
        rewards (torch.Tensor): Rewards per timestep. Shape: [seq_len, bsz].
        values (torch.Tensor): Value function estimates. Shape: [seq_len, bsz].
        dones (torch.Tensor): Done flags (1 if episode ended, else 0).
        gamma (float, optional): Discount factor. Defaults to 1.0.
        gae_lambda (float, optional): GAE smoothing factor. Defaults to 1.0.
        normalize_advantages (bool, optional): Whether to normalize advantages. Defaults to True.
        normalize_returns (bool, optional): Whether to normalize returns. Defaults to False.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: (advantages, returns)
    """
    T = rewards.shape[0]
    advantages = torch.zeros_like(rewards)
    returns = torch.zeros_like(rewards)
    gae = 0
    # 如果没有value，则不是 critic 计算方式，不需要加权平均
    critic_free = values is None
    if critic_free:
        gae_lambda = 1
        gamma = 1
    # 从后向前进行训练计算，rewards为T维，values和dones为T+1维
    for step in reversed(range(T)):
        if critic_free:
            delta = rewards[step]
        else:
            delta = (
                rewards[step]
                + gamma * values[step + 1] * (~dones[step + 1])
                - values[step]
            )

        gae = delta + gamma * gae_lambda * (~dones[step + 1]) * gae
        returns[step] = gae if critic_free else gae + values[step]

    advantages = returns - values[:-1] if not critic_free else returns
    # 对 loss_mask没有屏蔽的 adv 做归一化，为了让训练的只标准化，稳定训练，
    if normalize_advantages:
        advantages = safe_normalize(advantages, loss_mask=loss_mask)
    # 同样return也可以做归一化
    if normalize_returns:
        returns = safe_normalize(returns, loss_mask=loss_mask)

    return advantages, returns

def run_training(self) -> None:
    """
    Run the training process using the received rollout batch.
    """
    # 权重加载回 GPU
    if self.is_weight_offloaded:
        self.load_param_and_grad(self.device)
    if self.is_optimizer_offloaded:
        self.load_optimizer(self.device)
    # 切换训练模式
    self.model.train()
    rollout_size = (
        self.rollout_batch["prev_logprobs"].shape[0]
        * self.rollout_batch["prev_logprobs"].shape[1]
    )
    # 生成logprobs的维度矩阵，并对self.rollout_batch进行随机 shuffle，提升训练稳定性
    g = torch.Generator()
    g.manual_seed(self.cfg.actor.seed + self._rank)
    # 生成 [0, rollout_size) 的随机排列
    shuffle_id = torch.randperm(rollout_size, generator=g)

    with torch.no_grad():
        self.rollout_batch = process_nested_dict_for_train(
            self.rollout_batch, # rollout batch（常见为 [T, B, ...]）
            shuffle_id # 展平后的随机索引
        )

    # 检查全局 batch 是否能被 micro_batch 与 world_size 整除
    # 需要整除才能保证梯度累积步数为整数
    assert (
        self.cfg.actor.global_batch_size # 全局 batch size（跨所有 DP rank）
        % (self.cfg.actor.micro_batch_size * self._world_size) # 每次累积的总微批容量
        == 0
    ), "global_batch_size is not divisible by micro_batch_size * world_size"

    # 计算梯度累积步数（每个 rank）
    self.gradient_accumulation = (
        self.cfg.actor.global_batch_size
        // self.cfg.actor.micro_batch_size
        // self._world_size
    )

    # 将展平数据切分成可迭代的 mini-batch
    # PPO 通常对同一批数据做多轮更新. https://arxiv.org/abs/1707.06347
    rollout_size = self.rollout_batch["prev_logprobs"].size(0)
    batch_size_per_rank = self.cfg.actor.global_batch_size // self._world_size
    assert rollout_size % batch_size_per_rank == 0, (
        f"{rollout_size} is not divisible by {batch_size_per_rank}"
    )
    metrics = {}
    update_epoch = self.cfg.algorithm.get("update_epoch", 1)
    for _ in range(update_epoch):
        # 把展平后的数据切成若干个 per-rank global batch
        rollout_dataloader_iter = split_dict_to_chunk(
            self.rollout_batch,
            rollout_size // batch_size_per_rank,
        )
        for train_global_batch in rollout_dataloader_iter:
            # 将 global batch 再切成 micro batch 做梯度累积
            train_global_batch_size = train_global_batch["prev_logprobs"].shape[0]
            # 校验 global batch 的大小与配置匹配
            assert (
                train_global_batch_size
                == self.cfg.actor.global_batch_size
                // torch.distributed.get_world_size()
            )
            # 校验 micro batch 可整除
            assert train_global_batch_size % self.cfg.actor.micro_batch_size == 0, (
                f"{train_global_batch_size=}, {self.cfg.actor.micro_batch_size}"
            )
            # 将该 global batch 切成 micro batch
            train_micro_batch = split_dict_to_chunk(
                train_global_batch,
                train_global_batch_size // self.cfg.actor.micro_batch_size,
            )
            # 清空优化器梯度，为本 global batch 的累积做准备
            self.optimizer.zero_grad()
            # 遍历每个 micro batch
            for idx, batch in enumerate(train_micro_batch):
                # 将该 micro batch 的所有 tensor 移动到本地 GPU
                batch = put_tensor_device(
                    batch,
                    f"{Worker.torch_device_type}:{int(os.environ['LOCAL_RANK'])}",
                )
                # 设置 FSDP/no_sync 等上下文（最后一次 micro batch 才同步）
                backward_ctx = self.before_micro_batch(
                    self.model,
                    is_last_micro_batch=(idx + 1) == self.gradient_accumulation,
                )
                advantages = batch["advantages"]
                prev_logprobs = batch["prev_logprobs"]
                returns = batch.get("returns", None)
                prev_values = batch.get("prev_values", None)
                loss_mask = batch.get("loss_mask", None)
                loss_mask_sum = batch.get("loss_mask_sum", None)

                forward_inputs = batch.get("forward_inputs", None)
                # 针对不同 model_type 的额外前向参数
                kwargs = {}
                if SupportedModel(self.cfg.actor.model.model_type) in [
                    SupportedModel.OPENVLA,
                    SupportedModel.OPENVLA_OFT,
                ]:
                    kwargs["temperature"] = (
                        self.cfg.algorithm.sampling_params.temperature_train
                    )
                    kwargs["top_k"] = self.cfg.algorithm.sampling_params.top_k
                elif (
                    SupportedModel(self.cfg.actor.model.model_type)
                    == SupportedModel.GR00T
                ):
                    kwargs["prev_logprobs"] = prev_logprobs
                # 是否需要在 actor 前向中计算 values
                compute_values = (
                    True if self.cfg.algorithm.adv_type == "gae" else False
                )
                # 自动混合精度上下文（降低显存/加速）
                with self.amp_context:
                    output_dict = self.model(
                        forward_inputs=forward_inputs, # rollout的直接输入
                        compute_logprobs=True, # 需要当前策略 logprob（PPO loss）
                        compute_entropy=self.cfg.algorithm.entropy_bonus > 0, # 是否启用熵奖励则计算 entropy
                        compute_values=compute_values,
                        use_cache=False, # 训练禁用 cache，避免显存增长
                        **kwargs,
                    )

                # GR00T会在前向中重算/更新 prev_logprobs
                if (
                    SupportedModel(self.cfg.actor.model.model_type)
                    == SupportedModel.GR00T
                ):
                    prev_logprobs = output_dict["prev_logprobs"]

                kwargs = {  # 构造统一 loss 入口所需参数字典
                    "loss_type": self.cfg.algorithm.loss_type,  # 策略/价值损失类型
                    "logprob_type": self.cfg.algorithm.logprob_type,  # logprob 的组织方式（action/token）
                    "reward_type": self.cfg.algorithm.reward_type,  # reward 的组织方式（action/chunk）
                    "single_action_dim": self.cfg.actor.model.get("action_dim", 7),  # 单步 action 维度
                    "logprobs": output_dict["logprobs"],  # 当前策略 logprobs
                    "values": output_dict.get("values", None),  # 当前 value 预测（可能为 None）
                    "old_logprobs": prev_logprobs,  # 旧策略 logprobs（重要性采样）
                    "advantages": advantages,  # 优势
                    "returns": returns,  # 回报（用于 value loss）
                    "prev_values": prev_values,  # 旧 values（用于 value clipping）
                    "clip_ratio_high": self.cfg.algorithm.clip_ratio_high,  # PPO clip 上界
                    "clip_ratio_low": self.cfg.algorithm.clip_ratio_low,  # PPO clip 下界
                    "value_clip": self.cfg.algorithm.get("value_clip", None),  # value clip 系数（可选）
                    "huber_delta": self.cfg.algorithm.get("huber_delta", None),  # huber loss delta（可选）
                    "loss_mask": loss_mask,  # 有效位置 mask（屏蔽无效步）
                    "loss_mask_sum": loss_mask_sum,  # mask 长度/计数（用于统计/缩放）
                    "max_episode_steps": self.cfg.env.train.max_episode_steps,  # 环境最大步数（用于某些指标/缩放）
                    "task_type": self.cfg.runner.task_type,  # task 类型（embodied/reasoning）
                    "critic_warmup": (  # critic warmup 期间可能只训练 value 或跳过某些项
                        self.optimizer_steps < self.critic_warmup_steps  # 是否仍处于 warmup 步数内
                    ),  # 得到布尔值
                }  # loss 计算参数字典构造完成
                # 计算损失并返回一组指标
                loss, metrics_data = policy_loss(**kwargs)
                # 初始化熵项（默认 0）
                entropy_loss = torch.tensor(
                    0.0, device=Worker.torch_platform.current_device()
                )
                # 若启用熵奖励且不在 critic warmup 期
                if (
                    self.cfg.algorithm.entropy_bonus > 0
                    and not kwargs["critic_warmup"]
                ):
                    # 满足则计算熵项
                    entropy = output_dict["entropy"]  # 从模型输出取 entropy
                    entropy = reshape_entropy(  # 根据 entropy_type/action_dim 重新整理形状
                        entropy,  # 输入 entropy
                        entropy_type=self.cfg.algorithm.entropy_type,  # 熵的组织方式
                        action_dim=self.cfg.actor.model.get("action_dim", 7),  # action 维度
                        batch_size=output_dict["logprobs"].shape[0],  # batch 大小
                    )  # 输出 reshape 后 entropy
                    entropy_loss = masked_mean(entropy, mask=loss_mask)  # 在有效位置上求熵均值
                    loss -= self.cfg.algorithm.entropy_bonus * entropy_loss  # 将熵奖励加入总 loss（减去负熵）
                metrics_data["actor/entropy_loss"] = entropy_loss.detach().item()  # 记录熵项指标
                # 若启用 SFT 协同训练
                if self.enable_sft_co_train:
                    # 在同一次迭代里加入 SFT loss/指标
                    self._train_sft_epoch(metrics_data, loss)
                # 按累积步数缩放 loss，保证等效于大 batch
                loss /= self.gradient_accumulation
                # 在 FSDP/no_sync 上下文里做反向传播
                with backward_ctx:
                    self.grad_scaler.scale(loss).backward()
                # 记录缩放后的 total_loss
                metrics_data["actor/total_loss"] = loss.detach().item()
                # 将本 micro batch 的指标追加到 metrics 容器
                append_to_dict(metrics, metrics_data)
            # micro batch 结束后清理缓存，降低峰值显存
            self.torch_platform.empty_cache()
            # 执行一次优化器 step，并返回梯度范数与学习率
            grad_norm, lr_list = self.optimizer_step()
            # 构造优化相关指标
            data = {
                "actor/grad_norm": grad_norm,
                "actor/lr": lr_list[0],
            }
            if len(lr_list) > 1:
                # 记录 critic 学习率
                data["critic/lr"] = lr_list[1]
            append_to_dict(metrics, data)
    # 更新学习率调度器
    self.lr_scheduler.step()
    # 再次清空梯度，避免残留
    self.optimizer.zero_grad()
    # 释放临时张量/缓存（自定义工具）
    clear_memory()
    # 对指标列表取均值
    mean_metric_dict = {key: np.mean(value) for key, value in metrics.items()}
    # 跨分布式 rank 聚合指标（取平均）
    mean_metric_dict = all_reduce_dict(
        mean_metric_dict, op=torch.distributed.ReduceOp.AVG
    )
    # 返回本次训练迭代的全局平均指标
    return mean_metric_dict