@ignitionai/backend-tfjs/dist/tools/trainer.js

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

import { DQNAgent } from "../agents/dqn";
import * as tf from '@tensorflow/tfjs';
export async function trainAgent(options) {
    const { config, maxSteps = 100, checkpointEvery = 10, repoId, token, onStep, getEnvStep } = options;
    const agent = new DQNAgent(config);
    let bestReward = -Infinity;
    for (let step = 1; step <= maxSteps; step++) {
        const { state, correctAction, nextState } = getEnvStep();
        const action = await agent.getAction(state);
        const reward = action === correctAction ? 1 : -1;
        const done = false;
        agent.remember({ state, action, reward, nextState, done });
        await agent.train();
        onStep?.(step, reward, action, state);
        if (checkpointEvery && step % checkpointEvery === 0 && repoId && token) {
            await agent.saveCheckpoint(repoId, token, `step-${step}`);
        }
        if (reward > bestReward && repoId && token) {
            bestReward = reward;
            await agent.saveCheckpoint(repoId, token, 'best');
        }
    }
    agent.dispose();
    tf.disposeVariables();
    tf.dispose();
}