TrainingSet.h

#ifndef TRAININGSET_H
#define TRAININGSET_H
#include <iostream>
#include "../../alignment/State.h"
#include "../../alignment/Graph.h"
#include "../../alignment/Node.h"
#include "../../alignment/Edge.h"
#include "../SimpleAgent/Agent.h"
#include "../SimpleAgent/Episode.h"
#include "valueMLmodel.h"
#include "Policy.h"
#include "LinearNet.h"
#include "RLDataset.h"
#include <torch/torch.h>
#include <vector>

using std::vector;
using namespace std;

class TrainingSet {

    public:
        TrainingSet(unsigned int learningRepetitions, unsigned int epochs, unsigned int numbEpisodes, unsigned int batchSize, float learningRate); 
    
        unsigned int learningRepetitions;   
        unsigned int epochs;    
        unsigned int batchSize;    
        float learningRate;   
        unsigned int numbEpisodes;   
        
        void setLearningRepetitions(unsigned int learningRepetitions);

        void setEpochs(unsigned int epochs);

        void setNumbEpisodes(unsigned int numbEpisodes);

        void setBatchSize(unsigned int batchSize);

        void setLearningRate(float learningRate);

        void train(Agent* agent);
};
#endif