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 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