A3CModel.cpp

#include "A3CModel.h"
#include <stdexcept>
#include <memory>

// Oletetaan että A3CWorker-luokka on määritelty muualla
// Tässä tarvitaan forward-deklaraatio A3CWorker-luokalle
class A3CNetwork;
class A3CWorker;

A3CModel::A3CModel(A3CWorker& w) : worker(w) {}

void A3CModel::forward(float state[/*STATE_SIZE*/], float action[/*ACTION_SIZE*/], float& value) {
    worker.forward(state, action, value);
}

void A3CModel::collectExperience(float state[/*STATE_SIZE*/], float action[/*ACTION_SIZE*/], 
                               float reward, float value) {
    worker.collectExperience(state, action, reward, value);
}

void A3CModel::update(float nextStateValue, bool isTerminal) {
    worker.updateGlobalNetwork(nextStateValue, isTerminal);
}

bool A3CModel::shouldUpdate() const {
    return worker.shouldUpdate();
}

size_t A3CModel::getRewardsSize() const {
    return worker.getRewardsSize();
}

std::string A3CModel::getName() const {
    return "A3C";
}

float A3CModel::getLastAdvantage() const {
    return worker.getLocalNetwork().lastAdvantage;
}

float A3CModel::getLastWeightUpdate() const {
    return worker.getLocalNetwork().lastWeightUpdate;
}

float A3CModel::getTDError() const {
    return worker.getLocalNetwork().prevTDError;
}

int A3CModel::getUpdateCounter() const {
    return worker.getLocalNetwork().updateCounter;
}

void A3CModel::setUpdateFrequency(int freq) {
    worker.setUpdateFrequency(freq);
}

const HyperParameters& A3CModel::getHyperParams() const {
    return worker.getLocalNetwork().hyperParams;
}

void A3CModel::setHyperParams(const HyperParameters& params) {
    worker.setHyperParams(params);
}

bool A3CModel::saveModel(const std::string& filename) const {
    try {
        // Tallenna A3C-verkon painot
        worker.getLocalNetwork().saveWeights(filename);
        return true;
    } catch (const std::exception& e) {
        std::cerr << "Virhe tallennettaessa A3C-mallia: " << e.what() << std::endl;
        return false;
    }
}

bool A3CModel::loadModel(const std::string& filename) {
    try {
        // Lataa A3C-verkon painot
        worker.getLocalNetwork().loadWeights(filename);
        return true;
    } catch (const std::exception& e) {
        std::cerr << "Virhe ladattaessa A3C-mallia: " << e.what() << std::endl;
        return false;
    }
}

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

// A3CNetwork-luokan toteutukset
A3CNetwork::A3CNetwork(float lr, float g, float t) 
    : learningRate(lr), gamma(g), tau(t) {
    // Alusta globaali verkko
    globalNetwork.initialize(42);
}

void A3CNetwork::initializeWorker(ActorCritic& worker) {
    std::lock_guard<std::mutex> lock(globalNetworkMutex);
    copyNetworkWeights(globalNetwork, worker);
}

void A3CNetwork::copyNetworkWeights(const ActorCritic& source, ActorCritic& target) {
    // Kopioi actor-verkon painot
    for (int i = 0; i < STATE_SIZE; i++) {
        for (int j = 0; j < HIDDEN_SIZE; j++) {
            target.actorInputWeights[i][j] = source.actorInputWeights[i][j];
            target.criticInputWeights[i][j] = source.criticInputWeights[i][j];
            
            // Kopioi myös target-verkkojen painot
            target.targetActorInputWeights[i][j] = source.targetActorInputWeights[i][j];
            target.targetCriticInputWeights[i][j] = source.targetCriticInputWeights[i][j];
        }
    }
    
    for (int i = 0; i < HIDDEN_SIZE; i++) {
        for (int j = 0; j < ACTION_SIZE; j++) {
            target.actorHiddenWeights[i][j] = source.actorHiddenWeights[i][j];
            target.targetActorHiddenWeights[i][j] = source.targetActorHiddenWeights[i][j];
        }
        target.criticHiddenWeights[i][0] = source.criticHiddenWeights[i][0];
        target.targetCriticHiddenWeights[i][0] = source.targetCriticHiddenWeights[i][0];
    }
}

void A3CNetwork::updateGlobalNetwork(ActorCritic& worker) {
    std::lock_guard<std::mutex> lock(globalNetworkMutex);
    
    // Laske gradienttien keskiarvo ja päivitä globaali verkko
    for (int i = 0; i < STATE_SIZE; i++) {
        for (int j = 0; j < HIDDEN_SIZE; j++) {
            globalNetwork.actorInputWeights[i][j] += worker.momentum[i][j] * learningRate;
            globalNetwork.criticInputWeights[i][j] += worker.criticMomentum[i][j] * learningRate;
        }
    }
    
    for (int i = 0; i < HIDDEN_SIZE; i++) {
        for (int j = 0; j < ACTION_SIZE; j++) {
            globalNetwork.actorHiddenWeights[i][j] += worker.momentumHidden[i][j] * learningRate;
        }
        globalNetwork.criticHiddenWeights[i][0] += worker.criticMomentumHidden[i][0] * learningRate;
    }
    
    // Päivitä target-verkot soft-update-menetelmällä
    updateTargetNetworks();
    
    // Päivitä tilastot
    totalUpdates++;
}

