DDPGModel.cpp

#include "DDPGModel.h"
#include <algorithm>
#include <iostream>
#include <fstream>
#include <random>
#include <memory>

// Määritellään STATE_SIZE ja ACTION_SIZE, jos niitä ei ole määritelty
#ifndef STATE_SIZE
#define STATE_SIZE 10
#endif

#ifndef ACTION_SIZE
#define ACTION_SIZE 3
#endif

// OUNoise toteutus
DDPGModel::OUNoise::OUNoise(int size, float mu, float theta, float sigma) 
    : mu(mu), theta(theta), sigma(sigma) {
    state.resize(size, 0.0f);
    std::random_device rd;
    rng = std::mt19937(rd());
    reset();
}

void DDPGModel::OUNoise::reset() {
    for (auto& s : state) {
        s = mu;
    }
}

std::vector<float> DDPGModel::OUNoise::sample() {
    std::normal_distribution<float> dist(0.0f, 1.0f);
    
    for (size_t i = 0; i < state.size(); ++i) {
        float dx = theta * (mu - state[i]);
        float noise = sigma * dist(rng);
        state[i] += dx + noise;
    }
    
    return state;
}

// DDPGNetwork toteutus
DDPGNetwork::DDPGNetwork(float lr, float g, float t, size_t bufferCapacity, size_t batch)
    : actorNetwork(), criticNetwork(), targetActorNetwork(), targetCriticNetwork(),
      learningRate(lr), gamma(g), tau(t), 
      replayBufferCapacity(bufferCapacity), batchSize(batch) {
    
    // Alusta verkot
    actorNetwork.initialize(42);  // Käytä seed-arvoa 42
    criticNetwork.initialize(43);  // Käytä eri seed-arvoa
    
    // Alusta kohina
    actorNetwork.initializeNoise();
    criticNetwork.initializeNoise();
    
    // Kopioidaan painot target-verkkoihin
    targetActorNetwork.copyWeightsFrom(actorNetwork);
    targetCriticNetwork.copyWeightsFrom(criticNetwork);
}

void DDPGNetwork::copyNetworkWeights(const ActorCritic& source, ActorCritic& target) {
    // Käytetään ActorCritic-luokan copyWeightsFrom-metodia
    target.copyWeightsFrom(source);
}

void DDPGNetwork::updateTargetNetworks() {
    // Käytetään try_lock lukituksen sijaan, jotta ei jäädä jumiin
    if (!networkMutex.try_lock()) {
        // Jos lukitus ei onnistu, palataan ilman päivitystä
        return;
    }
    
    // Lukitus onnistui, käytetään std::lock_guard lukituksen automaattiseen vapauttamiseen
    std::lock_guard<std::mutex> lock(networkMutex, std::adopt_lock);
    
    // Soft-update target-verkkoihin
    // θ' = τθ + (1-τ)θ'
    
    // Käytetään ActorCritic-luokan copyWeightsFrom-metodia
    // Yksinkertaisuuden vuoksi käytetään suoraa kopiointia
    targetActorNetwork.copyWeightsFrom(actorNetwork);
    targetCriticNetwork.copyWeightsFrom(criticNetwork);
}

void DDPGNetwork::addExperience(const float* state, const float* action, float reward, 
                              const float* nextState, bool done) {
    // Käytetään try_lock lukituksen sijaan, jotta ei jäädä jumiin
    if (!networkMutex.try_lock()) {
        // Jos lukitus ei onnistu, palataan ilman päivitystä
        return;
    }
    
    // Lukitus onnistui, käytetään std::lock_guard lukituksen automaattiseen vapauttamiseen
    std::lock_guard<std::mutex> lock(networkMutex, std::adopt_lock);
    
    // Luo uusi kokemus käyttäen DDPGExperience-rakennetta
    DDPGExperience exp;
    exp.state.resize(STATE_SIZE);
    exp.action.resize(ACTION_SIZE);
    exp.nextState.resize(STATE_SIZE);
    exp.reward = reward;
    exp.done = done;
    
    // Kopioi tiedot
    std::copy(state, state + STATE_SIZE, exp.state.begin());
    std::copy(action, action + ACTION_SIZE, exp.action.begin());
    std::copy(nextState, nextState + STATE_SIZE, exp.nextState.begin());
    
    // Lisää kokemus puskuriin
    replayBuffer.push_back(exp);
    
    // Rajoita puskurin kokoa
    if (replayBuffer.size() > replayBufferCapacity) {
        replayBuffer.pop_front();
    }
}

