每日算法 — 使用java实现UCB算法:多臂老虎机与探索利用最优决策

算法概述

多臂老虎机(Multi-Armed Bandit)问题是强化学习中最基础也最重要的序贯决策问题。想象你面前有K台老虎机,每台机器的真实收益率未知,你的目标是在有限次数内最大化总收益。这需要在”探索”(尝试未知机器以获取信息)和”利用”(选择当前估计收益最高的机器)之间找到最佳平衡。

UCB(Upper Confidence Bound)算法通过为每台机器的估计收益添加上界置信区间来优雅地解决这一问题。其核心思想是:尝试次数越少的机器,其真实收益的不确定性越大,因此应给予更多的探索机会。

UCB1公式

UCB_i = meanReward_i + sqrt(2 * ln(totalPulls) / pullCount_i)

第一项是利用项(平均收益),第二项是探索项(置信区间上界)。

完整Java实现

import java.util.*;

/**
 * 多臂老虎机问题与UCB算法实现
 * 包含UCB1、UCB1-Tuned和Epsilon-Greedy对比
 */
public class MultiArmedBandit {

    interface SlotMachine {
        double pull();
        double getTrueMean();
        String getName();
    }

    static class GaussianMachine implements SlotMachine {
        String name;
        double mean;
        double stdDev;
        Random random;

        GaussianMachine(String name, double mean, double stdDev, long seed) {
            this.name = name;
            this.mean = mean;
            this.stdDev = stdDev;
            this.random = new Random(seed);
        }

        @Override
        public double pull() {
            return mean + random.nextGaussian() * stdDev;
        }

        @Override
        public double getTrueMean() {
            return mean;
        }

        @Override
        public String getName() {
            return name;
        }
    }

    static class UCB1Agent {
        int numMachines;
        double[] totalRewards;
        int[] pullCounts;
        int totalPulls;
        String name;

        UCB1Agent(int numMachines, String name) {
            this.numMachines = numMachines;
            this.name = name;
            this.totalRewards = new double[numMachines];
            this.pullCounts = new int[numMachines];
            this.totalPulls = 0;
        }

        int selectMachine() {
            for (int i = 0; i < numMachines; i++) {
                if (pullCounts[i] == 0) {
                    return i;
                }
            }
            double bestUCB = Double.NEGATIVE_INFINITY;
            int bestMachine = 0;
            for (int i = 0; i < numMachines; i++) {
                double meanReward = totalRewards[i] / pullCounts[i];
                double explorationTerm = Math.sqrt(
                    2 * Math.log(totalPulls) / pullCounts[i]
                );
                double ucb = meanReward + explorationTerm;
                if (ucb > bestUCB) {
                    bestUCB = ucb;
                    bestMachine = i;
                }
            }
            return bestMachine;
        }

        void update(int machine, double reward) {
            totalRewards[machine] += reward;
            pullCounts[machine]++;
            totalPulls++;
        }
    }

    static class UCB1TunedAgent {
        int numMachines;
        double[] totalRewards;
        double[] squaredRewards;
        int[] pullCounts;
        int totalPulls;

        UCB1TunedAgent(int numMachines) {
            this.numMachines = numMachines;
            this.totalRewards = new double[numMachines];
            this.squaredRewards = new double[numMachines];
            this.pullCounts = new int[numMachines];
        }

        int selectMachine() {
            for (int i = 0; i < numMachines; i++) {
                if (pullCounts[i] == 0) return i;
            }
            double bestUCB = Double.NEGATIVE_INFINITY;
            int bestMachine = 0;
            for (int i = 0; i < numMachines; i++) {
                double mean = totalRewards[i] / pullCounts[i];
                double meanSq = squaredRewards[i] / pullCounts[i];
                double variance = meanSq - mean * mean;
                double V = variance + Math.sqrt(2 * Math.log(totalPulls) / pullCounts[i]);
                double explorationTerm = Math.sqrt(
                    Math.log(totalPulls) / pullCounts[i] * Math.min(0.25, V)
                );
                double ucb = mean + explorationTerm;
                if (ucb > bestUCB) {
                    bestUCB = ucb;
                    bestMachine = i;
                }
            }
            return bestMachine;
        }

        void update(int machine, double reward) {
            totalRewards[machine] += reward;
            squaredRewards[machine] += reward * reward;
            pullCounts[machine]++;
            totalPulls++;
        }
    }

    static class EpsilonGreedyAgent {
        int numMachines;
        double[] totalRewards;
        int[] pullCounts;
        double epsilon;
        Random random;