void A3CNetwork::updateTargetNetworks() {
    for (int i = 0; i < STATE_SIZE; i++) {
        for (int j = 0; j < HIDDEN_SIZE; j++) {
            globalNetwork.targetActorInputWeights[i][j] = 
                (1 - tau) * globalNetwork.targetActorInputWeights[i][j] + 
                tau * globalNetwork.actorInputWeights[i][j];
            
            globalNetwork.targetCriticInputWeights[i][j] = 
                (1 - tau) * globalNetwork.targetCriticInputWeights[i][j] + 
                tau * globalNetwork.criticInputWeights[i][j];
        }
    }
    
    for (int i = 0; i < HIDDEN_SIZE; i++) {
        for (int j = 0; j < ACTION_SIZE; j++) {
            globalNetwork.targetActorHiddenWeights[i][j] = 
                (1 - tau) * globalNetwork.targetActorHiddenWeights[i][j] + 
                tau * globalNetwork.actorHiddenWeights[i][j];
        }
        
        globalNetwork.targetCriticHiddenWeights[i][0] = 
            (1 - tau) * globalNetwork.targetCriticHiddenWeights[i][0] + 
            tau * globalNetwork.criticHiddenWeights[i][0];
    }
}

void A3CNetwork::pullGlobalNetworkWeights(ActorCritic& worker) {
    std::lock_guard<std::mutex> lock(globalNetworkMutex);
    copyNetworkWeights(globalNetwork, worker);
}

void A3CNetwork::saveGlobalWeights(const std::string& filename) {
    std::lock_guard<std::mutex> lock(globalNetworkMutex);
    globalNetwork.saveWeights(filename);
    
    // Tallenna myös tilastot erilliseen tiedostoon
    std::string statsFilename = filename + ".stats";
    std::ofstream statsFile(statsFilename);
    if (statsFile.is_open()) {
        statsFile << totalUpdates.load() << " " << averageReward.load() << " " 
                 << learningRate << " " << gamma << " " << tau;
        statsFile.close();
    }
}

void A3CNetwork::loadGlobalWeights(const std::string& filename) {
    std::lock_guard<std::mutex> lock(globalNetworkMutex);
    try {
        globalNetwork.loadWeights(filename);
        
        // Yritä ladata myös tilastot
        std::string statsFilename = filename + ".stats";
        std::ifstream statsFile(statsFilename);
        if (statsFile.is_open()) {
            int updates;
            float avgReward, lr, g, t;
            statsFile >> updates >> avgReward >> lr >> g >> t;
            totalUpdates.store(updates);
            averageReward.store(avgReward);
            learningRate = lr;
            gamma = g;
            tau = t;
            statsFile.close();
        }
    } catch (const std::exception& e) {
        std::cerr << "Virhe ladattaessa painoja: " << e.what() << std::endl;
    }
}

void A3CNetwork::updateLearningRate(float newLR) {
    // Varmista että oppimisnopeudella on minimi
    if (newLR < 0.0003f) {  // Nostettu minimiä 0.0001f -> 0.0003f
        newLR = 0.0003f;
    }
    learningRate = newLR;
}

void A3CNetwork::updateAverageReward(float reward) {
    // Tarkista onko palkkio NaN
    if (std::isnan(reward)) {
        reward = 0.0f;
    }
    
    float currentAvg = averageReward.load();
    // Alusta keskimääräinen palkkio ensimmäisellä kerralla
    if (std::isnan(currentAvg)) {
        averageReward.store(reward);
        return;
    }
    
    // Käytä liukuvaa keskiarvoa
    float newAvg = currentAvg * 0.99f + reward * 0.01f;
    
    // Varmista ettei uusi arvo ole NaN
    if (!std::isnan(newAvg)) {
        averageReward.store(newAvg);
    }
}

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

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

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

// A3CWorker-luokan toteutukset
A3CWorker::A3CWorker(int id, A3CNetwork& global, float g, int updateFreq)
    : workerId(id), globalNetwork(global), gamma(g) {
    // Alusta paikallinen verkko globaalin verkon painoilla
    globalNetwork.initializeWorker(localNetwork);
    
    // Aseta päivitystaajuus
    localNetwork.hyperParams.updateFrequency = updateFreq;
    
    // Varmista että kohinavektorit on alustettu
    localNetwork.initializeNoise();
    
    // Alusta gradienttien keräämiseen tarvittavat vektorit
    states.reserve(localNetwork.hyperParams.updateFrequency * STATE_SIZE);
    actions.reserve(localNetwork.hyperParams.updateFrequency * ACTION_SIZE);
    rewards.reserve(localNetwork.hyperParams.updateFrequency);
    values.reserve(localNetwork.hyperParams.updateFrequency);
}

