Gaussian action distribution for unbounded action spaces. More...

#include <LibtorchActionDistribution.h>

Inheritance diagram for Moose::LibtorchGaussianActionDistribution:

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

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

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

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

virtual torch::Tensor	deterministicAction () const override
	Return the Gaussian 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 Gaussian distribution.

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

virtual bool	isBounded () const override
	Return false because the Gaussian distribution is unbounded.

bool	stateIndependentStd () const
	Return whether the Gaussian standard deviation ignores the current actor features.

torch::nn::Linear &	meanModule ()
	Return the Gaussian mean head.

const torch::nn::Linear &	meanModule () const
	Return the Gaussian mean head.

torch::nn::Linear &	stdModule ()
	Return the state-dependent Gaussian standard-deviation head.

const torch::nn::Linear &	stdModule () const
	Return the state-dependent Gaussian standard-deviation head.

torch::Tensor &	logStd ()
	Return the state-independent Gaussian log-standard-deviation parameter.

const torch::Tensor &	logStd () const
	Return the state-independent Gaussian log-standard-deviation parameter.

const torch::Tensor &	stdTensor () const
	Return the cached Gaussian standard-deviation 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 Gaussian distribution heads.

Private Attributes
const bool	_state_independent_std
	Whether the Gaussian std ignores the current actor features.

torch::nn::Linear	_mean_module {nullptr}
	Linear head that produces the Gaussian action mean.

torch::nn::Linear	_std_module {nullptr}
	Linear head that produces state-dependent Gaussian log-std inputs.

torch::Tensor	_log_std_parameter
	Trainable state-independent Gaussian log-std.

torch::Tensor	_mean
	Cached Gaussian action mean from the latest reset.

torch::Tensor	_std_tensor
	Cached Gaussian action standard deviation from the latest reset.

torch::Tensor	_log_std_tensor
	Cached Gaussian action log standard deviation from the latest reset.

Detailed Description

Gaussian action distribution for unbounded action spaces.

Definition at line 126 of file LibtorchActionDistribution.h.

Constructor & Destructor Documentation

◆ LibtorchGaussianActionDistribution()

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

Build the Gaussian action distribution used for unbounded 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.
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 per-action scaling applied in physical units.
state_independent_std	Whether the learned std should ignore the current state.

Definition at line 102 of file LibtorchActionDistribution.C.

  : LibtorchActionDistribution(
        name, num_inputs, num_outputs, device_type, data_type, output_scaling_factors),
    _state_independent_std(state_independent_std)
{
  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 deterministicAction(), Moose::LibtorchBetaActionDistribution::deterministicAction(), entropy(), Moose::LibtorchBetaActionDistribution::entropy(), logProbability(), Moose::LibtorchBetaActionDistribution::logProbability(), sample(), and Moose::LibtorchBetaActionDistribution::sample().

◆ constructDistribution()

void Moose::LibtorchGaussianActionDistribution::constructDistribution ( )

private

Build and register the Gaussian distribution heads.

Definition at line 120 of file LibtorchActionDistribution.C.

{
  _mean_module = register_module(
      "mean", torch::nn::Linear(torch::nn::LinearOptions(_num_inputs, _num_outputs).bias(true)));
  _mean_module->to(_device_type, _data_type);
 
  if (_state_independent_std)
    _log_std_parameter = register_parameter(
        "log_std",
        torch::zeros({_num_outputs},
                     torch::TensorOptions().dtype(_data_type).device(_device_type)));
  else
  {
    _std_module = register_module(
        "std", torch::nn::Linear(torch::nn::LinearOptions(_num_inputs, _num_outputs).bias(true)));
    _std_module->to(_device_type, _data_type);
  }
}

Referenced by LibtorchGaussianActionDistribution().

◆ deterministicAction()

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

overridevirtual

Return the Gaussian mean action used for deterministic evaluation.

Implements Moose::LibtorchActionDistribution.

Definition at line 182 of file LibtorchActionDistribution.C.

{
  return _mean * actionScaleTensor();
}

◆ entropy()

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

overridevirtual

Compute the entropy of the current Gaussian distribution.

Implements Moose::LibtorchActionDistribution.

Definition at line 201 of file LibtorchActionDistribution.C.

{
  const auto log_action_scale = torch::log(torch::abs(actionScaleTensor()));
  constexpr Real pi = 3.14159265358979323846;
  return 0.5 * std::log(2.0 * pi) + _log_std_tensor + 0.5 + log_action_scale;
}

◆ initialize()

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

overridevirtual

Initialize the trainable Gaussian distribution parameters.

Parameters

generator Optional random-number generator used for reproducible initialization.

Implements Moose::LibtorchActionDistribution.

Definition at line 140 of file LibtorchActionDistribution.C.

{
  Moose::orthogonalInitializeTensor(_mean_module->weight, 1.0, generator);
  torch::nn::init::zeros_(_mean_module->bias);
 
  if (_state_independent_std)
  {
    torch::nn::init::zeros_(_log_std_parameter);
    return;
  }
 
  Moose::orthogonalInitializeTensor(_std_module->weight, 1.0, generator);
  torch::nn::init::zeros_(_std_module->bias);
}

◆ isBounded()

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

inlineoverridevirtual

Return false because the Gaussian distribution is unbounded.

Implements Moose::LibtorchActionDistribution.

Definition at line 187 of file LibtorchActionDistribution.h.

187{ return false; }

◆ logProbability()

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

overridevirtual

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

Parameters

action Action tensor in physical units.

Implements Moose::LibtorchActionDistribution.

Definition at line 188 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();
 
  constexpr Real pi = 3.14159265358979323846;
  const torch::Tensor var = _std_tensor * _std_tensor;
  return -((unscaled_action - _mean) * (unscaled_action - _mean)) / (2.0 * var) - _log_std_tensor -
         0.5 * std::log(2.0 * pi) - log_action_scale;
}

