Feed-forward actor network coupled to a Gaussian or Beta action distribution. More...

#include <LibtorchActorNeuralNet.h>

Inheritance diagram for Moose::LibtorchActorNeuralNet:

Public Member Functions
	LibtorchActorNeuralNet (const std::string name, const unsigned int num_inputs, const unsigned int num_outputs, const std::vector< unsigned int > &num_neurons_per_layer, const std::vector< std::string > &activation_function={"relu"}, const std::vector< Real > &minimum_values={}, const std::vector< Real > &maximum_values={}, const torch::DeviceType device_type=torch::kCPU, const torch::ScalarType scalar_type=torch::kDouble, const bool build_on_construct=true, const std::vector< Real > &input_shift_factors={}, const std::vector< Real > &input_scaling_factors={}, const std::vector< Real > &output_scaling_factors={}, const bool state_independent_std=true)
	Build an actor network with either a Gaussian or Beta action distribution.

	LibtorchActorNeuralNet (const Moose::LibtorchActorNeuralNet &nn, const bool build_on_construct=true)
	Copy-construct an actor network.

virtual torch::Tensor	forward (const torch::Tensor &x) override
	Run a forward pass through the actor body.

virtual torch::Tensor	evaluate (torch::Tensor &input, bool sampled, c10::optional< at::Generator > generator=c10::nullopt)
	Evaluate the actor and either sample from it or use its deterministic action.

virtual torch::Tensor	sample (c10::optional< at::Generator > generator=c10::nullopt)
	Sample an action from the already-reset distribution.

virtual void	constructNeuralNetwork () override
	Construct the actor body and action-distribution modules.

const LibtorchActionDistribution &	actionDistribution () const
	Return the active action distribution as the common base type.

LibtorchActionDistribution &	actionDistribution ()
	Return the active action distribution as the common base type.

const LibtorchGaussianActionDistribution *	gaussianActionDistributionPtr () const
	Return the Gaussian action distribution pointer, or nullptr for bounded actors.

LibtorchGaussianActionDistribution *	gaussianActionDistributionPtr ()
	Return the Gaussian action distribution pointer, or nullptr for bounded actors.

const LibtorchGaussianActionDistribution &	gaussianActionDistribution () const
	Return the Gaussian action distribution reference.

LibtorchGaussianActionDistribution &	gaussianActionDistribution ()
	Return the Gaussian action distribution reference.

const LibtorchBetaActionDistribution *	betaActionDistributionPtr () const
	Return the Beta action distribution pointer, or nullptr for Gaussian actors.

LibtorchBetaActionDistribution *	betaActionDistributionPtr ()
	Return the Beta action distribution pointer, or nullptr for Gaussian actors.

const LibtorchBetaActionDistribution &	betaActionDistribution () const
	Return the Beta action distribution reference.

LibtorchBetaActionDistribution &	betaActionDistribution ()
	Return the Beta action distribution reference.

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

const std::vector< Real > &	minValues () const
	Return the configured lower action bounds.

const std::vector< Real > &	maxValues () const
	Return the configured upper action bounds.

void	resetDistributionParams (torch::Tensor input)
	Refresh the cached distribution parameters from a fresh input tensor.

torch::Tensor	logProbability (const torch::Tensor &other)
	Evaluate the log-probability of an action under the current actor state.

torch::Tensor	entropy ()
	Compute the entropy of the current action distribution.

virtual void	initializeNeuralNetwork (c10::optional< at::Generator > generator=c10::nullopt) override
	Initialize the actor body and action-distribution parameters.

virtual void	addLayer (const std::string &layer_name, const std::unordered_map< std::string, unsigned int > &parameters)

const std::string &	name () const

unsigned int	numInputs () const

unsigned int	numOutputs () const

unsigned int	numHiddenLayers () const

const std::vector< unsigned int > &	numNeuronsPerLayer () const

const MultiMooseEnum &	activationFunctions () const

torch::DeviceType	deviceType () const

torch::ScalarType	dataType () const

const std::vector< Real > &	inputShiftFactors () const

