LibtorchDRLControlTrainer.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 <torch/torch.h>
#include "LibtorchActorNeuralNet.h"
#include "LibtorchObservationHistoryHelper.h"
#include "LibtorchRLMiniBatchSampler.h"
#include "LibtorchRLPPOLoss.h"
#include "LibtorchRLTrajectoryAssembler.h"
#include "LibtorchRLTrajectoryBuffer.h"
#include "LibtorchRLValueEstimator.h"
 
#include "libmesh/utility.h"
#include "SurrogateTrainer.h"
 
class LibtorchDRLControlTrainer : public SurrogateTrainerBase
{
public:
  static InputParameters validParams();
 
  LibtorchDRLControlTrainer(const InputParameters & parameters);
 
  virtual void execute() override;
 
  Real averageEpisodeReward() { return _average_episode_reward; }
  Real stdEpisodeReward() { return _std_episode_reward; }
 
  std::vector<Real> sampleAverageEpsiodeRewards() { return _sample_average_episode_reward; }
 
  std::vector<Real> sampleStdEpsiodeRewards() { return _sample_std_episode_reward; }
 
  void trainController(const LibtorchRLTrajectoryBuffer::TensorBatch & batch);
 
  const Moose::LibtorchActorNeuralNet & controlNeuralNet() const { return *_control_nn; }
 
  unsigned int seed() const { return _seed; }
 
  unsigned int timestepWindow() const { return getParam<unsigned int>("timestep_window"); }
 
  unsigned int inputTimesteps() const { return _input_timesteps; }
 
  unsigned int numberOfObservations() const { return _state_names.size(); }
 
  unsigned int numberOfControlSignals() const { return _action_names.size(); }
 
protected:
  void computeEpisodeRewardStatistics();
 
  void resetData();
 
  const std::vector<ReporterName> _state_names;
 
  std::vector<const std::vector<std::vector<Real>> *> _state_value_pointers;
 
  const std::vector<Real> _state_shift_factors;
 
  const std::vector<Real> _state_scaling_factors;
 
  const std::vector<ReporterName> _action_names;
 
  const std::vector<Real> _action_scaling_factors;
 
  std::vector<const std::vector<std::vector<Real>> *> _action_value_pointers;
 
  const std::vector<ReporterName> _log_probability_names;
 
  std::vector<const std::vector<std::vector<Real>> *> _log_probability_value_pointers;
 
  const ReporterName _reward_name;
 
  const std::vector<std::vector<Real>> * _reward_value_pointer;
 
  const unsigned int _input_timesteps;
 
  unsigned int _num_inputs;
  unsigned int _num_outputs;
 
  const unsigned int _num_epochs;
 
  const std::vector<unsigned int> _num_critic_neurons_per_layer;
 
  const Real _critic_learning_rate;
 
  const std::vector<unsigned int> _num_control_neurons_per_layer;
 
  const Real _control_learning_rate;
 
  const unsigned int _update_frequency;
 
  const Real _clip_param;
 
  const Real _decay_factor;
  const Real _lambda_factor;
 
  const std::string _filename_base;
 
  const bool _read_from_file;
 
  const bool _shift_outputs;
  const bool _average_reward_over_timestep_window;
 
  Real _average_episode_reward;
  Real _std_episode_reward;
 
  std::vector<Real> _sample_average_episode_reward;
  std::vector<Real> _sample_std_episode_reward;
 
  const bool _standardize_advantage;
 
  const unsigned int _loss_print_frequency;
 
  const unsigned int _seed;
 
  std::vector<Real> _min_values;
  std::vector<Real> _max_values;
 
  std::shared_ptr<Moose::LibtorchActorNeuralNet> _control_nn;
  std::shared_ptr<Moose::LibtorchArtificialNeuralNet> _critic_nn;
 
  Real _highest_reward;
  Real _entropy_coeff;
 
  std::unique_ptr<torch::optim::Adam> _actor_optimizer;
  std::unique_ptr<torch::optim::Adam> _critic_optimizer;
 
private:
  void getReporterPointers(const std::vector<ReporterName> & reporter_names,
                           std::vector<const std::vector<std::vector<Real>> *> & pointer_storage);
 
  bool collectTrajectoriesFromReporters();
 
  unsigned int _update_counter;
 
  const LibtorchRLTrajectoryAssembler _trajectory_assembler;
  LibtorchRLTrajectoryBuffer _trajectory_buffer;
  const LibtorchRLMiniBatchSampler _sampler;
  const LibtorchRLValueEstimator _value_estimator;
  const LibtorchRLPPOLoss _ppo_loss;
};
 
#endif