LibtorchRLMiniBatchSampler Class Reference

Samples shuffled mini-batches from a flattened on-policy trajectory batch. More...

#include <LibtorchRLMiniBatchSampler.h>

Public Member Functions
std::vector< LibtorchRLMiniBatch >	sample (const LibtorchRLTrajectoryBuffer::TensorBatch &batch, unsigned int batch_size, bool standardize_advantage, c10::optional< at::Generator > generator=c10::nullopt) const
	Shuffle a flattened rollout batch into PPO-sized chunks.

Static Private Member Functions
static void	validateBatch (const LibtorchRLTrajectoryBuffer::TensorBatch &batch)
	Sanity-check that the flattened rollout tensors all line up.
static LibtorchRLMiniBatch	makeMiniBatch (const LibtorchRLTrajectoryBuffer::TensorBatch &batch, const torch::Tensor &indices, bool standardize_advantage)
	Slice one shuffled mini-batch out of the flattened rollout tensors.

Detailed Description

Samples shuffled mini-batches from a flattened on-policy trajectory batch.

Definition at line 46 of file LibtorchRLMiniBatchSampler.h.

Member Function Documentation

makeMiniBatch()

LibtorchRLMiniBatch LibtorchRLMiniBatchSampler::makeMiniBatch ( const LibtorchRLTrajectoryBuffer::TensorBatch &  batch, const torch::Tensor &  indices, bool  standardize_advantage  )

staticprivate

Slice one shuffled mini-batch out of the flattened rollout tensors.

Parameters

batch	Flattened rollout tensors.
indices	Row indices assigned to this mini-batch.
standardize_advantage	Whether to normalize the advantages in this slice.

Definition at line 71 of file LibtorchRLMiniBatchSampler.C.

{
  LibtorchRLMiniBatch mini_batch;
 
  mini_batch.observations = batch.observations.index({indices});
  mini_batch.actions = batch.actions.index({indices});
  mini_batch.old_log_probabilities = batch.log_probabilities.index({indices});
  mini_batch.value_targets = batch.value_targets.index({indices});
  mini_batch.advantages = batch.advantages.index({indices});
 
  if (standardize_advantage)
  {
    const auto std = mini_batch.advantages.std(false);
    mini_batch.advantages = (mini_batch.advantages - mini_batch.advantages.mean()) / (std + 1e-10);
  }
 
  return mini_batch;
}

Referenced by sample().

sample()

std::vector< LibtorchRLMiniBatch > LibtorchRLMiniBatchSampler::sample	(	const LibtorchRLTrajectoryBuffer::TensorBatch &	batch,
		unsigned int	batch_size,
		bool	standardize_advantage,
		c10::optional< at::Generator >	generator = c10::nullopt
	)		const

Shuffle a flattened rollout batch into PPO-sized chunks.

Parameters

batch	Flattened rollout tensors ready for PPO updates.
batch_size	Preferred number of rows per mini-batch.
standardize_advantage	Whether to normalize the advantages inside each mini-batch.
generator	Optional random-number generator used for the row permutation.

Definition at line 20 of file LibtorchRLMiniBatchSampler.C.

{
  std::vector<LibtorchRLMiniBatch> mini_batches;
 
  if (!batch.size())
    return mini_batches;
 
  validateBatch(batch);
 
  const auto effective_batch_size = std::max<unsigned int>(1, batch_size);
  auto permutation =
      at::randperm(batch.size(), generator, torch::TensorOptions().dtype(torch::kLong));
 
  for (std::int64_t batch_begin = 0; batch_begin < batch.size();
       batch_begin += effective_batch_size)
  {
    const auto batch_end = std::min<std::int64_t>(
        batch.size(), batch_begin + static_cast<std::int64_t>(effective_batch_size));
    const auto indices = permutation.narrow(0, batch_begin, batch_end - batch_begin);
    mini_batches.push_back(makeMiniBatch(batch, indices, standardize_advantage));
  }
 
  return mini_batches;
}

Referenced by LibtorchDRLControlTrainer::trainController().

validateBatch()

void LibtorchRLMiniBatchSampler::validateBatch ( const LibtorchRLTrajectoryBuffer::TensorBatch &  batch )

staticprivate

Sanity-check that the flattened rollout tensors all line up.

Parameters

batch

Flattened rollout tensors to validate.

Definition at line 49 of file LibtorchRLMiniBatchSampler.C.

{
  if (!batch.actions.defined() || !batch.log_probabilities.defined() ||
      !batch.value_targets.defined() || !batch.advantages.defined())
    mooseError("RL tensor batches must define observations, actions, log probabilities, value "
               "targets, and advantages before mini-batch sampling.");
 
  const auto batch_size = batch.size();
  const auto validate_rows = [batch_size](const torch::Tensor & tensor, std::string_view name)
  {
    if (!tensor.defined() || tensor.size(0) != batch_size)
      mooseError(
          "RL tensor batch field ", name, " must have the same number of rows as observations.");
  };
 
  validate_rows(batch.actions, "actions");
  validate_rows(batch.log_probabilities, "log_probabilities");
  validate_rows(batch.value_targets, "value_targets");
  validate_rows(batch.advantages, "advantages");
}

Referenced by sample().

---

The documentation for this class was generated from the following files:

• stochastic_tools/include/libtorch/utils/LibtorchRLMiniBatchSampler.h
• stochastic_tools/src/libtorch/utils/LibtorchRLMiniBatchSampler.C