const std::vector< Real > &	inputScalingFactors () const

const std::vector< Real > &	outputScalingFactors () const

void	synchronizeAffineFactorsFromBuffers ()

Real	determineGain (const std::string &activation)

void	store (nlohmann::json &json) const

Protected Member Functions
void	initializeAffineBuffers ()

virtual torch::Tensor	preprocessInput (const torch::Tensor &x) const

virtual torch::Tensor	scaleOutput (const torch::Tensor &y) const

Static Protected Member Functions
static std::vector< Real >	setAffineFactors (const std::vector< Real > &factors, unsigned int expected_size, Real default_value, const std::string &factor_name)

Protected Attributes
const std::vector< Real >	_minimum_values
	Lower action bounds used by bounded actor distributions.

const std::vector< Real >	_maximum_values
	Upper action bounds used by bounded actor distributions.

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

std::shared_ptr< LibtorchActionDistribution >	_action_distribution
	Action-distribution module attached to the actor output.

const std::string	_name

std::vector< torch::nn::Linear >	_weights

const unsigned int	_num_inputs

const unsigned int	_num_outputs

const std::vector< unsigned int >	_num_neurons_per_layer

MultiMooseEnum	_activation_function

const torch::DeviceType	_device_type

const torch::ScalarType	_data_type

std::vector< Real >	_input_shift_factors

std::vector< Real >	_input_scaling_factors

std::vector< Real >	_output_scaling_factors

torch::Tensor	_input_shift_tensor

torch::Tensor	_input_scale_tensor

torch::Tensor	_output_scale_tensor

Detailed Description

Feed-forward actor network coupled to a Gaussian or Beta action distribution.

Definition at line 24 of file LibtorchActorNeuralNet.h.

Constructor & Destructor Documentation

◆ LibtorchActorNeuralNet() [1/2]

Moose::LibtorchActorNeuralNet::LibtorchActorNeuralNet	(	const std::string	name,
		const unsigned int	num_inputs,
		const unsigned int	num_outputs,
		const std::vector< unsigned int > &	num_neurons_per_layer,
		const std::vector< std::string > &	activation_function = `{"relu"}`,
		const std::vector< Real > &	minimum_values = `{}`,
		const std::vector< Real > &	maximum_values = `{}`,
		const torch::DeviceType	device_type = `torch::kCPU`,
		const torch::ScalarType	scalar_type = `torch::kDouble`,
		const bool	build_on_construct = `true`,
		const std::vector< Real > &	input_shift_factors = `{}`,
		const std::vector< Real > &	input_scaling_factors = `{}`,
		const std::vector< Real > &	output_scaling_factors = `{}`,
		const bool	state_independent_std = `true`
	)

Build an actor network with either a Gaussian or Beta action distribution.

Parameters

name	Name of the neural network module.
num_inputs	Number of actor inputs.
num_outputs	Number of action outputs.
num_neurons_per_layer	Hidden-layer widths.
activation_function	Hidden-layer activation names.
minimum_values	Optional lower action bounds. Leave empty for Gaussian actions.
maximum_values	Optional upper action bounds. Leave empty for Gaussian actions.
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.
input_shift_factors	Optional affine input shifts.
input_scaling_factors	Optional affine input scales.
output_scaling_factors	Optional per-action scaling in physical units.
state_independent_std	Whether the Gaussian std should ignore the current state.

Definition at line 64 of file LibtorchActorNeuralNet.C.

  : LibtorchArtificialNeuralNet(name,
                                num_inputs,
                                num_outputs,
                                num_neurons_per_layer,
                                activation_function,
                                device_type,
                                data_type,
                                false,
                                input_shift_factors,
                                input_scaling_factors,
                                output_scaling_factors),
    _minimum_values(minimum_values),
    _maximum_values(maximum_values),
    _state_independent_std(state_independent_std)
{
  const bool has_minimum_values = !_minimum_values.empty();
  const bool has_maximum_values = !_maximum_values.empty();
  if (has_minimum_values != has_maximum_values)
    mooseError("Bounded action distributions require both minimum_values and maximum_values.");
 
  if (has_minimum_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.");
  }
 
  if (build_on_construct)
    constructNeuralNetwork();
}