void A3CWorker::forward(float state[STATE_SIZE], float action[ACTION_SIZE], float& value) {
    // Lisää satunnaisuutta toimintoihin
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<float> dist(-0.4f, 0.4f);  // Pienennetty satunnaisuutta (-0.8f -> -0.4f)
    std::uniform_real_distribution<float> smallDist(-0.2f, 0.2f);  // Pienennetty satunnaisuutta (-0.5f -> -0.2f)
    std::uniform_real_distribution<float> shootDist(0.0f, 1.0f);
    
    // Suorita eteenpäinkulku
    localNetwork.forward(state, action, value);
    
    // Varmista että value ei ole nolla tai NaN
    if (std::isnan(value) || std::abs(value) < 0.001f) {
        value = 0.01f;  // Aseta pieni ei-nolla arvo
    }
    
    // Lisää satunnaisuutta toimintoihin
    for (int i = 0; i < ACTION_SIZE; i++) {
        // Lisää satunnaisuutta vain jos toiminto on lähellä nollaa
        if (std::abs(action[i]) < 0.3f) {  // Säilytetään kynnys
            action[i] += dist(gen);
        } else {
            // Lisää vähemmän satunnaisuutta jos toiminto on jo merkittävä
            action[i] += smallDist(gen);
        }
        
        // Varmista että toiminto on järkevällä alueella
        action[i] = std::tanh(action[i]);
    }
    
    // Lisää erityisesti ampumistoimintoon satunnaisuutta
    // Tämä auttaa tankkia oppimaan ampumisen hyödyt
    if (shootDist(gen) < 0.2f) {  // Pienennetty todennäköisyyttä (25% -> 20%)
        action[2] = 1.0f;  // Pakota ampumistoiminto
    }
}

void A3CWorker::collectExperience(float state[STATE_SIZE], float action[ACTION_SIZE], 
                               float reward, float value) {
    // Tallenna tila
    for (int i = 0; i < STATE_SIZE; i++) {
        states.push_back(state[i]);
    }
    
    // Tallenna toiminto
    for (int i = 0; i < ACTION_SIZE; i++) {
        actions.push_back(action[i]);
    }
    
    // Tallenna palkkio ja arvo
    rewards.push_back(reward);
    values.push_back(value);
}

void A3CWorker::updateGlobalNetwork(float nextStateValue, bool isTerminal) {
    // Laske advantage ja päivitä gradientit
    float R = isTerminal ? 0.0f : nextStateValue;
    
    // Käy läpi kerätyt kokemukset käänteisessä järjestyksessä
    for (int t = rewards.size() - 1; t >= 0; t--) {
        R = rewards[t] + gamma * R;
        float advantage = R - values[t];
        
        // Päivitä gradientit paikallisessa verkossa
        int stateOffset = t * STATE_SIZE;
        int actionOffset = t * ACTION_SIZE;
        
        float state[STATE_SIZE];
        float action[ACTION_SIZE];
        
        // Kopioi tila ja toiminto vektoreista
        for (int i = 0; i < STATE_SIZE; i++) {
            state[i] = states[stateOffset + i];
        }
        
        for (int i = 0; i < ACTION_SIZE; i++) {
            action[i] = actions[actionOffset + i];
        }
        
        // Päivitä gradientit
        localNetwork.calculateGradients(state, action, advantage, R);
    }
    
    // Päivitä globaali verkko
    globalNetwork.updateGlobalNetwork(localNetwork);
    
    // Päivitä paikallinen verkko globaalin verkon painoilla
    globalNetwork.pullGlobalNetworkWeights(localNetwork);
    
    // Varmista että updateCounter kasvaa
    localNetwork.updateCounter++;
    
    // Päivitä keskimääräinen palkkio
    float avgReward = 0.0f;
    for (float r : rewards) {
        avgReward += r;
    }
    avgReward /= rewards.size();
    globalNetwork.updateAverageReward(avgReward);
    
    // Tyhjennä gradienttien keräämiseen käytetyt vektorit
    states.clear();
    actions.clear();
    rewards.clear();
    values.clear();
}

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

int A3CWorker::getId() const {
    return workerId;
}

ActorCritic& A3CWorker::getLocalNetwork() const {
    return const_cast<ActorCritic&>(localNetwork);
}

void A3CWorker::setUpdateFrequency(int freq) {
    localNetwork.hyperParams.updateFrequency = freq;
}

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

void A3CWorker::setHyperParams(const HyperParameters& params) {
    localNetwork.hyperParams = params;
}