130 lines
4.0 KiB
TypeScript
130 lines
4.0 KiB
TypeScript
import type { Population } from './evolution';
|
|
import type { EvolutionConfig } from './evolution';
|
|
import { evaluatePopulation, DEFAULT_MATCH_CONFIG } from './selfPlay';
|
|
import { evolveGeneration, createPopulation, getPopulationStats } from './evolution';
|
|
|
|
/**
|
|
* NEAT Training Worker
|
|
*
|
|
* Runs training in a background thread to prevent UI blocking.
|
|
* The main thread only handles visualization and UI updates.
|
|
*/
|
|
|
|
export interface TrainingWorkerMessage {
|
|
type: 'start' | 'pause' | 'step' | 'reset' | 'init';
|
|
config?: EvolutionConfig;
|
|
}
|
|
|
|
export interface TrainingWorkerResponse {
|
|
type: 'update' | 'error' | 'ready';
|
|
population?: Population;
|
|
stats?: ReturnType<typeof getPopulationStats>;
|
|
error?: string;
|
|
}
|
|
|
|
let population: Population | null = null;
|
|
let isRunning = false;
|
|
let config: EvolutionConfig | null = null;
|
|
|
|
/**
|
|
* Handle messages from main thread
|
|
*/
|
|
self.onmessage = async (e: MessageEvent<TrainingWorkerMessage>) => {
|
|
const message = e.data;
|
|
|
|
try {
|
|
switch (message.type) {
|
|
case 'init':
|
|
if (message.config) {
|
|
config = message.config;
|
|
population = createPopulation(config);
|
|
sendUpdate();
|
|
self.postMessage({ type: 'ready' } as TrainingWorkerResponse);
|
|
}
|
|
break;
|
|
|
|
case 'start':
|
|
isRunning = true;
|
|
runTrainingLoop();
|
|
break;
|
|
|
|
case 'pause':
|
|
isRunning = false;
|
|
break;
|
|
|
|
case 'step':
|
|
if (population && config) {
|
|
const stats = await runSingleGeneration();
|
|
sendUpdate(stats);
|
|
}
|
|
break;
|
|
|
|
case 'reset':
|
|
if (config) {
|
|
population = createPopulation(config);
|
|
isRunning = false;
|
|
sendUpdate();
|
|
}
|
|
break;
|
|
}
|
|
} catch (error) {
|
|
self.postMessage({
|
|
type: 'error',
|
|
error: error instanceof Error ? error.message : 'Unknown error',
|
|
} as TrainingWorkerResponse);
|
|
}
|
|
};
|
|
|
|
/**
|
|
* Run continuous training loop
|
|
*/
|
|
async function runTrainingLoop() {
|
|
while (isRunning && population && config) {
|
|
const stats = await runSingleGeneration();
|
|
sendUpdate(stats);
|
|
|
|
// Yield to allow pause/stop messages to be processed
|
|
await new Promise(resolve => setTimeout(resolve, 0));
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Run a single generation
|
|
*/
|
|
async function runSingleGeneration(): Promise<ReturnType<typeof getPopulationStats> | null> {
|
|
if (!population || !config) return null;
|
|
|
|
console.log('[Worker] Starting generation', population.generation);
|
|
|
|
// Evaluate population
|
|
const evaluatedPop = evaluatePopulation(population, DEFAULT_MATCH_CONFIG);
|
|
|
|
// Check fitness after evaluation
|
|
const fitnesses = evaluatedPop.genomes.map(g => g.fitness);
|
|
const avgFit = fitnesses.reduce((a, b) => a + b, 0) / fitnesses.length;
|
|
const maxFit = Math.max(...fitnesses);
|
|
console.log('[Worker] After evaluation - Avg fitness:', avgFit.toFixed(2), 'Max:', maxFit.toFixed(2));
|
|
|
|
// Evolve to next generation
|
|
population = evolveGeneration(evaluatedPop, config);
|
|
|
|
console.log('[Worker] Generation', population.generation, 'complete');
|
|
|
|
// IMPORTANT: Send stats from the EVALUATED population, not the evolved one
|
|
// (evolved population has fitness reset to 0)
|
|
return getPopulationStats(evaluatedPop);
|
|
}
|
|
|
|
/**
|
|
* Send population update to main thread
|
|
*/
|
|
function sendUpdate(stats?: ReturnType<typeof getPopulationStats> | null) {
|
|
if (!population) return;
|
|
|
|
self.postMessage({
|
|
type: 'update',
|
|
population,
|
|
stats: stats || undefined,
|
|
} as TrainingWorkerResponse);
|
|
}
|