◆ LibtorchActorNeuralNet() [2/2]

Moose::LibtorchActorNeuralNet::LibtorchActorNeuralNet	(	const Moose::LibtorchActorNeuralNet &	nn,
		const bool	build_on_construct = `true`
	)

Copy-construct an actor network.

Parameters

nn	Actor network to copy.
build_on_construct	Whether to rebuild the module structure during the copy.

Definition at line 114 of file LibtorchActorNeuralNet.C.

  : LibtorchArtificialNeuralNet(dynamic_cast<const LibtorchArtificialNeuralNet &>(nn), false),
    _minimum_values(nn.minValues()),
    _maximum_values(nn.maxValues()),
    _state_independent_std(nn.stateIndependentStd())
{
  // We construct the NN architecture
  if (build_on_construct)
  {
    constructNeuralNetwork();
    // We fill it up with the current parameter values
    const auto & from_params = nn.named_parameters();
    auto to_params = this->named_parameters();
    for (unsigned int param_i : make_range(from_params.size()))
      to_params[param_i].value().data() = from_params[param_i].value().data().clone();
 
    const auto & from_buffers = nn.named_buffers();
    auto to_buffers = this->named_buffers();
    for (unsigned int buffer_i : make_range(from_buffers.size()))
      to_buffers[buffer_i].value().data() = from_buffers[buffer_i].value().data().clone();
  }
}

Member Function Documentation

◆ actionDistribution() [1/2]

LibtorchActionDistribution & Moose::LibtorchActorNeuralNet::actionDistribution ( )

inline

Return the active action distribution as the common base type.

Definition at line 102 of file LibtorchActorNeuralNet.h.

102{ return *_action_distribution; }

Moose::LibtorchActorNeuralNet::_action_distribution

std::shared_ptr< LibtorchActionDistribution > _action_distribution

Action-distribution module attached to the actor output.

Definition LibtorchActorNeuralNet.h:191

◆ actionDistribution() [2/2]

const LibtorchActionDistribution & Moose::LibtorchActorNeuralNet::actionDistribution ( ) const

inline

Return the active action distribution as the common base type.

Definition at line 97 of file LibtorchActorNeuralNet.h.

97{ return *_action_distribution; }

Referenced by Moose::loadLibtorchActorNeuralNetState().

◆ betaActionDistribution() [1/2]

LibtorchBetaActionDistribution & Moose::LibtorchActorNeuralNet::betaActionDistribution ( )

Return the Beta action distribution reference.

Errors if the actor is unbounded.

Definition at line 253 of file LibtorchActorNeuralNet.C.

{
  auto * distribution = betaActionDistributionPtr();
  if (!distribution)
    mooseError("Requested a Beta action distribution from an unbounded actor.");
  return *distribution;
}

◆ betaActionDistribution() [2/2]

const LibtorchBetaActionDistribution & Moose::LibtorchActorNeuralNet::betaActionDistribution ( ) const

Return the Beta action distribution reference.

Errors if the actor is unbounded.

Definition at line 244 of file LibtorchActorNeuralNet.C.

{
  const auto * distribution = betaActionDistributionPtr();
  if (!distribution)
    mooseError("Requested a Beta action distribution from an unbounded actor.");
  return *distribution;
}

◆ betaActionDistributionPtr() [1/2]

LibtorchBetaActionDistribution * Moose::LibtorchActorNeuralNet::betaActionDistributionPtr ( )

Return the Beta action distribution pointer, or nullptr for Gaussian actors.

Definition at line 238 of file LibtorchActorNeuralNet.C.

{
  return dynamic_cast<LibtorchBetaActionDistribution *>(_action_distribution.get());
}

◆ betaActionDistributionPtr() [2/2]

const LibtorchBetaActionDistribution * Moose::LibtorchActorNeuralNet::betaActionDistributionPtr ( ) const

