@ignitionai/backend-tfjs/dist/agents/dqn.js

TensorFlow.js backend for IgnitionAI - browser-based reinforcement learning framework

import * as tf from '@tensorflow/tfjs';
import { ReplayBuffer } from '../memory/ReplayBuffer';
import { buildQNetwork } from '../model/BuildMLP';
import { saveModelToHub } from '../io/saveModelToHub';
import { loadModelFromHub } from '../io/loadModel';
export class DQNAgent {
    config;
    model;
    targetModel;
    memory;
    epsilon;
    epsilonDecay;
    minEpsilon;
    gamma;
    batchSize;
    targetUpdateFrequency;
    trainStepCounter = 0;
    actionSize;
    bestReward = -Infinity;
    constructor(config) {
        this.config = config;
        const { inputSize, actionSize, hiddenLayers = [24, 24], gamma = 0.99, epsilon = 1.0, epsilonDecay = 0.995, minEpsilon = 0.01, lr = 0.001, batchSize = 32, memorySize = 10000, targetUpdateFrequency = 1000, } = config;
        this.actionSize = actionSize;
        this.gamma = gamma;
        this.epsilon = epsilon;
        this.epsilonDecay = epsilonDecay;
        this.minEpsilon = minEpsilon;
        this.batchSize = batchSize;
        this.targetUpdateFrequency = targetUpdateFrequency;
        this.model = buildQNetwork(inputSize, actionSize, hiddenLayers, lr);
        this.targetModel = buildQNetwork(inputSize, actionSize, hiddenLayers, lr);
        this.updateTargetModel();
        this.memory = new ReplayBuffer(memorySize);
    }
    async getAction(state) {
        if (Math.random() < this.epsilon) {
            return Math.floor(Math.random() * this.actionSize);
        }
        const stateTensor = tf.tensor2d([state]);
        const qValues = this.model.predict(stateTensor);
        const action = (await qValues.argMax(1).data())[0];
        tf.dispose([stateTensor, qValues]);
        return action;
    }
    remember(exp) {
        this.memory.add(exp);
    }
    async updateTargetModel() {
        this.targetModel.setWeights(this.model.getWeights());
    }
    async train() {
        if (this.memory.size() < this.batchSize)
            return;
        const batch = this.memory.sample(this.batchSize);
        const states = batch.map(e => e.state);
        const nextStates = batch.map(e => e.nextState);
        const stateTensor = tf.tensor2d(states);
        const nextStateTensor = tf.tensor2d(nextStates);
        const qValues = this.model.predict(stateTensor);
        const nextQValues = this.targetModel.predict(nextStateTensor);
        const qArray = qValues.arraySync();
        const nextQArray = nextQValues.arraySync();
        const updatedQ = qArray.map((q, i) => {
            const { action, reward, done } = batch[i];
            q[action] = done ? reward : reward + this.gamma * Math.max(...nextQArray[i]);
            return q;
        });
        const targetTensor = tf.tensor2d(updatedQ);
        await this.model.fit(stateTensor, targetTensor, { epochs: 1, verbose: 0 });
        tf.dispose([stateTensor, nextStateTensor, qValues, nextQValues, targetTensor]);
        if (this.epsilon > this.minEpsilon) {
            this.epsilon *= this.epsilonDecay;
        }
        this.trainStepCounter++;
        if (this.trainStepCounter % this.targetUpdateFrequency === 0) {
            await this.updateTargetModel();
        }
    }
    reset() {
        this.epsilon = this.config.epsilon ?? 1.0;
        this.memory = new ReplayBuffer(this.config.memorySize);
        this.trainStepCounter = 0;
    }
    async saveToHub(repoId, token, modelName = 'model', checkpointName = 'last') {
        console.log(`[DQN] Saving model to HF Hub: ${repoId}`);
        await saveModelToHub(this.model, repoId, token, `model_${checkpointName}`);
    }
    async loadFromHub(repoId, modelPath = 'model.json') {
        console.log(`[DQN] Loading model from HF Hub: ${repoId}`);
        const loadedModel = await loadModelFromHub(repoId, modelPath);
        this.model = loadedModel;
        await this.updateTargetModel();
    }
    /**
   * Save the model under a checkpoint name to Hugging Face Hub.
   * e.g., checkpointName = "last", "best", "step-1000"
   */
    async saveCheckpoint(repoId, token, checkpointName) {
        const folder = `model_${checkpointName}`;
        console.log(`[DQN] Saving checkpoint "${checkpointName}" to HF Hub...`);
        await saveModelToHub(this.model, repoId, token, folder);
        console.log(`[DQN] ✅ Checkpoint "${checkpointName}" saved`);
    }
    async maybeSaveBestCheckpoint(repoId, token, reward, step) {
        console.log(`[DQN] Current best: ${this.bestReward.toFixed(4)}, new reward: ${reward.toFixed(4)}`);
        if (reward > this.bestReward) {
            console.log(`[DQN] 🏆 New best reward: ${reward.toFixed(3)} > ${this.bestReward.toFixed(3)}`);
            this.bestReward = reward;
            const checkpointName = step !== undefined ? `step-${step}` : 'best';
            await this.saveCheckpoint(repoId, token, checkpointName);
        }
    }
    /**
     * Load a checkpointed model from Hugging Face Hub.
     */
    async loadCheckpoint(repoId, checkpointName) {
        const modelPath = `model_${checkpointName}/model.json`;
        console.log(`[DQN] Loading checkpoint "${checkpointName}" from HF Hub...`);
        const model = await loadModelFromHub(repoId, modelPath);
        this.model = model;
        await this.updateTargetModel();
        console.log(`[DQN] ✅ Checkpoint "${checkpointName}" loaded`);
    }
    dispose() {
        console.log(`[DQN] Disposing model...`);
        this.model?.dispose();
        console.log(`[DQN] Model disposed`);
        this.targetModel?.dispose();
        console.log(`[DQN] Target model disposed`);
        this.memory = new ReplayBuffer(0);
        console.log(`[DQN] Memory disposed`);
        console.log(`[DQN] ✅ DQNAgent disposed`);
    }
}