Beta action distribution for bounded action spaces. More...

#include <LibtorchActionDistribution.h>

Inheritance diagram for Moose::LibtorchBetaActionDistribution:

Public Member Functions
	LibtorchBetaActionDistribution (const std::string &name, unsigned int num_inputs, unsigned int num_outputs, const std::vector< Real > &minimum_values, const std::vector< Real > &maximum_values, torch::DeviceType device_type=torch::kCPU, torch::ScalarType scalar_type=torch::kDouble, bool build_on_construct=true, const std::vector< Real > &output_scaling_factors={})
	Build the Beta action distribution used for bounded controls.

virtual void	initialize (c10::optional< at::Generator > generator=c10::nullopt) override
	Initialize the trainable Beta distribution parameters.

virtual void	reset (const torch::Tensor &input) override
	Refresh the Beta distribution parameters from the latest actor features.

virtual torch::Tensor	sample (c10::optional< at::Generator > generator=c10::nullopt) const override
	Draw a stochastic Beta action sample in physical units.

virtual torch::Tensor	deterministicAction () const override
	Return the Beta mean action used for deterministic evaluation.

virtual torch::Tensor	logProbability (const torch::Tensor &action) const override
	Evaluate the log-probability of an action under the current Beta distribution.

virtual torch::Tensor	entropy () const override
	Compute the entropy of the current Beta distribution.

virtual bool	isBounded () const override
	Return true because the Beta distribution enforces action bounds.

torch::nn::Linear &	alphaModule ()
	Return the Beta alpha head.

const torch::nn::Linear &	alphaModule () const
	Return the Beta alpha head.

torch::nn::Linear &	betaModule ()
	Return the Beta beta head.

const torch::nn::Linear &	betaModule () const
	Return the Beta beta head.

const torch::Tensor &	alphaTensor () const
	Return the cached Beta alpha tensor.

const torch::Tensor &	betaTensor () const
	Return the cached Beta beta tensor.

void	synchronizeScalingFactorsFromBuffer ()
	Sync cached scaling metadata from the registered buffers after loading state.

Protected Member Functions
torch::Tensor	prepareFeatures (const torch::Tensor &input) const
	Convert actor features to the configured device and scalar type.

torch::Tensor	prepareAction (const torch::Tensor &action) const
	Convert actions to the configured device and scalar type.

const torch::Tensor &	actionScaleTensor () const
	Return the registered tensor that stores per-action scaling factors.

Protected Attributes
const std::string	_name
	Module name used for registration and serialization.

const unsigned int	_num_inputs
	Number of actor features feeding this distribution.

const unsigned int	_num_outputs
	Number of action dimensions produced by this distribution.

const torch::DeviceType	_device_type
	Torch device used by this distribution.

const torch::ScalarType	_data_type
	Torch scalar type used by this distribution.

std::vector< Real >	_output_scaling_factors
	Cached per-action scaling factors applied in physical units.

torch::Tensor	_action_scale_tensor
	Registered libtorch buffer holding the per-action scaling factors.

Private Member Functions
void	constructDistribution ()
	Build and register the Beta distribution heads.

Private Attributes
const std::vector< Real >	_minimum_values
	Lower action bounds in physical units.

const std::vector< Real >	_maximum_values
	Upper action bounds in physical units.

torch::nn::Linear	_alpha_module {nullptr}
	Linear head that produces the Beta alpha parameters.

torch::nn::Linear	_beta_module {nullptr}
	Linear head that produces the Beta beta parameters.

torch::Tensor	_min_tensor
	Tensor form of the lower action bounds.

torch::Tensor	_max_tensor
	Tensor form of the upper action bounds.

torch::Tensor	_alpha_tensor
	Cached Beta alpha parameters from the latest reset.

torch::Tensor	_beta_tensor
	Cached Beta beta parameters from the latest reset.

torch::Tensor	_alpha_beta_tensor
	Cached sum of the alpha and beta parameters.

torch::Tensor	_log_norm
	Cached log normalization factor for the latest Beta distribution state.

torch::Tensor	_mean
	Cached normalized Beta mean from the latest reset.

Detailed Description

Beta action distribution for bounded action spaces.

Definition at line 254 of file LibtorchActionDistribution.h.

Constructor & Destructor Documentation

◆ LibtorchBetaActionDistribution()