void DDPGNetwork::actorForward(const float* state, float* action) {
    // Käytetään try_lock lukituksen sijaan, jotta ei jäädä jumiin
    if (!networkMutex.try_lock()) {
        // Jos lukitus ei onnistu, generoidaan satunnainen toiminto
        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
        for (int i = 0; i < ACTION_SIZE; i++) {
            action[i] = dist(gen);
        }
        return;
    }
    
    // Lukitus onnistui, käytetään std::lock_guard lukituksen automaattiseen vapauttamiseen
    std::lock_guard<std::mutex> lock(networkMutex, std::adopt_lock);
    
    // Käytä ActorCritic-luokan forward-metodia
    float dummyValue;
    
    // Kopioi const float* -> float*
    float stateCopy[STATE_SIZE];
    std::copy(state, state + STATE_SIZE, stateCopy);
    
    actorNetwork.forward(stateCopy, action, dummyValue);
}

float DDPGNetwork::criticForward(const float* state, const float* action) {
    // Käytetään try_lock lukituksen sijaan, jotta ei jäädä jumiin
    if (!networkMutex.try_lock()) {
        // Jos lukitus ei onnistu, palautetaan oletusarvo
        return 0.0f;
    }
    
    // Lukitus onnistui, käytetään std::lock_guard lukituksen automaattiseen vapauttamiseen
    std::lock_guard<std::mutex> lock(networkMutex, std::adopt_lock);
    
    // Laske critic-verkon arvo
    float value = 0.0f;
    
    // Kopioi tila ja toiminto piilokerrokseen
    for (int i = 0; i < HIDDEN_SIZE; i++) {
        float sum = 0.0f;
        for (int j = 0; j < STATE_SIZE; j++) {
            sum += state[j] * criticNetwork.criticInputWeights[j][i];
        }
        criticNetwork.hiddenLayer[i] = std::tanh(sum);
    }
    
    // Laske ulostulokerroksen arvo
    for (int i = 0; i < 1; i++) {
        float sum = 0.0f;
        for (int j = 0; j < HIDDEN_SIZE; j++) {
            sum += criticNetwork.hiddenLayer[j] * criticNetwork.criticHiddenWeights[j][i];
        }
        value = sum;
    }
    
    return value;
}

