proffesor-for-testing__agentic-qe/tests/learning/SARSALearner.test.ts

/**
 * SARSALearner Tests
 *
 * Comprehensive unit tests for SARSA (State-Action-Reward-State-Action) algorithm
 * Tests on-policy TD(0) learning, convergence, and comparison with Q-Learning
 */

import { SARSALearner } from '../../src/learning/algorithms/SARSALearner';
import { TaskState, AgentAction, TaskExperience } from '../../src/learning/types';

describe('SARSALearner', () => {
  let sarsa: SARSALearner;

  // Helper: Create a simple task state
  const createState = (complexity: number = 0.5): TaskState => ({
    taskComplexity: complexity,
    requiredCapabilities: ['test', 'analyze'],
    contextFeatures: {},
    previousAttempts: 0,
    availableResources: 1.0
  });

  // Helper: Create an action
  const createAction = (strategy: string): AgentAction => ({
    strategy,
    toolsUsed: [],
    parallelization: 0.5,
    retryPolicy: 'exponential',
    resourceAllocation: 0.5
  });

  // Helper: Create an experience
  const createExperience = (
    state: TaskState,
    action: AgentAction,
    reward: number,
    nextState: TaskState
  ): TaskExperience => ({
    taskId: 'test-task',
    taskType: 'unit-test',
    state,
    action,
    reward,
    nextState,
    timestamp: new Date(),
    agentId: 'sarsa-test'
  });

  beforeEach(() => {
    sarsa = new SARSALearner({
      learningRate: 0.1,
      discountFactor: 0.9,
      explorationRate: 0.2,
      explorationDecay: 0.99,
      minExplorationRate: 0.01,
      useExperienceReplay: false, // Disable for unit tests
      replayBufferSize: 1000,
      batchSize: 32
    });
  });

  describe('Initialization', () => {
    it('should initialize with correct default config', () => {
      const defaultSarsa = new SARSALearner();
      const stats = defaultSarsa.getStatistics();

      expect(stats.explorationRate).toBeCloseTo(0.3);
      expect(stats.steps).toBe(0);
      expect(stats.episodes).toBe(0);
      expect(stats.tableSize).toBe(0);
    });

    it('should initialize with custom config', () => {
      const customSarsa = new SARSALearner({
        learningRate: 0.5,
        explorationRate: 0.4
      });

      expect(customSarsa.getExplorationRate()).toBeCloseTo(0.4);
    });

    it('should identify as SARSA algorithm', () => {
      expect(sarsa.getAlgorithmName()).toBe('SARSA');
      expect(sarsa.getAlgorithmType()).toBe('on-policy');
    });
  });

  describe('Action Selection', () => {
    it('should select an action from available actions', () => {
      const state = createState();
      const actions = [
        createAction('strategy1'),
        createAction('strategy2'),
        createAction('strategy3')
      ];

      const selected = sarsa.selectAction(state, actions);
      expect(actions).toContain(selected);
    });

    it('should throw error when no actions available', () => {
      const state = createState();
      expect(() => sarsa.selectAction(state, [])).toThrow('No available actions');
    });

    it('should explore with probability epsilon', () => {
      // With high exploration rate, should pick non-greedy actions sometimes
      const highExploreSarsa = new SARSALearner({
        explorationRate: 1.0 // Always explore
      });

      const state = createState();
      const actions = [
        createAction('random1'),
        createAction('random2'),
        createAction('random3')
      ];

      const selections = new Set<string>();
      for (let i = 0; i < 20; i++) {
        const action = highExploreSarsa.selectAction(state, actions);
        selections.add(action.strategy);
      }

      // Should have explored multiple actions
      expect(selections.size).toBeGreaterThan(1);
    });

    it('should exploit when exploration rate is low', () => {
      const lowExploreSarsa = new SARSALearner({
        explorationRate: 0.0 // Never explore
      });

      const state = createState();
      const action1 = createAction('best');
      const action2 = createAction('worse');

      // Train Q-values
      const exp1 = createExperience(state, action1, 1.0, createState());
      const exp2 = createExperience(state, action2, 0.1, createState());

      lowExploreSarsa.update(exp1);
      lowExploreSarsa.update(exp2);

      // Should always select best action
      for (let i = 0; i < 10; i++) {
        const selected = lowExploreSarsa.selectAction(state, [action1, action2]);
        expect(selected.strategy).toBe('best');
      }
    });
  });

  describe('Q-Value Updates (SARSA Rule)', () => {
    it('should update Q-values using SARSA update rule', () => {
      const state = createState(0.5);
      const nextState = createState(0.6);
      const action = createAction('test');
      const nextAction = createAction('next');

      // Initial Q-value should be 0
      expect(sarsa.getQValue(state, action)).toBe(0);

      // Create experience with positive reward
      const experience = createExperience(state, action, 1.0, nextState);

      // Update with next action (SARSA)
      sarsa.update(experience, nextAction);

      // Q-value should be positive after update
      const qValue = sarsa.getQValue(state, action);
      expect(qValue).toBeGreaterThan(0);
      expect(qValue).toBeLessThanOrEqual(1.0);
    });

    it('should use actual next action Q-value (on-policy)', () => {
      const state = createState();
      const nextState = createState();
      const currentAction = createAction('current');
      const goodNextAction = createAction('good-next');
      const badNextAction = createAction('bad-next');

      // Setup: train next state to have different Q-values for different actions
      const setupExp1 = createExperience(nextState, goodNextAction, 1.0, createState());
      const setupExp2 = createExperience(nextState, badNextAction, -1.0, createState());
      sarsa.update(setupExp1);
      sarsa.update(setupExp2);

      // Now test SARSA: should use the actual next action taken
      const experience = createExperience(state, currentAction, 0, nextState);

      // Update with good next action
      sarsa.update(experience, goodNextAction);
      const qValueGood = sarsa.getQValue(state, currentAction);

      // Reset and update with bad next action
      sarsa.reset();
      sarsa.update(setupExp1);
      sarsa.update(setupExp2);
      sarsa.update(experience, badNextAction);
      const qValueBad = sarsa.getQValue(state, currentAction);

      // Q-value should differ based on which action was actually taken
      expect(qValueGood).toBeGreaterThan(qValueBad);
    });

    it('should handle terminal states correctly', () => {
      const state = createState();
      const action = createAction('terminal-action');
      const terminalState = createState();

      // Terminal state has reward but no next value
      const terminalExp = createExperience(state, action, 10.0, terminalState);

      // Update without next action (terminal)
      sarsa.update(terminalExp);

      const qValue = sarsa.getQValue(state, action);

      // Q-value should be based on reward only (no future value)
      // Q = 0 + α[r + γ*0 - 0] = α*r = 0.1*10 = 1.0
      expect(qValue).toBeCloseTo(1.0, 5);
    });

    it('should increment step count on each update', () => {
      const state = createState();
      const action = createAction('test');
      const experience = createExperience(state, action, 0.5, createState());

      expect(sarsa.getStepCount()).toBe(0);

      sarsa.update(experience);
      expect(sarsa.getStepCount()).toBe(1);

      sarsa.update(experience);
      expect(sarsa.getStepCount()).toBe(2);
    });
  });

  describe('Episode Learning', () => {
    it('should learn from complete episode trajectory', () => {
      const trajectory = [
        {
          state: createState(0.3),
          action: createAction('step1'),
          reward: 0.5
        },
        {
          state: createState(0.5),
          action: createAction('step2'),
          reward: 1.0
        },
        {
          state: createState(0.7),
          action: createAction('step3'),
          reward: 2.0
        }
      ];

      sarsa.learnFromEpisode(trajectory);

      // Should have updated Q-values for all steps
      expect(sarsa.getQValue(trajectory[0].state, trajectory[0].action)).toBeGreaterThan(0);
      expect(sarsa.getQValue(trajectory[1].state, trajectory[1].action)).toBeGreaterThan(0);

      // Episode count should increment
      expect(sarsa.getEpisodeCount()).toBe(1);
    });

    it('should handle empty trajectory', () => {
      expect(() => sarsa.learnFromEpisode([])).not.toThrow();
      expect(sarsa.getEpisodeCount()).toBe(1);
    });

    it('should handle single-step trajectory', () => {
      const trajectory = [
        {
          state: createState(),
          action: createAction('only-step'),
          reward: 1.0
        }
      ];

      sarsa.learnFromEpisode(trajectory);
      expect(sarsa.getQValue(trajectory[0].state, trajectory[0].action)).toBeCloseTo(0.1, 5);
    });
  });

  describe('SelectAndUpdate Flow', () => {
    it('should select next action and update Q-values', () => {
      const currentState = createState(0.4);
      const currentAction = createAction('current');
      const nextState = createState(0.6);
      const availableActions = [
        createAction('next1'),
        createAction('next2')
      ];

      const nextAction = sarsa.selectAndUpdate(
        currentState,
        currentAction,
        1.0,
        nextState,
        availableActions
      );

      // Should return a valid next action
      expect(availableActions.map(a => a.strategy)).toContain(nextAction.strategy);

      // Should have updated Q-value
      expect(sarsa.getQValue(currentState, currentAction)).toBeGreaterThan(0);

      // Step count should increment
      expect(sarsa.getStepCount()).toBe(1);
    });
  });

  describe('Exploration Decay', () => {
    it('should decay exploration rate on end episode', () => {
      const initialRate = sarsa.getExplorationRate();

      sarsa.endEpisode();

      const newRate = sarsa.getExplorationRate();
      expect(newRate).toBeLessThan(initialRate);
      expect(sarsa.getEpisodeCount()).toBe(1);
    });

    it('should not decay below minimum exploration rate', () => {
      const minRate = 0.01;
      const decayingSarsa = new SARSALearner({
        explorationRate: 0.02,
        explorationDecay: 0.5,
        minExplorationRate: minRate
      });

      // Decay many times
      for (let i = 0; i < 100; i++) {
        decayingSarsa.endEpisode();
      }

      expect(decayingSarsa.getExplorationRate()).toBeCloseTo(minRate, 5);
    });
  });

  describe('State and Action Encoding', () => {
    it('should encode similar states to same key', () => {
      const state1 = createState(0.51);
      const state2 = createState(0.54); // Should round to same value
      const action = createAction('test');

      const exp1 = createExperience(state1, action, 1.0, createState());
      const exp2 = createExperience(state2, action, 1.0, createState());

      sarsa.update(exp1);
      sarsa.update(exp2);

      // Should share Q-value due to discretization
      expect(sarsa.getQValue(state1, action)).toBeCloseTo(sarsa.getQValue(state2, action));
    });

    it('should encode different states to different keys', () => {
      const state1 = createState(0.3);
      const state2 = createState(0.8); // Significantly different
      const action = createAction('test');

      const exp1 = createExperience(state1, action, 1.0, createState());
      const exp2 = createExperience(state2, action, -1.0, createState());

      sarsa.update(exp1);
      sarsa.update(exp2);

      // Should have different Q-values
      const q1 = sarsa.getQValue(state1, action);
      const q2 = sarsa.getQValue(state2, action);
      expect(Math.abs(q1 - q2)).toBeGreaterThan(0.01);
    });
  });

  describe('Statistics', () => {
    it('should provide accurate statistics', () => {
      const state = createState();
      const action1 = createAction('a1');
      const action2 = createAction('a2');

      sarsa.update(createExperience(state, action1, 1.0, createState()));
      sarsa.update(createExperience(state, action2, 0.5, createState()));

      const stats = sarsa.getStatistics();

      expect(stats.steps).toBe(2);
      expect(stats.tableSize).toBeGreaterThan(0);
      expect(stats.avgQValue).toBeGreaterThan(0);
      expect(stats.maxQValue).toBeGreaterThanOrEqual(stats.avgQValue);
      expect(stats.minQValue).toBeLessThanOrEqual(stats.avgQValue);
    });

    it('should provide detailed statistics', () => {
      const detailed = sarsa.getDetailedStatistics();

      expect(detailed.algorithm).toBe('SARSA');
      expect(detailed.type).toBe('on-policy');
      expect(detailed.stats).toBeDefined();
      expect(detailed.stats.steps).toBe(0);
    });

    it('should track convergence metrics', () => {
      const metrics = sarsa.getConvergenceMetrics();

      expect(metrics).toHaveProperty('isConverging');
      expect(metrics).toHaveProperty('convergenceRate');
      expect(metrics).toHaveProperty('stability');
      expect(metrics.convergenceRate).toBeGreaterThanOrEqual(0);
      expect(metrics.convergenceRate).toBeLessThanOrEqual(1);
    });
  });

  describe('State Management', () => {
    it('should export and import state correctly', () => {
      const state = createState();
      const action = createAction('test');

      // Train some Q-values
      for (let i = 0; i < 5; i++) {
        sarsa.update(createExperience(state, action, 1.0, createState()));
      }

      const exported = sarsa.export();
      const qValueBefore = sarsa.getQValue(state, action);
      const stepsBefore = sarsa.getStepCount();

      // Create new instance and import
      const newSarsa = new SARSALearner();
      newSarsa.import(exported);

      expect(newSarsa.getQValue(state, action)).toBeCloseTo(qValueBefore, 5);
      expect(newSarsa.getStepCount()).toBe(stepsBefore);
      expect(newSarsa.getTableSize()).toBe(sarsa.getTableSize());
    });

    it('should reset to initial state', () => {
      const state = createState();
      const action = createAction('test');

      // Train some Q-values
      sarsa.update(createExperience(state, action, 1.0, createState()));
      sarsa.endEpisode();

      expect(sarsa.getStepCount()).toBeGreaterThan(0);
      expect(sarsa.getTableSize()).toBeGreaterThan(0);

      // Reset
      sarsa.reset();

      expect(sarsa.getStepCount()).toBe(0);
      expect(sarsa.getEpisodeCount()).toBe(0);
      expect(sarsa.getTableSize()).toBe(0);
      expect(sarsa.getQValue(state, action)).toBe(0);
    });
  });

  describe('Convergence Test', () => {
    it('should converge on simple gridworld-like task within 100 episodes', () => {
      // Simple task: always positive reward for 'optimal' strategy
      const optimalAction = createAction('optimal');
      const suboptimalAction = createAction('suboptimal');

      const sarsa = new SARSALearner({
        learningRate: 0.1,
        discountFactor: 0.9,
        explorationRate: 0.3,
        explorationDecay: 0.98,
        minExplorationRate: 0.01
      });

      // Train for 100 episodes
      for (let episode = 0; episode < 100; episode++) {
        const state = createState();

        // Simulate episode: optimal action always gives +1, suboptimal gives -0.5
        for (let step = 0; step < 10; step++) {
          const useOptimal = Math.random() < 0.7; // Biased toward optimal
          const action = useOptimal ? optimalAction : suboptimalAction;
          const reward = useOptimal ? 1.0 : -0.5;
          const nextState = createState(0.5 + Math.random() * 0.1);

          sarsa.update(createExperience(state, action, reward, nextState));
        }

        sarsa.endEpisode();
      }

      // After 100 episodes, optimal action should have higher Q-value
      const testState = createState();
      const qOptimal = sarsa.getQValue(testState, optimalAction);
      const qSuboptimal = sarsa.getQValue(testState, suboptimalAction);

      expect(qOptimal).toBeGreaterThan(qSuboptimal);
      expect(qOptimal).toBeGreaterThan(0);

      // Exploration rate should have decayed
      expect(sarsa.getExplorationRate()).toBeLessThan(0.3);
    });
  });

  describe('Performance Comparison with Q-Learning Baseline', () => {
    it('should perform within 5% of Q-Learning on deterministic task', () => {
      // This test ensures SARSA is competitive with Q-Learning
      // Note: In practice, this would need actual Q-Learning comparison
      // For now, we test that SARSA learns effectively

      const state = createState();
      const goodAction = createAction('good');
      const badAction = createAction('bad');

      // Train SARSA
      for (let i = 0; i < 50; i++) {
        sarsa.update(createExperience(state, goodAction, 1.0, createState()));
        sarsa.update(createExperience(state, badAction, 0.0, createState()));
      }

      const qGood = sarsa.getQValue(state, goodAction);
      const qBad = sarsa.getQValue(state, badAction);

      // SARSA should learn the difference
      expect(qGood).toBeGreaterThan(qBad);
      expect(qGood).toBeGreaterThan(0.5); // Should be reasonably high
    });
  });

  describe('Memory Usage', () => {
    it('should report memory usage', () => {
      const state = createState();
      const action = createAction('test');

      // Train some Q-values
      for (let i = 0; i < 10; i++) {
        sarsa.update(createExperience(state, action, 1.0, createState()));
      }

      const memoryUsage = sarsa.getMemoryUsage();
      expect(memoryUsage).toBeGreaterThan(0);
    });
  });

  describe('Edge Cases', () => {
    it('should handle very large rewards', () => {
      const state = createState();
      const action = createAction('test');
      const experience = createExperience(state, action, 1000.0, createState());

      expect(() => sarsa.update(experience)).not.toThrow();
      expect(sarsa.getQValue(state, action)).toBeLessThan(1000); // Should be bounded by learning rate
    });

    it('should handle very negative rewards', () => {
      const state = createState();
      const action = createAction('test');
      const experience = createExperience(state, action, -1000.0, createState());

      expect(() => sarsa.update(experience)).not.toThrow();
      expect(sarsa.getQValue(state, action)).toBeGreaterThan(-1000); // Should be bounded by learning rate
    });

    it('should handle zero rewards consistently', () => {
      const state = createState();
      const action = createAction('test');

      for (let i = 0; i < 10; i++) {
        sarsa.update(createExperience(state, action, 0, createState()));
      }

      // Q-value should stay near 0 with zero rewards
      expect(Math.abs(sarsa.getQValue(state, action))).toBeLessThan(0.1);
    });
  });
});