Moose::LibtorchBetaActionDistribution::LibtorchBetaActionDistribution	(	const std::string &	name,
		unsigned int	num_inputs,
		unsigned int	num_outputs,
		const std::vector< Real > &	minimum_values,
		const std::vector< Real > &	maximum_values,
		torch::DeviceType	device_type = `torch::kCPU`,
		torch::ScalarType	scalar_type = `torch::kDouble`,
		bool	build_on_construct = `true`,
		const std::vector< Real > &	output_scaling_factors = `{}`
	)

Build the Beta action distribution used for bounded controls.

Parameters

name	Module name used for registration and serialization.
num_inputs	Number of actor features feeding the distribution.
num_outputs	Number of action dimensions produced by the distribution.
minimum_values	Lower action bounds in physical units.
maximum_values	Upper action bounds in physical units.
device_type	Torch device used by the module.
scalar_type	Torch scalar type used by the module.
build_on_construct	Whether to build the torch modules right away.
output_scaling_factors	Optional extra per-action scaling in physical units.

Definition at line 208 of file LibtorchActionDistribution.C.

  : LibtorchActionDistribution(
        name, num_inputs, num_outputs, device_type, data_type, output_scaling_factors),
    _minimum_values(minimum_values),
    _maximum_values(maximum_values)
{
  if (_minimum_values.size() != _num_outputs || _maximum_values.size() != _num_outputs)
    mooseError("The number of minimum_values and maximum_values entries must match the number "
               "of action outputs.");
 
  for (const auto i : make_range(_minimum_values.size()))
    if (!(_maximum_values[i] > _minimum_values[i]))
      mooseError("maximum_values entries must be strictly greater than minimum_values entries.");
 
  auto min_value = _minimum_values;
  LibtorchUtils::vectorToTensor(min_value, _min_tensor);
  _min_tensor = _min_tensor.transpose(0, 1).to(_data_type).to(_device_type);
  auto max_value = _maximum_values;
  LibtorchUtils::vectorToTensor(max_value, _max_tensor);
  _max_tensor = _max_tensor.transpose(0, 1).to(_data_type).to(_device_type);
 
  if (build_on_construct)
    constructDistribution();
}

Member Function Documentation

◆ actionScaleTensor()

const torch::Tensor & Moose::LibtorchActionDistribution::actionScaleTensor ( ) const

inlineprotectedinherited

Return the registered tensor that stores per-action scaling factors.

Definition at line 104 of file LibtorchActionDistribution.h.

104{ return _action_scale_tensor; }

Moose::LibtorchActionDistribution::_action_scale_tensor

torch::Tensor _action_scale_tensor

Registered libtorch buffer holding the per-action scaling factors.

Definition LibtorchActionDistribution.h:120

Referenced by Moose::LibtorchGaussianActionDistribution::deterministicAction(), deterministicAction(), Moose::LibtorchGaussianActionDistribution::entropy(), entropy(), Moose::LibtorchGaussianActionDistribution::logProbability(), logProbability(), Moose::LibtorchGaussianActionDistribution::sample(), and sample().

◆ alphaModule() [1/2]

torch::nn::Linear & Moose::LibtorchBetaActionDistribution::alphaModule ( )

inline

Return the Beta alpha head.

Definition at line 322 of file LibtorchActionDistribution.h.

322{ return _alpha_module; }

Moose::LibtorchBetaActionDistribution::_alpha_module

torch::nn::Linear _alpha_module

Linear head that produces the Beta alpha parameters.

Definition LibtorchActionDistribution.h:361

◆ alphaModule() [2/2]

const torch::nn::Linear & Moose::LibtorchBetaActionDistribution::alphaModule ( ) const

inline

Return the Beta alpha head.

Definition at line 327 of file LibtorchActionDistribution.h.

327{ return _alpha_module; }

◆ alphaTensor()

const torch::Tensor & Moose::LibtorchBetaActionDistribution::alphaTensor ( ) const

inline

Return the cached Beta alpha tensor.

Definition at line 342 of file LibtorchActionDistribution.h.

342{ return _alpha_tensor; }

Moose::LibtorchBetaActionDistribution::_alpha_tensor

torch::Tensor _alpha_tensor

Cached Beta alpha parameters from the latest reset.

Definition LibtorchActionDistribution.h:369

◆ betaModule() [1/2]

torch::nn::Linear & Moose::LibtorchBetaActionDistribution::betaModule ( )

inline

Return the Beta beta head.

Definition at line 332 of file LibtorchActionDistribution.h.

332{ return _beta_module; }

