ray.rllib.execution.train_ops.multi_gpu_train_one_step#

ray.rllib.execution.train_ops.multi_gpu_train_one_step(algorithm, train_batch) → Dict[源代码]#

多GPU版本的train_one_step。

使用策略的 load_batch_into_buffer 和 learn_on_loaded_batch 方法，以提高CPU/GPU数据传输的效率。例如，当对训练批次进行多次传递（例如用于PPO）时，使用 config.num_sgd_iter，实际的训练批次仅被分割一次并加载到GPU中一次。

from ray.rllib.execution.rollout_ops import synchronous_parallel_sample
algo = [...]
train_batch = synchronous_parallel_sample(algo.env_runner_group)
# This trains the policy on one batch.
print(multi_gpu_train_one_step(algo, train_batch)))

{"default_policy": ...}

更新算法实例的 NUM_ENV_STEPS_TRAINED 和 NUM_AGENT_STEPS_TRAINED 计数器，以及 LOAD_BATCH_TIMER 和 LEARN_ON_BATCH_TIMER 计时器。