Return the Beta action distribution pointer, or nullptr for Gaussian actors.

Definition at line 232 of file LibtorchActorNeuralNet.C.

{
  return dynamic_cast<const LibtorchBetaActionDistribution *>(_action_distribution.get());
}

Referenced by betaActionDistribution(), and betaActionDistribution().

◆ constructNeuralNetwork()

void Moose::LibtorchActorNeuralNet::constructNeuralNetwork ( )

overridevirtual

Construct the actor body and action-distribution modules.

Reimplemented from Moose::LibtorchArtificialNeuralNet.

Definition at line 154 of file LibtorchActorNeuralNet.C.

{
  // Adding hidden layers
  unsigned int inp_neurons = _num_inputs;
  for (unsigned int i = 0; i < numHiddenLayers(); ++i)
  {
    std::unordered_map<std::string, unsigned int> parameters = {
        {"inp_neurons", inp_neurons}, {"out_neurons", _num_neurons_per_layer[i]}};
    addLayer("hidden_layer_" + std::to_string(i + 1), parameters);
 
    // Necessary to retain double precision (and error-free runs)
    _weights[i]->to(_device_type, _data_type);
    inp_neurons = _num_neurons_per_layer[i];
  }
 
  if (_minimum_values.empty() && _maximum_values.empty())
    _action_distribution =
        std::make_shared<LibtorchGaussianActionDistribution>("action_distribution",
                                                             inp_neurons,
                                                             _num_outputs,
                                                             _device_type,
                                                             _data_type,
                                                             true,
                                                             _output_scaling_factors,
                                                             _state_independent_std);
  else
    _action_distribution =
        std::make_shared<LibtorchBetaActionDistribution>("action_distribution",
                                                         inp_neurons,
                                                         _num_outputs,
                                                         _minimum_values,
                                                         _maximum_values,
                                                         _device_type,
                                                         _data_type,
                                                         true,
                                                         _output_scaling_factors);
 
  // Keep the serialized module name stable so existing checkpoints continue to load.
  register_module("action_head", _action_distribution);
}

Referenced by LibtorchActorNeuralNet(), and LibtorchActorNeuralNet().

◆ entropy()

torch::Tensor Moose::LibtorchActorNeuralNet::entropy ( )

Compute the entropy of the current action distribution.

Definition at line 196 of file LibtorchActorNeuralNet.C.

{
  return _action_distribution->entropy();
}

Referenced by LibtorchRLPPOLoss::compute().

◆ evaluate()

torch::Tensor Moose::LibtorchActorNeuralNet::evaluate	(	torch::Tensor &	input,
		bool	sampled,
		c10::optional< at::Generator >	generator = `c10::nullopt`
	)

virtual

Evaluate the actor and either sample from it or use its deterministic action.

Parameters

input	Input tensor for the evaluation.
sampled	Whether to draw a stochastic sample.
generator	Optional random-number generator used for sampling.

Definition at line 294 of file LibtorchActorNeuralNet.C.

{
  torch::Tensor output = forward(x);
 
  resetDistributionParams(output);
 
  if (sampled)
    return sample(generator);
 
  return _action_distribution->deterministicAction();
}

Referenced by LibtorchRLPPOLoss::compute().

◆ forward()

torch::Tensor Moose::LibtorchActorNeuralNet::forward ( const torch::Tensor & x )

overridevirtual

Run a forward pass through the actor body.

Parameters

x	Input tensor for the evaluation.

Reimplemented from Moose::LibtorchArtificialNeuralNet.

Definition at line 268 of file LibtorchActorNeuralNet.C.

{
  torch::Tensor output = preprocessInput(x);
 
  for (unsigned int i = 0; i < _weights.size(); ++i)
  {
    std::string activation =
        _activation_function.size() > 1 ? _activation_function[i] : _activation_function[0];
    if (activation == "relu")
      output = torch::relu(_weights[i]->forward(output));
    else if (activation == "sigmoid")
      output = torch::sigmoid(_weights[i]->forward(output));
    else if (activation == "tanh")
      output = torch::tanh(_weights[i]->forward(output));
    else if (activation == "elu")
      output = torch::elu(_weights[i]->forward(output));
    else if (activation == "gelu")
      output = torch::gelu(_weights[i]->forward(output));
    else if (activation == "linear")
      output = _weights[i]->forward(output);
  }
 
  return output;
}

