LibtorchDRLControl.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://mooseframework.inl.gov
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html
 
#ifdef MOOSE_LIBTORCH_ENABLED
 
#pragma once
 
#include <cstdint>
 
#include "LibtorchActorNeuralNet.h"
#include "LibtorchNeuralNetControl.h"
 
class LibtorchDRLControl : public LibtorchNeuralNetControl
{
public:
  static InputParameters validParams();
 
  LibtorchDRLControl(const InputParameters & parameters);
 
  virtual void initialSetup() override;
 
  virtual void execute() override;
 
  Real getSignalLogProbability(const unsigned int signal_index) const;
 
  void loadControlNeuralNet(const Moose::LibtorchActorNeuralNet & input_nn);
 
  virtual void loadControlNeuralNet(const Moose::LibtorchArtificialNeuralNet & input_nn) override;
 
  virtual void loadControlNeuralNetFromFile() override;
 
  void setPolicySampleSeed(uint64_t seed);
 
  unsigned int numStepsInPeriod() const { return _num_steps_in_period; }
 
  unsigned int inputTimesteps() const { return _input_timesteps; }
 
  unsigned int numberOfObservations() const { return _observation_names.size(); }
 
protected:
  void applyControlSignals();
 
  Real computeControlOffset(unsigned int control_i) const;
 
  std::vector<Real> & _current_control_signal_log_probabilities;
 
  std::vector<Real> & _previous_control_signal;
  std::vector<Real> & _current_smoothed_signal;
 
  std::vector<Real> _control_offsets;
  bool _has_control_offsets;
 
  std::shared_ptr<Moose::LibtorchActorNeuralNet> _actor_nn;
  at::Generator _policy_generator;
  std::vector<std::uint8_t> & _policy_generator_state;
 
  const unsigned int _num_steps_in_period;
  const Real _control_smoothing_factor;
  const bool _stochastic;
 
private:
  bool shouldEvaluatePolicy() const;
 
  void restorePolicyGeneratorState();
 
  void savePolicyGeneratorState();
};
 
#endif