        EpsilonGreedyAgent(int numMachines, double epsilon) {
            this.numMachines = numMachines;
            this.epsilon = epsilon;
            this.totalRewards = new double[numMachines];
            this.pullCounts = new int[numMachines];
            this.random = new Random();
        }

        int selectMachine() {
            if (random.nextDouble() < epsilon) {
                return random.nextInt(numMachines);
            }
            double bestMean = Double.NEGATIVE_INFINITY;
            int bestMachine = 0;
            for (int i = 0; i < numMachines; i++) {
                double mean = pullCounts[i] == 0 ? 0 : totalRewards[i] / pullCounts[i];
                if (mean > bestMean) {
                    bestMean = mean;
                    bestMachine = i;
                }
            }
            return bestMachine;
        }

        void update(int machine, double reward) {
            totalRewards[machine] += reward;
            pullCounts[machine]++;
        }
    }

    public static void runExperiment(int numRounds) {
        SlotMachine[] machines = new SlotMachine[5];
        machines[0] = new GaussianMachine("机器A", 0.3, 1.0, 1);
        machines[1] = new GaussianMachine("机器B", 0.5, 1.0, 2);
        machines[2] = new GaussianMachine("机器C", 0.7, 1.0, 3);
        machines[3] = new GaussianMachine("机器D", 0.4, 1.0, 4);
        machines[4] = new GaussianMachine("机器E", 0.2, 1.0, 5);

        double optimalMean = Arrays.stream(machines)
            .mapToDouble(SlotMachine::getTrueMean)
            .max().orElse(0);

        UCB1Agent ucb1 = new UCB1Agent(5, "UCB1");
        UCB1TunedAgent ucbTuned = new UCB1TunedAgent(5);
        EpsilonGreedyAgent egreedy = new EpsilonGreedyAgent(5, 0.1);

        double ucb1Cum = 0, tunedCum = 0, greedyCum = 0;
        double ucb1Regret = 0, tunedRegret = 0, greedyRegret = 0;

        System.out.println("轮次\tUCB1收益\tUCB1遗憾\tTuned收益\tTuned遗憾\tE-Greedy收益\tE-Greedy遗憾");

        for (int round = 0; round < numRounds; round++) {
            int m1 = ucb1.selectMachine();
            double r1 = machines[m1].pull();
            ucb1.update(m1, r1);
            ucb1Cum += r1;
            ucb1Regret += (optimalMean - machines[m1].getTrueMean());

            int m2 = ucbTuned.selectMachine();
            double r2 = machines[m2].pull();
            ucbTuned.update(m2, r2);
            tunedCum += r2;
            tunedRegret += (optimalMean - machines[m2].getTrueMean());

            int m3 = egreedy.selectMachine();
            double r3 = machines[m3].pull();
            egreedy.update(m3, r3);
            greedyCum += r3;
            greedyRegret += (optimalMean - machines[m3].getTrueMean());

            if (round % 200 == 0 || round == numRounds - 1) {
                System.out.printf("%d\t%.2f\t\t%.2f\t\t%.2f\t\t%.2f\t\t%.2f\t\t%.2f%n",
                    round, ucb1Cum, ucb1Regret, tunedCum, tunedRegret, greedyCum, greedyRegret);
            }
        }

        System.out.println("\n=== 最终统计 ===");
        System.out.println("算法\t\t总收益\t\t总遗憾");
        System.out.printf("UCB1\t\t%.2f\t\t%.2f%n", ucb1Cum, ucb1Regret);
        System.out.printf("UCB1-Tuned\t%.2f\t\t%.2f%n", tunedCum, tunedRegret);
        System.out.printf("Epsilon-Greedy\t%.2f\t\t%.2f%n", greedyCum, greedyRegret);

        System.out.println("\nUCB1各机器选择次数:");
        for (int i = 0; i < 5; i++) {
            System.out.printf("%s(真实收益%.1f): %d次%n",
                machines[i].getName(), machines[i].getTrueMean(), ucb1.pullCounts[i]);
        }
    }

    public static void main(String[] args) {
        System.out.println("多臂老虎机UCB算法实验\n");
        runExperiment(2000);
    }
}

理论保证

UCB1算法的累积遗憾(Regret)上界为O(sqrt(KT log T)),其中K为机器数量,T为总轮数。这意味着随着轮数增加,UCB1的平均收益会收敛到最优策略。UCB1-Tuned进一步考虑了奖励方差,在实际中通常表现更好。

应用场景

UCB算法广泛应用于在线广告投放(选择最优广告位)、推荐系统(探索用户兴趣)、临床试验(选择最优治疗方案)、网络路由(选择最优路径)等需要在线决策的场景。