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 "LibtorchArtificialNeuralNet.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; }
 
  void trainController();
 
  const Moose::LibtorchArtificialNeuralNet & controlNeuralNet() const { return *_control_nn; }
 
protected:
  void computeAverageEpisodeReward();
 
  void convertDataToTensor(std::vector<std::vector<Real>> & vector_data,
                           torch::Tensor & tensor_data,
                           const bool detach = false);
 
  torch::Tensor evaluateValue(torch::Tensor & input);
 
  torch::Tensor evaluateAction(torch::Tensor & input, torch::Tensor & output);
 
  void computeRewardToGo();
 
  void resetData();
 
  const std::vector<ReporterName> _response_names;
 
  std::vector<const std::vector<Real> *> _response_value_pointers;
 
  const std::vector<Real> _response_shift_factors;
 
  const std::vector<Real> _response_scaling_factors;
 
  const std::vector<ReporterName> _control_names;
 
  std::vector<const std::vector<Real> *> _control_value_pointers;
 
  const std::vector<ReporterName> _log_probability_names;
 
  std::vector<const std::vector<Real> *> _log_probability_value_pointers;
 
  const ReporterName _reward_name;
 
  const std::vector<Real> * _reward_value_pointer;
 
  const unsigned int _input_timesteps;
 
  unsigned int _num_inputs;
  unsigned int _num_outputs;
 
  std::vector<std::vector<Real>> _input_data;
  std::vector<std::vector<Real>> _output_data;
  std::vector<std::vector<Real>> _log_probability_data;
 
  std::vector<Real> _reward_data;
  std::vector<Real> _return_data;
 
  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 std::vector<Real> _action_std;
 
  const std::string _filename_base;
 
  const bool _read_from_file;
 
  const bool _shift_outputs;
 
  Real _average_episode_reward;
 
  const bool _standardize_advantage;
 
  const unsigned int _loss_print_frequency;
 
  std::shared_ptr<Moose::LibtorchArtificialNeuralNet> _control_nn;
  std::shared_ptr<Moose::LibtorchArtificialNeuralNet> _critic_nn;
 
  torch::Tensor _std;
 
  torch::Tensor _input_tensor;
  torch::Tensor _output_tensor;
  torch::Tensor _return_tensor;
  torch::Tensor _log_probability_tensor;
 
private:
  void getInputDataFromReporter(std::vector<std::vector<Real>> & data,
                                const std::vector<const std::vector<Real> *> & reporter_links,
                                const unsigned int num_timesteps);
  void getOutputDataFromReporter(std::vector<std::vector<Real>> & data,
                                 const std::vector<const std::vector<Real> *> & reporter_links);
 
  void getRewardDataFromReporter(std::vector<Real> & data,
                                 const std::vector<Real> * const reporter_link);
 
  void getReporterPointers(const std::vector<ReporterName> & reporter_names,
                           std::vector<const std::vector<Real> *> & pointer_storage);
 
  unsigned int _update_counter;
};
 
#endif