void DDPGNetwork::updateNetworks() {
    // Käytetään try_lock lukituksen sijaan, jotta ei jäädä jumiin
    if (!networkMutex.try_lock()) {
        // Jos lukitus ei onnistu, palataan ilman päivitystä
        return;
    }
    
    // Lukitus onnistui, käytetään std::lock_guard lukituksen automaattiseen vapauttamiseen
    std::lock_guard<std::mutex> lock(networkMutex, std::adopt_lock);
    
    // Tarkista onko tarpeeksi kokemuksia
    if (replayBuffer.size() < batchSize) {
        return;
    }
    
    // Valitse satunnainen batch kokemuksia
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<size_t> dist(0, replayBuffer.size() - 1);
    
    for (size_t i = 0; i < batchSize; ++i) {
        size_t idx = dist(gen);
        const DDPGExperience& exp = replayBuffer[idx];
        
        // Suorita eteenpäinkulku target-verkoilla
        float nextAction[ACTION_SIZE];
        float nextValue;
        
        // Kopioi const float* -> float*
        float nextStateCopy[STATE_SIZE];
        std::copy(exp.nextState.data(), exp.nextState.data() + STATE_SIZE, nextStateCopy);
        
        targetActorNetwork.forward(nextStateCopy, nextAction, nextValue);
        float targetQ = criticForward(exp.nextState.data(), nextAction);
        
        // Laske TD-virhe
        float target = exp.reward + (exp.done ? 0.0f : gamma * targetQ);
        float currentQ = criticForward(exp.state.data(), exp.action.data());
        float tdError = target - currentQ;
        
        // Päivitä critic-verkko suoraan
        // Laske gradientit ja päivitä painot
        float learningRateC = 0.001f;  // Critic-verkon oppimisnopeuskerroin
        
        // Päivitä critic-verkon painot
        for (int i = 0; i < STATE_SIZE; i++) {
            for (int j = 0; j < HIDDEN_SIZE; j++) {
                criticNetwork.criticInputWeights[i][j] += learningRateC * tdError * exp.state[i] * 
                    (1.0f - criticNetwork.hiddenLayer[j] * criticNetwork.hiddenLayer[j]);
            }
        }
        
        for (int i = 0; i < HIDDEN_SIZE; i++) {
            criticNetwork.criticHiddenWeights[i][0] += learningRateC * tdError * criticNetwork.hiddenLayer[i];
        }
        
        // Päivitä actor-verkko suoraan
        // Laske actor-verkon gradientit ja päivitä painot
        float learningRateA = 0.0001f;  // Actor-verkon oppimisnopeuskerroin
        
        // Suorita eteenpäinkulku actor-verkolla
        float actions[ACTION_SIZE];
        
        // Kopioi const float* -> float*
        float stateCopy[STATE_SIZE];
        std::copy(exp.state.data(), exp.state.data() + STATE_SIZE, stateCopy);
        
        actorForward(exp.state.data(), actions);
        
        // Laske actor-verkon gradientit
        float actorGradients[ACTION_SIZE];
        for (int i = 0; i < ACTION_SIZE; i++) {
            actorGradients[i] = 0.0f;
        }
        
        // Yksinkertaistettu gradienttien laskenta
        for (int i = 0; i < ACTION_SIZE; i++) {
            actorGradients[i] = tdError * (actions[i] - exp.action[i]);
        }
        
        // Päivitä actor-verkon painot
        for (int i = 0; i < STATE_SIZE; i++) {
            for (int j = 0; j < HIDDEN_SIZE; j++) {
                actorNetwork.actorInputWeights[i][j] += learningRateA * actorGradients[0] * exp.state[i] * 
                    (1.0f - actorNetwork.hiddenLayer[j] * actorNetwork.hiddenLayer[j]);
            }
        }
        
        for (int i = 0; i < HIDDEN_SIZE; i++) {
            for (int j = 0; j < ACTION_SIZE; j++) {
                actorNetwork.actorHiddenWeights[i][j] += learningRateA * actorGradients[j] * actorNetwork.hiddenLayer[i];
            }
        }
    }
    
    // Päivitä target-verkot
    updateTargetNetworks();
    
    // Päivitä tilastot
    totalUpdates++;
}

void DDPGNetwork::saveWeights(const std::string& filename) {
    // Käytetään try_lock lukituksen sijaan, jotta ei jäädä jumiin
    if (!networkMutex.try_lock()) {
        // Jos lukitus ei onnistu, heitetään poikkeus
        throw std::runtime_error("Ei voitu lukita verkkoa tallennusta varten");
    }
    
    // Lukitus onnistui, käytetään std::lock_guard lukituksen automaattiseen vapauttamiseen
    std::lock_guard<std::mutex> lock(networkMutex, std::adopt_lock);
    
    std::ofstream file(filename, std::ios::binary);
    if (!file.is_open()) {
        throw std::runtime_error("Ei voitu avata tiedostoa tallennusta varten: " + filename);
    }
    
    // Tallenna actor-verkon painot
    file.write(reinterpret_cast<const char*>(actorNetwork.actorInputWeights), 
              sizeof(actorNetwork.actorInputWeights));
    file.write(reinterpret_cast<const char*>(actorNetwork.actorHiddenWeights), 
              sizeof(actorNetwork.actorHiddenWeights));
    
    // Tallenna critic-verkon painot
    file.write(reinterpret_cast<const char*>(criticNetwork.criticInputWeights), 
              sizeof(criticNetwork.criticInputWeights));
    file.write(reinterpret_cast<const char*>(criticNetwork.criticHiddenWeights), 
              sizeof(criticNetwork.criticHiddenWeights));
    
    // Tallenna hyperparametrit
    file.write(reinterpret_cast<const char*>(&learningRate), sizeof(learningRate));
    file.write(reinterpret_cast<const char*>(&gamma), sizeof(gamma));
    file.write(reinterpret_cast<const char*>(&tau), sizeof(tau));
    
    file.close();
}

