LibtorchRLTrajectoryAssembler.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 "LibtorchObservationHistoryHelper.h"
#include "LibtorchRLTrajectoryBuffer.h"
 
#include <string>
#include <vector>
 
class LibtorchRLTrajectoryAssembler
{
public:
  using ReporterData = std::vector<const std::vector<std::vector<Real>> *>;
 
  LibtorchRLTrajectoryAssembler(unsigned int input_timesteps,
                                unsigned int timestep_window,
                                bool shift_outputs,
                                bool average_reward_over_timestep_window);
 
  std::vector<LibtorchRLTrajectoryBuffer::Trajectory>
  assemble(const ReporterData & observations,
           const ReporterData & actions,
           const ReporterData & log_probabilities,
           const std::vector<std::vector<Real>> & rewards) const;
 
  unsigned int numTransitions(std::size_t raw_sequence_size) const;
 
private:
  static void checkSampleCount(const ReporterData & data,
                               const std::string & quantity,
                               std::size_t num_samples);
 
  std::vector<Real>
  downsample(const std::vector<Real> & sample, unsigned int offset, unsigned int num_entries) const;
 
  std::vector<Real> windowAverage(const std::vector<Real> & sample, unsigned int num_entries) const;
 
  const unsigned int _timestep_window;
  const bool _shift_outputs;
  const bool _average_reward_over_timestep_window;
  const LibtorchObservationHistoryHelper _observation_history;
};
 
#endif