Moose::LibtorchBetaActionDistribution::_beta_module

torch::nn::Linear _beta_module

Linear head that produces the Beta beta parameters.

Definition LibtorchActionDistribution.h:363

◆ betaModule() [2/2]

const torch::nn::Linear & Moose::LibtorchBetaActionDistribution::betaModule ( ) const

inline

Return the Beta beta head.

Definition at line 337 of file LibtorchActionDistribution.h.

337{ return _beta_module; }

◆ betaTensor()

const torch::Tensor & Moose::LibtorchBetaActionDistribution::betaTensor ( ) const

inline

Return the cached Beta beta tensor.

Definition at line 347 of file LibtorchActionDistribution.h.

347{ return _beta_tensor; }

Moose::LibtorchBetaActionDistribution::_beta_tensor

torch::Tensor _beta_tensor

Cached Beta beta parameters from the latest reset.

Definition LibtorchActionDistribution.h:371

◆ constructDistribution()

void Moose::LibtorchBetaActionDistribution::constructDistribution ( )

private

Build and register the Beta distribution heads.

Definition at line 243 of file LibtorchActionDistribution.C.

{
  _alpha_module = register_module(
      "alpha", torch::nn::Linear(torch::nn::LinearOptions(_num_inputs, _num_outputs).bias(true)));
  _beta_module = register_module(
      "beta", torch::nn::Linear(torch::nn::LinearOptions(_num_inputs, _num_outputs).bias(true)));
 
  _alpha_module->to(_device_type, _data_type);
  _beta_module->to(_device_type, _data_type);
}

Referenced by LibtorchBetaActionDistribution().

◆ deterministicAction()

torch::Tensor Moose::LibtorchBetaActionDistribution::deterministicAction ( ) const

overridevirtual

Return the Beta mean action used for deterministic evaluation.

Implements Moose::LibtorchActionDistribution.

Definition at line 288 of file LibtorchActionDistribution.C.

{
  return (_min_tensor + (_max_tensor - _min_tensor) * _mean) * actionScaleTensor();
}

◆ entropy()

torch::Tensor Moose::LibtorchBetaActionDistribution::entropy ( ) const

overridevirtual

Compute the entropy of the current Beta distribution.

Implements Moose::LibtorchActionDistribution.

Definition at line 316 of file LibtorchActionDistribution.C.

{
  const auto log_action_scale = torch::log(torch::abs(actionScaleTensor()));
  const auto scale = torch::clamp_min(_max_tensor - _min_tensor, 1e-8);
  return _log_norm - (_beta_tensor - 1.0) * torch::digamma(_beta_tensor) -
         (_alpha_tensor - 1.0) * torch::digamma(_alpha_tensor) +
         (_alpha_beta_tensor - 2.0) * torch::digamma(_alpha_beta_tensor) + torch::log(scale) +
         log_action_scale;
}

◆ initialize()

void Moose::LibtorchBetaActionDistribution::initialize ( c10::optional< at::Generator > generator = c10::nullopt )

overridevirtual

Initialize the trainable Beta distribution parameters.

Parameters

generator Optional random-number generator used for reproducible initialization.

Implements Moose::LibtorchActionDistribution.

Definition at line 255 of file LibtorchActionDistribution.C.

{
  Moose::orthogonalInitializeTensor(_alpha_module->weight, 1.0, generator);
  torch::nn::init::zeros_(_alpha_module->bias);
 
  Moose::orthogonalInitializeTensor(_beta_module->weight, 1.0, generator);
  torch::nn::init::zeros_(_beta_module->bias);
}

◆ isBounded()

virtual bool Moose::LibtorchBetaActionDistribution::isBounded ( ) const

inlineoverridevirtual

Return true because the Beta distribution enforces action bounds.

Implements Moose::LibtorchActionDistribution.

Definition at line 317 of file LibtorchActionDistribution.h.

317{ return true; }

◆ logProbability()

torch::Tensor Moose::LibtorchBetaActionDistribution::logProbability ( const torch::Tensor & action ) const

overridevirtual

Evaluate the log-probability of an action under the current Beta distribution.

Parameters

action Action tensor in physical units.

Implements Moose::LibtorchActionDistribution.

Definition at line 294 of file LibtorchActionDistribution.C.