◆ logStd() [1/2]

torch::Tensor & Moose::LibtorchGaussianActionDistribution::logStd ( )

inline

Return the state-independent Gaussian log-standard-deviation parameter.

Definition at line 217 of file LibtorchActionDistribution.h.

217{ return _log_std_parameter; }

◆ logStd() [2/2]

const torch::Tensor & Moose::LibtorchGaussianActionDistribution::logStd ( ) const

inline

Return the state-independent Gaussian log-standard-deviation parameter.

Definition at line 222 of file LibtorchActionDistribution.h.

222{ return _log_std_parameter; }

◆ meanModule() [1/2]

torch::nn::Linear & Moose::LibtorchGaussianActionDistribution::meanModule ( )

inline

Return the Gaussian mean head.

Definition at line 197 of file LibtorchActionDistribution.h.

197{ return _mean_module; }

◆ meanModule() [2/2]

const torch::nn::Linear & Moose::LibtorchGaussianActionDistribution::meanModule ( ) const

inline

Return the Gaussian mean head.

Definition at line 202 of file LibtorchActionDistribution.h.

202{ return _mean_module; }

◆ 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 logProbability(), and Moose::LibtorchBetaActionDistribution::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 reset(), and Moose::LibtorchBetaActionDistribution::reset().

◆ reset()

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

overridevirtual

Refresh the Gaussian distribution parameters from the latest actor features.

Parameters

input Feature tensor coming from the actor body.

Implements Moose::LibtorchActionDistribution.

Definition at line 156 of file LibtorchActionDistribution.C.

{
  const auto features = prepareFeatures(input);
  _mean = _mean_module->forward(features);
 
  if (_state_independent_std)
  {
    if (_mean.dim() <= 1)
      _log_std_tensor = _log_std_parameter;
    else
      _log_std_tensor = _log_std_parameter.view({1, -1}).expand(_mean.sizes());
  }
  else
    _log_std_tensor = _std_module->forward(features);
 
  _log_std_tensor = torch::clamp(_log_std_tensor, std::log(1e-12), -std::log(1e-12));
  _std_tensor = torch::exp(_log_std_tensor);
}

◆ sample()

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

overridevirtual

Draw a stochastic Gaussian action sample in physical units.

Parameters

generator Optional random-number generator used for sampling.

Implements Moose::LibtorchActionDistribution.

Definition at line 176 of file LibtorchActionDistribution.C.

{
  return at::normal(_mean, _std_tensor, generator) * actionScaleTensor();
}

◆ stateIndependentStd()

bool Moose::LibtorchGaussianActionDistribution::stateIndependentStd ( ) const

inline

Return whether the Gaussian standard deviation ignores the current actor features.

Definition at line 192 of file LibtorchActionDistribution.h.

192{ return _state_independent_std; }

◆ stdModule() [1/2]

torch::nn::Linear & Moose::LibtorchGaussianActionDistribution::stdModule ( )

inline

Return the state-dependent Gaussian standard-deviation head.

Definition at line 207 of file LibtorchActionDistribution.h.

207{ return _std_module; }

◆ stdModule() [2/2]

const torch::nn::Linear & Moose::LibtorchGaussianActionDistribution::stdModule ( ) const

inline

Return the state-dependent Gaussian standard-deviation head.

Definition at line 212 of file LibtorchActionDistribution.h.

212{ return _std_module; }

◆ stdTensor()

const torch::Tensor & Moose::LibtorchGaussianActionDistribution::stdTensor ( ) const

inline

Return the cached Gaussian standard-deviation tensor.

Definition at line 227 of file LibtorchActionDistribution.h.

227{ return _std_tensor; }

◆ 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);
}