@sparkleideas/neural 3.5.2-patch.1

package/__tests__/algorithms.test.ts

@@ -0,0 +1,582 @@
+/**
+ * RL Algorithms Tests
+ *
+ * Tests for reinforcement learning algorithms:
+ * - Q-Learning
+ * - SARSA
+ * - DQN
+ * - PPO
+ * - Decision Transformer
+ *
+ * Performance target: <10ms per update
+ */
+
+import { describe, it, expect, beforeEach } from 'vitest';
+import { QLearning, createQLearning } from '../src/algorithms/q-learning.js';
+import { SARSAAlgorithm, createSARSA } from '../src/algorithms/sarsa.js';
+import { DQNAlgorithm, createDQN } from '../src/algorithms/dqn.js';
+import { PPOAlgorithm, createPPO } from '../src/algorithms/ppo.js';
+import { DecisionTransformer, createDecisionTransformer } from '../src/algorithms/decision-transformer.js';
+import type { Trajectory } from '../src/types.js';
+
+// Helper function to create test trajectories
+function createTestTrajectory(steps: number = 5): Trajectory {
+  return {
+    trajectoryId: `test-traj-${Date.now()}`,
+    context: 'Test task',
+    domain: 'code',
+    steps: Array.from({ length: steps }, (_, i) => ({
+      stepId: `step-${i}`,
+      timestamp: Date.now() + i * 100,
+      action: `action-${i % 4}`, // 4 discrete actions
+      stateBefore: new Float32Array(768).fill(i * 0.1),
+      stateAfter: new Float32Array(768).fill((i + 1) * 0.1),
+      reward: 0.5 + (i / steps) * 0.5, // Increasing rewards
+    })),
+    qualityScore: 0.75,
+    isComplete: true,
+    startTime: Date.now() - 1000,
+    endTime: Date.now(),
+  };
+}
+
+describe('Q-Learning Algorithm', () => {
+  let qlearning: QLearning;
+
+  beforeEach(() => {
+    qlearning = createQLearning({
+      learningRate: 0.1,
+      gamma: 0.99,
+      explorationInitial: 1.0,
+      explorationFinal: 0.01,
+      explorationDecay: 1000,
+    });
+  });
+
+  it('should initialize correctly', () => {
+    expect(qlearning).toBeDefined();
+    const stats = qlearning.getStats();
+    expect(stats.updateCount).toBe(0);
+    expect(stats.qTableSize).toBe(0);
+    expect(stats.epsilon).toBeCloseTo(1.0);
+  });
+
+  it('should update Q-values from trajectory', () => {
+    const trajectory = createTestTrajectory(5);
+    const result = qlearning.update(trajectory);
+
+    expect(result.tdError).toBeGreaterThanOrEqual(0);
+    const stats = qlearning.getStats();
+    expect(stats.updateCount).toBe(1);
+    expect(stats.qTableSize).toBeGreaterThan(0);
+  });
+
+  it('should update under performance target (<1ms)', () => {
+    const trajectory = createTestTrajectory(10);
+
+    const startTime = performance.now();
+    qlearning.update(trajectory);
+    const elapsed = performance.now() - startTime;
+
+    expect(elapsed).toBeLessThan(10); // Reasonable target for small trajectories
+  });
+
+  it('should decay exploration rate', () => {
+    const trajectory = createTestTrajectory(5);
+    const initialEpsilon = qlearning.getStats().epsilon;
+
+    for (let i = 0; i < 10; i++) {
+      qlearning.update(trajectory);
+    }
+
+    const finalEpsilon = qlearning.getStats().epsilon;
+    expect(finalEpsilon).toBeLessThan(initialEpsilon);
+  });
+
+  it('should select actions with epsilon-greedy', () => {
+    const state = new Float32Array(768).fill(0.5);
+
+    // First call should be random (high epsilon)
+    const action1 = qlearning.getAction(state, true);
+    expect(action1).toBeGreaterThanOrEqual(0);
+    expect(action1).toBeLessThan(4);
+
+    // Without exploration, should be deterministic
+    const action2 = qlearning.getAction(state, false);
+    expect(action2).toBeDefined();
+  });
+
+  it('should return Q-values for a state', () => {
+    const trajectory = createTestTrajectory(5);
+    qlearning.update(trajectory);
+
+    const state = new Float32Array(768).fill(0.5);
+    const qValues = qlearning.getQValues(state);
+
+    expect(qValues).toBeInstanceOf(Float32Array);
+    expect(qValues.length).toBe(4);
+  });
+
+  it('should handle eligibility traces', () => {
+    const qlearningWithTraces = createQLearning({
+      useEligibilityTraces: true,
+      traceDecay: 0.9,
+    });
+
+    const trajectory = createTestTrajectory(10);
+    expect(() => qlearningWithTraces.update(trajectory)).not.toThrow();
+  });
+
+  it('should prune Q-table when over capacity', () => {
+    const smallQLearning = createQLearning({
+      maxStates: 10,
+    });
+
+    // Add many different trajectories to fill Q-table
+    for (let i = 0; i < 20; i++) {
+      const trajectory = createTestTrajectory(5);
+      smallQLearning.update(trajectory);
+    }
+
+    const stats = smallQLearning.getStats();
+    expect(stats.qTableSize).toBeLessThanOrEqual(10);
+  });
+
+  it('should reset correctly', () => {
+    const trajectory = createTestTrajectory(5);
+    qlearning.update(trajectory);
+
+    qlearning.reset();
+    const stats = qlearning.getStats();
+
+    expect(stats.updateCount).toBe(0);
+    expect(stats.qTableSize).toBe(0);
+    expect(stats.epsilon).toBeCloseTo(1.0);
+  });
+});
+
+describe('SARSA Algorithm', () => {
+  let sarsa: SARSAAlgorithm;
+
+  beforeEach(() => {
+    sarsa = createSARSA({
+      learningRate: 0.1,
+      gamma: 0.99,
+      explorationInitial: 1.0,
+      explorationFinal: 0.01,
+      explorationDecay: 1000,
+    });
+  });
+
+  it('should initialize correctly', () => {
+    expect(sarsa).toBeDefined();
+    const stats = sarsa.getStats();
+    expect(stats.updateCount).toBe(0);
+    expect(stats.qTableSize).toBe(0);
+  });
+
+  it('should update using SARSA rule', () => {
+    const trajectory = createTestTrajectory(5);
+    const result = sarsa.update(trajectory);
+
+    expect(result.tdError).toBeGreaterThanOrEqual(0);
+    const stats = sarsa.getStats();
+    expect(stats.updateCount).toBe(1);
+  });
+
+  it('should handle expected SARSA variant', () => {
+    const expectedSARSA = createSARSA({
+      useExpectedSARSA: true,
+    });
+
+    const trajectory = createTestTrajectory(5);
+    expect(() => expectedSARSA.update(trajectory)).not.toThrow();
+  });
+
+  it('should return action probabilities', () => {
+    const state = new Float32Array(768).fill(0.5);
+    const probs = sarsa.getActionProbabilities(state);
+
+    expect(probs).toBeInstanceOf(Float32Array);
+    expect(probs.length).toBe(4);
+
+    // Probabilities should sum to ~1
+    const sum = Array.from(probs).reduce((a, b) => a + b, 0);
+    expect(sum).toBeCloseTo(1.0, 2);
+  });
+
+  it('should select actions with epsilon-greedy policy', () => {
+    const state = new Float32Array(768).fill(0.5);
+    const action = sarsa.getAction(state, true);
+
+    expect(action).toBeGreaterThanOrEqual(0);
+    expect(action).toBeLessThan(4);
+  });
+
+  it('should handle eligibility traces (SARSA-lambda)', () => {
+    const sarsaLambda = createSARSA({
+      useEligibilityTraces: true,
+      traceDecay: 0.9,
+    });
+
+    const trajectory = createTestTrajectory(10);
+    expect(() => sarsaLambda.update(trajectory)).not.toThrow();
+  });
+
+  it('should handle short trajectories gracefully', () => {
+    const shortTrajectory = createTestTrajectory(1);
+    const result = sarsa.update(shortTrajectory);
+
+    expect(result.tdError).toBe(0); // Not enough steps for SARSA
+  });
+
+  it('should reset algorithm state', () => {
+    const trajectory = createTestTrajectory(5);
+    sarsa.update(trajectory);
+
+    sarsa.reset();
+    const stats = sarsa.getStats();
+
+    expect(stats.updateCount).toBe(0);
+    expect(stats.qTableSize).toBe(0);
+  });
+});
+
+describe('DQN Algorithm', () => {
+  let dqn: DQNAlgorithm;
+
+  beforeEach(() => {
+    dqn = createDQN({
+      learningRate: 0.0001,
+      bufferSize: 1000,
+      miniBatchSize: 32,
+      doubleDQN: true,
+      targetUpdateFreq: 100,
+    });
+  });
+
+  it('should initialize correctly', () => {
+    expect(dqn).toBeDefined();
+    const stats = dqn.getStats();
+    expect(stats.updateCount).toBe(0);
+    expect(stats.bufferSize).toBe(0);
+  });
+
+  it('should add experience to replay buffer', () => {
+    const trajectory = createTestTrajectory(10);
+    dqn.addExperience(trajectory);
+
+    const stats = dqn.getStats();
+    expect(stats.bufferSize).toBe(10);
+  });
+
+  it('should perform DQN update', () => {
+    // Add enough experiences
+    for (let i = 0; i < 5; i++) {
+      dqn.addExperience(createTestTrajectory(10));
+    }
+
+    const result = dqn.update();
+    expect(result.loss).toBeGreaterThanOrEqual(0);
+    expect(result.epsilon).toBeGreaterThan(0);
+  });
+
+  it('should update under performance target (<10ms)', () => {
+    // Add experiences
+    for (let i = 0; i < 5; i++) {
+      dqn.addExperience(createTestTrajectory(10));
+    }
+
+    const startTime = performance.now();
+    dqn.update();
+    const elapsed = performance.now() - startTime;
+
+    // Allow generous overhead for neural network in test environment
+    // (actual production target is <10ms, but tests run in CI may be slower)
+    expect(elapsed).toBeLessThan(500);
+  });
+
+  it('should use double DQN when enabled', () => {
+    const doubleDQN = createDQN({
+      doubleDQN: true,
+      miniBatchSize: 16,
+    });
+
+    for (let i = 0; i < 3; i++) {
+      doubleDQN.addExperience(createTestTrajectory(10));
+    }
+
+    expect(() => doubleDQN.update()).not.toThrow();
+  });
+
+  it('should select actions with epsilon-greedy', () => {
+    const state = new Float32Array(768).fill(0.5);
+    const action = dqn.getAction(state, true);
+
+    expect(action).toBeGreaterThanOrEqual(0);
+    expect(action).toBeLessThan(4);
+  });
+
+  it('should return Q-values for a state', () => {
+    const state = new Float32Array(768).fill(0.5);
+    const qValues = dqn.getQValues(state);
+
+    expect(qValues).toBeInstanceOf(Float32Array);
+    expect(qValues.length).toBe(4);
+  });
+
+  it('should update target network periodically', () => {
+    const dqnWithFreqUpdate = createDQN({
+      targetUpdateFreq: 5,
+      miniBatchSize: 16,
+    });
+
+    for (let i = 0; i < 3; i++) {
+      dqnWithFreqUpdate.addExperience(createTestTrajectory(10));
+    }
+
+    // Perform multiple updates to trigger target network update
+    for (let i = 0; i < 10; i++) {
+      dqnWithFreqUpdate.update();
+    }
+
+    const stats = dqnWithFreqUpdate.getStats();
+    expect(stats.stepCount).toBeGreaterThan(5);
+  });
+
+  it('should handle circular replay buffer correctly', () => {
+    const smallDQN = createDQN({
+      bufferSize: 10,
+      miniBatchSize: 4,
+    });
+
+    // Add more experiences than buffer size
+    for (let i = 0; i < 15; i++) {
+      smallDQN.addExperience(createTestTrajectory(2));
+    }
+
+    const stats = smallDQN.getStats();
+    expect(stats.bufferSize).toBe(10);
+  });
+});
+
+describe('PPO Algorithm', () => {
+  let ppo: PPOAlgorithm;
+
+  beforeEach(() => {
+    ppo = createPPO({
+      learningRate: 0.0003,
+      clipRange: 0.2,
+      gaeLambda: 0.95,
+      epochs: 4,
+      miniBatchSize: 64,
+    });
+  });
+
+  it('should initialize correctly', () => {
+    expect(ppo).toBeDefined();
+    const stats = ppo.getStats();
+    expect(stats.updateCount).toBe(0);
+  });
+
+  it('should add experience from trajectory', () => {
+    const trajectory = createTestTrajectory(10);
+    expect(() => ppo.addExperience(trajectory)).not.toThrow();
+
+    const stats = ppo.getStats();
+    expect(stats.bufferSize).toBe(10);
+  });
+
+  it('should perform PPO update with clipping', () => {
+    // Add enough experiences
+    for (let i = 0; i < 10; i++) {
+      ppo.addExperience(createTestTrajectory(10));
+    }
+
+    const result = ppo.update();
+
+    // Policy loss can be negative in PPO (we minimize -surrogate_objective)
+    expect(typeof result.policyLoss).toBe('number');
+    expect(result.valueLoss).toBeGreaterThanOrEqual(0);
+    expect(result.entropy).toBeGreaterThanOrEqual(0);
+  });
+
+  it('should update under performance target (<10ms for small batches)', () => {
+    const smallPPO = createPPO({
+      miniBatchSize: 16,
+      epochs: 1,
+    });
+
+    for (let i = 0; i < 3; i++) {
+      smallPPO.addExperience(createTestTrajectory(10));
+    }
+
+    const startTime = performance.now();
+    smallPPO.update();
+    const elapsed = performance.now() - startTime;
+
+    expect(elapsed).toBeLessThan(100); // Allow overhead for PPO complexity
+  });
+
+  it('should compute GAE advantages', () => {
+    const trajectory = createTestTrajectory(20);
+    expect(() => ppo.addExperience(trajectory)).not.toThrow();
+
+    // Verify experiences were added with advantages
+    const stats = ppo.getStats();
+    expect(stats.bufferSize).toBe(20);
+  });
+
+  it('should sample actions from policy', () => {
+    const state = new Float32Array(768).fill(0.5);
+    const result = ppo.getAction(state);
+
+    expect(result.action).toBeGreaterThanOrEqual(0);
+    expect(result.action).toBeLessThan(4);
+    expect(result.logProb).toBeDefined();
+    expect(result.value).toBeDefined();
+  });
+
+  it('should handle multiple training epochs', () => {
+    const multiEpochPPO = createPPO({
+      epochs: 8,
+      miniBatchSize: 32,
+    });
+
+    for (let i = 0; i < 5; i++) {
+      multiEpochPPO.addExperience(createTestTrajectory(10));
+    }
+
+    expect(() => multiEpochPPO.update()).not.toThrow();
+  });
+
+  it('should clear buffer after update', () => {
+    for (let i = 0; i < 10; i++) {
+      ppo.addExperience(createTestTrajectory(10));
+    }
+
+    ppo.update();
+    const stats = ppo.getStats();
+
+    expect(stats.bufferSize).toBe(0);
+  });
+});
+
+describe('Decision Transformer', () => {
+  let dt: DecisionTransformer;
+
+  beforeEach(() => {
+    dt = createDecisionTransformer({
+      contextLength: 20,
+      numHeads: 4,
+      numLayers: 2,
+      hiddenDim: 64,
+      embeddingDim: 32,
+    });
+  });
+
+  it('should initialize correctly', () => {
+    expect(dt).toBeDefined();
+    const stats = dt.getStats();
+    expect(stats.updateCount).toBe(0);
+    expect(stats.bufferSize).toBe(0);
+    expect(stats.contextLength).toBe(20);
+    expect(stats.numLayers).toBe(2);
+  });
+
+  it('should add complete trajectories to buffer', () => {
+    const trajectory = createTestTrajectory(10);
+    dt.addTrajectory(trajectory);
+
+    const stats = dt.getStats();
+    expect(stats.bufferSize).toBe(1);
+  });
+
+  it('should not add incomplete trajectories', () => {
+    const incompleteTrajectory: Trajectory = {
+      ...createTestTrajectory(5),
+      isComplete: false,
+    };
+
+    dt.addTrajectory(incompleteTrajectory);
+    const stats = dt.getStats();
+
+    expect(stats.bufferSize).toBe(0);
+  });
+
+  it('should train on buffered trajectories', () => {
+    // Add multiple trajectories
+    for (let i = 0; i < 5; i++) {
+      dt.addTrajectory(createTestTrajectory(10));
+    }
+
+    const result = dt.train();
+
+    expect(result.loss).toBeGreaterThanOrEqual(0);
+    expect(result.accuracy).toBeGreaterThanOrEqual(0);
+    expect(result.accuracy).toBeLessThanOrEqual(1);
+  });
+
+  it('should train under performance target (<10ms per batch)', () => {
+    for (let i = 0; i < 3; i++) {
+      dt.addTrajectory(createTestTrajectory(5));
+    }
+
+    const startTime = performance.now();
+    dt.train();
+    const elapsed = performance.now() - startTime;
+
+    expect(elapsed).toBeLessThan(100); // Allow overhead for transformer
+  });
+
+  it('should get action conditioned on target return', () => {
+    const states = [
+      new Float32Array(768).fill(0.1),
+      new Float32Array(768).fill(0.2),
+      new Float32Array(768).fill(0.3),
+    ];
+    const actions = [0, 1, 2];
+    const targetReturn = 0.9;
+
+    const action = dt.getAction(states, actions, targetReturn);
+
+    expect(action).toBeGreaterThanOrEqual(0);
+    expect(action).toBeLessThan(4);
+  });
+
+  it('should handle causal attention masking', () => {
+    // Train with sequence data
+    for (let i = 0; i < 5; i++) {
+      dt.addTrajectory(createTestTrajectory(15));
+    }
+
+    expect(() => dt.train()).not.toThrow();
+  });
+
+  it('should maintain bounded trajectory buffer', () => {
+    // Add more than max capacity (1000)
+    for (let i = 0; i < 1100; i++) {
+      dt.addTrajectory(createTestTrajectory(5));
+    }
+
+    const stats = dt.getStats();
+    expect(stats.bufferSize).toBe(1000);
+  });
+
+  it('should handle varying trajectory lengths', () => {
+    dt.addTrajectory(createTestTrajectory(3));
+    dt.addTrajectory(createTestTrajectory(10));
+    dt.addTrajectory(createTestTrajectory(25));
+
+    expect(() => dt.train()).not.toThrow();
+  });
+
+  it('should compute returns-to-go correctly', () => {
+    const trajectory = createTestTrajectory(5);
+    dt.addTrajectory(trajectory);
+
+    expect(() => dt.train()).not.toThrow();
+    const stats = dt.getStats();
+    expect(stats.avgLoss).toBeGreaterThanOrEqual(0);
+  });
+});