{
  const auto scaled_action = prepareAction(action);
  const auto log_action_scale = torch::log(torch::abs(actionScaleTensor()));
  const auto unscaled_action = scaled_action / actionScaleTensor();
  const auto scale = torch::clamp_min(_max_tensor - _min_tensor, 1e-8);
  const auto normalized = (unscaled_action - _min_tensor) / scale;
  const auto clipped = torch::clamp(normalized, 1e-8, 1.0 - 1e-8);
  auto log_prob = (_alpha_tensor - 1.0) * torch::log(clipped) +
                  (_beta_tensor - 1.0) * torch::log1p(-clipped) - _log_norm - torch::log(scale) -
                  log_action_scale;
 
  const auto out_of_bounds = (normalized < 0.0) | (normalized > 1.0);
  if (out_of_bounds.any().item<bool>())
    log_prob = torch::where(out_of_bounds,
                            torch::full_like(log_prob, -std::numeric_limits<Real>::infinity()),
                            log_prob);
 
  return log_prob;
}

◆ prepareAction()

torch::Tensor Moose::LibtorchActionDistribution::prepareAction ( const torch::Tensor & action ) const

protectedinherited

Convert actions to the configured device and scalar type.

Parameters

action Raw action tensor.

Definition at line 92 of file LibtorchActionDistribution.C.

{
  auto scaled_action = action;
  if (_data_type != scaled_action.scalar_type())
    scaled_action = scaled_action.to(_data_type);
  if (_device_type != scaled_action.device().type())
    scaled_action = scaled_action.to(_device_type);
  return scaled_action;
}

Referenced by Moose::LibtorchGaussianActionDistribution::logProbability(), and logProbability().

◆ prepareFeatures()

torch::Tensor Moose::LibtorchActionDistribution::prepareFeatures ( const torch::Tensor & input ) const

protectedinherited

Convert actor features to the configured device and scalar type.

Parameters

input Raw actor feature tensor.

Definition at line 81 of file LibtorchActionDistribution.C.

{
  auto features = input;
  if (_data_type != features.scalar_type())
    features = features.to(_data_type);
  if (_device_type != features.device().type())
    features = features.to(_device_type);
  return features;
}

Referenced by Moose::LibtorchGaussianActionDistribution::reset(), and reset().

◆ reset()

void Moose::LibtorchBetaActionDistribution::reset ( const torch::Tensor & input )

overridevirtual

Refresh the Beta distribution parameters from the latest actor features.

Parameters

input Feature tensor coming from the actor body.

Implements Moose::LibtorchActionDistribution.

Definition at line 265 of file LibtorchActionDistribution.C.

{
  const auto features = prepareFeatures(input);
  // softplus(x) + 1 keeps alpha, beta > 1 so the Beta density stays unimodal. torch::softplus
  // switches to the identity for large logits, where log(exp(x) + 1) would overflow to inf.
  _alpha_tensor = torch::softplus(_alpha_module->forward(features)) + 1.0;
  _beta_tensor = torch::softplus(_beta_module->forward(features)) + 1.0;
 
  _alpha_beta_tensor = torch::clamp_min(_alpha_tensor + _beta_tensor, 1e-8);
  _mean = _alpha_tensor / _alpha_beta_tensor;
  _log_norm = at::lgamma(_alpha_tensor) + at::lgamma(_beta_tensor) - at::lgamma(_alpha_beta_tensor);
}

◆ sample()

torch::Tensor Moose::LibtorchBetaActionDistribution::sample ( c10::optional< at::Generator > generator = c10::nullopt ) const

overridevirtual

Draw a stochastic Beta action sample in physical units.

Parameters

generator Optional random-number generator used for sampling.

Implements Moose::LibtorchActionDistribution.

Definition at line 279 of file LibtorchActionDistribution.C.

{
  const auto alpha_sample = at::_standard_gamma(_alpha_tensor, generator);
  const auto beta_sample = at::_standard_gamma(_beta_tensor, generator);
  const auto sampled = alpha_sample / (alpha_sample + beta_sample);
  return (_min_tensor + (_max_tensor - _min_tensor) * sampled) * actionScaleTensor();
}

◆ synchronizeScalingFactorsFromBuffer()

void Moose::LibtorchActionDistribution::synchronizeScalingFactorsFromBuffer ( )

inherited

Sync cached scaling metadata from the registered buffers after loading state.

Definition at line 73 of file LibtorchActionDistribution.C.

{
  auto action_scale =
      _action_scale_tensor.detach().reshape({-1}).to(torch::kCPU).to(torch::kDouble);
  LibtorchUtils::tensorToVector(action_scale, _output_scaling_factors);
}