Referenced by evaluate(), and forward().

◆ gaussianActionDistribution() [1/2]

LibtorchGaussianActionDistribution & Moose::LibtorchActorNeuralNet::gaussianActionDistribution ( )

Return the Gaussian action distribution reference.

Errors if the actor is bounded.

Definition at line 223 of file LibtorchActorNeuralNet.C.

{
  auto * distribution = gaussianActionDistributionPtr();
  if (!distribution)
    mooseError("Requested a Gaussian action distribution from a bounded actor.");
  return *distribution;
}

◆ gaussianActionDistribution() [2/2]

const LibtorchGaussianActionDistribution & Moose::LibtorchActorNeuralNet::gaussianActionDistribution ( ) const

Return the Gaussian action distribution reference.

Errors if the actor is bounded.

Definition at line 214 of file LibtorchActorNeuralNet.C.

{
  const auto * distribution = gaussianActionDistributionPtr();
  if (!distribution)
    mooseError("Requested a Gaussian action distribution from a bounded actor.");
  return *distribution;
}

◆ gaussianActionDistributionPtr() [1/2]

LibtorchGaussianActionDistribution * Moose::LibtorchActorNeuralNet::gaussianActionDistributionPtr ( )

Return the Gaussian action distribution pointer, or nullptr for bounded actors.

Definition at line 208 of file LibtorchActorNeuralNet.C.

{
  return dynamic_cast<LibtorchGaussianActionDistribution *>(_action_distribution.get());
}

◆ gaussianActionDistributionPtr() [2/2]

const LibtorchGaussianActionDistribution * Moose::LibtorchActorNeuralNet::gaussianActionDistributionPtr ( ) const

Return the Gaussian action distribution pointer, or nullptr for bounded actors.

Definition at line 202 of file LibtorchActorNeuralNet.C.

{
  return dynamic_cast<const LibtorchGaussianActionDistribution *>(_action_distribution.get());
}

Referenced by gaussianActionDistribution(), and gaussianActionDistribution().

◆ initializeNeuralNetwork()

void Moose::LibtorchActorNeuralNet::initializeNeuralNetwork ( c10::optional< at::Generator > generator = c10::nullopt )

overridevirtual

Initialize the actor body and action-distribution parameters.

Parameters

generator Optional random-number generator used for reproducible initialization.

Reimplemented from Moose::LibtorchArtificialNeuralNet.

Definition at line 139 of file LibtorchActorNeuralNet.C.

{
  for (unsigned int i = 0; i < numHiddenLayers(); ++i)
  {
    const auto & activation =
        _activation_function.size() > 1 ? _activation_function[i] : _activation_function[0];
    const Real gain = determineGain(activation);
    Moose::orthogonalInitializeTensor(_weights[i]->weight, gain, generator);
    torch::nn::init::zeros_(_weights[i]->bias);
  }
 
  _action_distribution->initialize(generator);
}

◆ logProbability()

torch::Tensor Moose::LibtorchActorNeuralNet::logProbability ( const torch::Tensor & other )

Evaluate the log-probability of an action under the current actor state.

Parameters

other Action tensor in physical units.

Definition at line 315 of file LibtorchActorNeuralNet.C.

{
  return _action_distribution->logProbability(action);
}

Referenced by LibtorchRLPPOLoss::compute().

◆ maxValues()

const std::vector< Real > & Moose::LibtorchActorNeuralNet::maxValues ( ) const

inline

Return the configured upper action bounds.

Definition at line 157 of file LibtorchActorNeuralNet.h.

157{ return _maximum_values; }

◆ minValues()

const std::vector< Real > & Moose::LibtorchActorNeuralNet::minValues ( ) const