void DDPGNetwork::loadWeights(const std::string& filename) {
    // Käytetään try_lock lukituksen sijaan, jotta ei jäädä jumiin
    if (!networkMutex.try_lock()) {
        // Jos lukitus ei onnistu, heitetään poikkeus
        throw std::runtime_error("Ei voitu lukita verkkoa latausta varten");
    }
    
    // Lukitus onnistui, käytetään std::lock_guard lukituksen automaattiseen vapauttamiseen
    std::lock_guard<std::mutex> lock(networkMutex, std::adopt_lock);
    
    std::ifstream file(filename, std::ios::binary);
    if (!file.is_open()) {
        throw std::runtime_error("Ei voitu avata tiedostoa latausta varten: " + filename);
    }
    
    // Lataa actor-verkon painot
    file.read(reinterpret_cast<char*>(actorNetwork.actorInputWeights), 
             sizeof(actorNetwork.actorInputWeights));
    file.read(reinterpret_cast<char*>(actorNetwork.actorHiddenWeights), 
             sizeof(actorNetwork.actorHiddenWeights));
    
    // Lataa critic-verkon painot
    file.read(reinterpret_cast<char*>(criticNetwork.criticInputWeights), 
             sizeof(criticNetwork.criticInputWeights));
    file.read(reinterpret_cast<char*>(criticNetwork.criticHiddenWeights), 
             sizeof(criticNetwork.criticHiddenWeights));
    
    // Lataa hyperparametrit
    file.read(reinterpret_cast<char*>(&learningRate), sizeof(learningRate));
    file.read(reinterpret_cast<char*>(&gamma), sizeof(gamma));
    file.read(reinterpret_cast<char*>(&tau), sizeof(tau));
    
    // Päivitä target-verkot
    targetActorNetwork.copyWeightsFrom(actorNetwork);
    targetCriticNetwork.copyWeightsFrom(criticNetwork);
    
    file.close();
}

void DDPGNetwork::updateLearningRate(float newLR) {
    // Käytetään try_lock lukituksen sijaan, jotta ei jäädä jumiin
    if (!networkMutex.try_lock()) {
        // Jos lukitus ei onnistu, palataan ilman päivitystä
        return;
    }
    
    // Lukitus onnistui, käytetään std::lock_guard lukituksen automaattiseen vapauttamiseen
    std::lock_guard<std::mutex> lock(networkMutex, std::adopt_lock);
    
    learningRate = newLR;
}

void DDPGNetwork::updateAverageReward(float reward) {
    // Päivitä keskimääräinen palkkio
    float oldAvg = averageReward.load();
    float newAvg = oldAvg * 0.99f + reward * 0.01f;
    averageReward.store(newAvg);
}

float DDPGNetwork::getAverageReward() const {
    return averageReward.load();
}

int DDPGNetwork::getTotalUpdates() const {
    return totalUpdates.load();
}

float DDPGNetwork::getLearningRate() const {
    return learningRate;
}

// DDPGModel toteutus
DDPGModel::DDPGModel(float lr, float gamma, float tau)
    : network(new DDPGNetwork(lr, gamma, tau)),
      lastAdvantage(0.0f),
      lastWeightUpdate(0.0f),
      tdError(0.0f),
      updateCounter(0),
      updateFrequency(20),
      noise(ACTION_SIZE) {
    
    // Alusta hyperparametrit
    hyperParams.learningRate = lr;
    hyperParams.momentum = 0.95f;
    hyperParams.noiseScale = 0.2f;
    hyperParams.priorityExponent = 0.6f;
    hyperParams.importanceSamplingBeta = 0.4f;
    hyperParams.updateFrequency = 20;
}

void DDPGModel::forward(float state[STATE_SIZE], float action[ACTION_SIZE], float& value) {
    // Suorita eteenpäinkulku actor-verkolla
    network->actorForward(state, action);
    
    // Lisää kohinaa tutkimista varten
    std::vector<float> noiseValues = noise.sample();
    for (int i = 0; i < ACTION_SIZE; ++i) {
        action[i] += hyperParams.noiseScale * noiseValues[i];
        // Rajoita arvot välille [-1, 1]
        action[i] = std::min(std::max(action[i], -1.0f), 1.0f);
    }
    
    // Suorita eteenpäinkulku critic-verkolla
    value = network->criticForward(state, action);
}