inline

Return the configured lower action bounds.

Definition at line 152 of file LibtorchActorNeuralNet.h.

152{ return _minimum_values; }

◆ resetDistributionParams()

void Moose::LibtorchActorNeuralNet::resetDistributionParams ( torch::Tensor input )

Refresh the cached distribution parameters from a fresh input tensor.

Parameters

input Input tensor used to update the distribution.

Definition at line 262 of file LibtorchActorNeuralNet.C.

{
  _action_distribution->reset(input);
}

Referenced by evaluate().

◆ sample()

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

virtual

Sample an action from the already-reset distribution.

Parameters

generator Optional random-number generator used for sampling.

Definition at line 309 of file LibtorchActorNeuralNet.C.

{
  return _action_distribution->sample(generator);
}

Referenced by evaluate().

◆ stateIndependentStd()

bool Moose::LibtorchActorNeuralNet::stateIndependentStd ( ) const

inline

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

Definition at line 147 of file LibtorchActorNeuralNet.h.

147{ return _state_independent_std; }

Member Data Documentation

◆ _action_distribution

std::shared_ptr<LibtorchActionDistribution> Moose::LibtorchActorNeuralNet::_action_distribution

protected

Action-distribution module attached to the actor output.

Definition at line 191 of file LibtorchActorNeuralNet.h.

Referenced by actionDistribution(), actionDistribution(), betaActionDistributionPtr(), betaActionDistributionPtr(), constructNeuralNetwork(), entropy(), evaluate(), gaussianActionDistributionPtr(), gaussianActionDistributionPtr(), initializeNeuralNetwork(), logProbability(), resetDistributionParams(), and sample().

◆ _maximum_values

const std::vector<Real> Moose::LibtorchActorNeuralNet::_maximum_values

protected

Upper action bounds used by bounded actor distributions.

Definition at line 187 of file LibtorchActorNeuralNet.h.

Referenced by constructNeuralNetwork(), LibtorchActorNeuralNet(), and maxValues().

◆ _minimum_values

const std::vector<Real> Moose::LibtorchActorNeuralNet::_minimum_values

protected

Lower action bounds used by bounded actor distributions.

Definition at line 185 of file LibtorchActorNeuralNet.h.

Referenced by constructNeuralNetwork(), LibtorchActorNeuralNet(), and minValues().

◆ _state_independent_std

const bool Moose::LibtorchActorNeuralNet::_state_independent_std

protected

Whether the Gaussian std ignores the current actor features.

Definition at line 189 of file LibtorchActorNeuralNet.h.

Referenced by constructNeuralNetwork(), and stateIndependentStd().

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

stochastic_tools/include/libtorch/utils/LibtorchActorNeuralNet.h
stochastic_tools/src/libtorch/utils/LibtorchActorNeuralNet.C

Public Member Functions

Protected Member Functions

Static Protected Member Functions

Protected Attributes

Detailed Description

Constructor & Destructor Documentation

◆ LibtorchActorNeuralNet() [1/2]

◆ LibtorchActorNeuralNet() [2/2]

Member Function Documentation

◆ actionDistribution() [1/2]

◆ actionDistribution() [2/2]

◆ betaActionDistribution() [1/2]

◆ betaActionDistribution() [2/2]

◆ betaActionDistributionPtr() [1/2]

◆ betaActionDistributionPtr() [2/2]

◆ constructNeuralNetwork()

◆ entropy()

◆ evaluate()

◆ forward()

◆ gaussianActionDistribution() [1/2]

◆ gaussianActionDistribution() [2/2]

◆ gaussianActionDistributionPtr() [1/2]

◆ gaussianActionDistributionPtr() [2/2]

◆ initializeNeuralNetwork()

◆ logProbability()

◆ maxValues()

◆ minValues()

◆ resetDistributionParams()

◆ sample()

◆ stateIndependentStd()

Member Data Documentation

◆ _action_distribution

◆ _maximum_values

◆ _minimum_values

◆ _state_independent_std