void DDPGModel::collectExperience(float state[STATE_SIZE], float action[ACTION_SIZE], 
                                float reward, float value) {
    // Tallenna kokemus
    std::vector<float> stateVec(state, state + STATE_SIZE);
    std::vector<float> actionVec(action, action + ACTION_SIZE);
    
    states.insert(states.end(), stateVec.begin(), stateVec.end());
    actions.insert(actions.end(), actionVec.begin(), actionVec.end());
    rewards.push_back(reward);
    values.push_back(value);
    
    // Päivitä keskimääräinen palkkio
    network->updateAverageReward(reward);
}

void DDPGModel::update(float nextStateValue, bool isTerminal) {
    // Tarkista onko kokemuksia
    if (rewards.empty()) {
        return;
    }
    
    // Luo seuraava tila
    std::vector<float> nextStateVec(STATE_SIZE);
    // Tässä pitäisi olla oikea seuraava tila, mutta yksinkertaisuuden vuoksi käytetään nollia
    
    // Lisää kokemus replay-puskuriin
    for (size_t i = 0; i < rewards.size(); ++i) {
        size_t stateIdx = i * STATE_SIZE;
        size_t actionIdx = i * ACTION_SIZE;
        
        if (stateIdx + STATE_SIZE <= states.size() && actionIdx + ACTION_SIZE <= actions.size()) {
            network->addExperience(
                &states[stateIdx], 
                &actions[actionIdx], 
                rewards[i], 
                &nextStateVec[0], 
                i == rewards.size() - 1 && isTerminal
            );
        }
    }
    
    // Päivitä verkot
    network->updateNetworks();
    
    // Päivitä target-verkot
    network->updateTargetNetworks();
    
    // Päivitä metriikat
    lastAdvantage = rewards.back() - values.back();
    lastWeightUpdate = 0.01f; // Tässä pitäisi olla oikea painojen muutos
    tdError = rewards.back() + hyperParams.noiseScale * nextStateValue - values.back();
    updateCounter++;
    
    // Tyhjennä kokemukset
    states.clear();
    actions.clear();
    rewards.clear();
    values.clear();
}

bool DDPGModel::shouldUpdate() const {
    return rewards.size() >= updateFrequency;
}

size_t DDPGModel::getRewardsSize() const {
    return rewards.size();
}

std::string DDPGModel::getName() const {
    return "DDPG";
}

float DDPGModel::getLastAdvantage() const {
    return lastAdvantage;
}

float DDPGModel::getLastWeightUpdate() const {
    return lastWeightUpdate;
}

float DDPGModel::getTDError() const {
    return tdError;
}

int DDPGModel::getUpdateCounter() const {
    return updateCounter;
}

void DDPGModel::setUpdateFrequency(int freq) {
    updateFrequency = freq;
}

const HyperParameters& DDPGModel::getHyperParams() const {
    return hyperParams;
}

bool DDPGModel::saveModel(const std::string& filename) const {
    try {
        network->saveWeights(filename);
        return true;
    } catch (const std::exception& e) {
        std::cerr << "Virhe tallennettaessa mallia: " << e.what() << std::endl;
        return false;
    }
}

bool DDPGModel::loadModel(const std::string& filename) {
    try {
        network->loadWeights(filename);
        return true;
    } catch (const std::exception& e) {
        std::cerr << "Virhe ladattaessa mallia: " << e.what() << std::endl;
        return false;
    }
}

// Staattinen metodi, joka luo ja lataa DDPG-mallin
std::unique_ptr<RLModel> DDPGModel::createAndLoad(const std::string& filename) {
    auto model = std::make_unique<DDPGModel>();
    if (model->loadModel(filename)) {
        std::cout << "DDPG-malli ladattu onnistuneesti." << std::endl;
    } else {
        std::cout << "DDPG-mallin lataus epäonnistui, käytetään oletusarvoja." << std::endl;
    }
    return model;
} 

void DDPGModel::setHyperParams(const HyperParameters& params) {
    hyperParams = params;
}