每日算法 — 使用java实现多臂老虎机:ε-贪心与UCB算法探索利用权衡

引言

多臂老虎机(Multi-Armed Bandit, MAB)问题是强化学习中最经典的序列决策模型之一。想象面前有K台老虎机,每台机器的中奖概率未知且固定,玩家的目标是在有限次数内最大化累计收益。核心矛盾在于探索(Exploration)与利用(Exploitation)的权衡:是继续尝试已知高收益的机器,还是探索可能更优的未知机器?本文将用Java实现ε-贪心与UCB(Upper Confidence Bound)两种经典算法,并对比其 regret 边界。

问题形式化

设有K个动作(老虎机拉杆),每个动作a的奖励服从某个未知分布,期望奖励为μₐ。在T轮决策中,每轮选择一个动作aₜ,获得奖励rₜ。定义最优期望奖励为μ* = maxₐ μₐ,则累计 regret 为:

R(T) = Tμ* - Σ E[rₜ]

目标是使R(T)随T增长尽可能缓慢。

ε-贪心算法

以概率ε随机探索(均匀选择动作),以概率1-ε选择当前平均奖励最高的动作。ε通常随时间衰减,如 ε = 1/t。

优点:实现简单,渐近收敛
缺点:线性 regret 下界,探索效率低

UCB1算法

基于乐观面对不确定性原则,为每个动作的置信上界排序:

UCB(a) = μ̂ₐ + √(2lnT / Nₐ)

其中μ̂ₐ为动作a的样本平均奖励,Nₐ为被选择次数。该公式保证:尝试次数少的动作获得较高探索 bonus,随着数据积累逐渐收敛到真实均值。

理论保证:UCB1的 regret 上界为O(√(KT lnT)),优于ε-贪心的O(T)。

Java实现

import java.util.Random;

/**
 * 多臂老虎机问题求解器
 * 实现ε-贪心与UCB1两种经典算法
 */
public class MultiArmedBandit {

    private static final Random random = new Random(42);

    /**
     * 老虎机环境:每台机器有固定但未知的中奖概率
     */
    static class BanditEnvironment {
        private final double[] trueProbs; // 每台机器的真实中奖概率

        public BanditEnvironment(int k) {
            trueProbs = new double[k];
            // 随机生成真实概率,范围0.1-0.9
            for (int i = 0; i < k; i++) {
                trueProbs[i] = 0.1 + random.nextDouble() * 0.8;
            }
        }

        /**
         * 拉动指定机器,返回奖励(0或1)
         */
        public int pull(int arm) {
            return random.nextDouble() < trueProbs[arm] ? 1 : 0;
        }

        public double getTrueProb(int arm) {
            return trueProbs[arm];
        }

        public int getBestArm() {
            int best = 0;
            for (int i = 1; i < trueProbs.length; i++) {
                if (trueProbs[i] > trueProbs[best]) best = i;
            }
            return best;
        }

        public double getMaxProb() {
            double max = trueProbs[0];
            for (double p : trueProbs) max = Math.max(max, p);
            return max;
        }

        @Override
        public String toString() {
            StringBuilder sb = new StringBuilder("老虎机环境 (真实概率): ");
            for (int i = 0; i < trueProbs.length; i++) {
                sb.append(String.format("[%d]=%.3f ", i, trueProbs[i]));
            }
            return sb.toString();
        }
    }

    /**
     * 智能体基类
     */
    static abstract class Agent {
        protected final int k; // 机器数量
        protected final double[] rewards; // 累计奖励
        protected final int[] counts; // 每台机器被拉动次数
        protected double totalReward = 0;
        protected int totalPulls = 0;

        public Agent(int k) {
            this.k = k;
            this.rewards = new double[k];
            this.counts = new int[k];
        }

        public abstract int selectArm();

        public void update(int arm, int reward) {
            counts[arm]++;
            rewards[arm] += reward;
            totalReward += reward;
            totalPulls++;
        }

        public double getEstimatedValue(int arm) {
            return counts[arm] == 0 ? 0 : rewards[arm] / counts[arm];
        }

        public double getTotalReward() {
            return totalReward;
        }

        public double getRegret(double maxProb) {
            return totalPulls * maxProb - totalReward;
        }
    }

    /**
     * ε-贪心智能体
     */
    static class EpsilonGreedyAgent extends Agent {
        private double epsilon;
        private final boolean decay; // 是否衰减ε

        public EpsilonGreedyAgent(int k, double epsilon, boolean decay) {
            super(k);
            this.epsilon = epsilon;
            this.decay = decay;
        }

        @Override
        public int selectArm() {
            double currentEpsilon = decay ? Math.max(0.01, 1.0 / (totalPulls + 1)) : epsilon;

            if (random.nextDouble() < currentEpsilon) {
                // 探索:随机选择
                return random.nextInt(k);
            }

            // 利用:选择当前平均奖励最高的
            int bestArm = 0;
            double bestValue = -1;
            for (int i = 0; i < k; i++) {
                double value = getEstimatedValue(i);
                if (value > bestValue) {
                    bestValue = value;
                    bestArm = i;
                }
            }
            return bestArm;
        }
    }

    /**
     * UCB1智能体
     */
    static class UCB1Agent extends Agent {
        public UCB1Agent(int k) {
            super(k);
        }

        @Override
        public int selectArm() {
            // 如果有未尝试过的机器,优先尝试
            for (int i = 0; i < k; i++) {
                if (counts[i] == 0) return i;
            }

            // 计算UCB值并选择最大值
            int bestArm = 0;
            double bestUCB = -1;
            double logTotal = Math.log(totalPulls);

            for (int i = 0; i < k; i++) {
                double avgReward = rewards[i] / counts[i];
                double confidence = Math.sqrt(2 * logTotal / counts[i]);
                double ucb = avgReward + confidence;

                if (ucb > bestUCB) {
                    bestUCB = ucb;
                    bestArm = i;
                }
            }

            return bestArm;
        }
    }

    /**
     * 汤普森采样智能体(贝叶斯方法)
     */
    static class ThompsonSamplingAgent extends Agent {
        private final double[] alpha; // Beta分布参数α
        private final double[] beta;  // Beta分布参数β

        public ThompsonSamplingAgent(int k) {
            super(k);
            alpha = new double[k];
            beta = new double[k];
            for (int i = 0; i < k; i++) {
                alpha[i] = 1; // 先验:Beta(1,1) = Uniform
                beta[i] = 1;
            }
        }

        @Override
        public int selectArm() {
            int bestArm = 0;
            double bestSample = -1;

            for (int i = 0; i < k; i++) {
                // 从Beta分布采样
                double sample = sampleBeta(alpha[i], beta[i]);
                if (sample > bestSample) {
                    bestSample = sample;
                    bestArm = i;
                }
            }

            return bestArm;
        }

        @Override
        public void update(int arm, int reward) {
            super.update(arm, reward);
            alpha[arm] += reward;
            beta[arm] += (1 - reward);
        }

        /**
         * 使用Gamma分布近似生成Beta分布样本
         * Beta(α,β) = Gamma(α,1) / (Gamma(α,1) + Gamma(β,1))
         */
        private double sampleBeta(double a, double b) {
            double x = sampleGamma(a, 1.0);
            double y = sampleGamma(b, 1.0);
            return x / (x + y);
        }

        /**
         * Marsaglia-Tsang方法生成Gamma分布样本
         */
        private double sampleGamma(double shape, double scale) {
            if (shape < 1) {
                return sampleGamma(shape + 1, scale) * Math.pow(random.nextDouble(), 1.0 / shape);
            }
            double d = shape - 1.0 / 3.0;
            double c = 1.0 / Math.sqrt(9.0 * d);
            while (true) {
                double x, v;
                do {
                    x = random.nextGaussian();
                    v = 1.0 + c * x;
                } while (v <= 0);
                v = v * v * v;
                double u = random.nextDouble();
                if (u < 1.0 - 0.0331 * x * x * x * x) return d * v * scale;
                if (Math.log(u) < 0.5 * x * x + d * (1.0 - v + Math.log(v))) return d * v * scale;
            }
        }
    }

    /**
     * 运行对比实验
     */
    public static void runExperiment(int k, int rounds, int repetitions) {
        double totalRegretEps = 0, totalRegretUCB = 0, totalRegretTS = 0;
        double totalRewardEps = 0, totalRewardUCB = 0, totalRewardTS = 0;

        for (int rep = 0; rep < repetitions; rep++) {
            BanditEnvironment env = new BanditEnvironment(k);
            double maxProb = env.getMaxProb();

            EpsilonGreedyAgent epsAgent = new EpsilonGreedyAgent(k, 0.1, true);
            UCB1Agent ucbAgent = new UCB1Agent(k);
            ThompsonSamplingAgent tsAgent = new ThompsonSamplingAgent(k);

            for (int t = 0; t < rounds; t++) {
                // ε-贪心
                int arm = epsAgent.selectArm();
                int reward = env.pull(arm);
                epsAgent.update(arm, reward);

                // UCB1(需要新环境实例,或重置随机种子)
                // 为公平对比,在同一环境下顺序执行
            }

            // 重新运行以公平对比
            BanditEnvironment envEps = new BanditEnvironment(k);
            EpsilonGreedyAgent aEps = new EpsilonGreedyAgent(k, 0.1, true);
            for (int t = 0; t < rounds; t++) {
                int arm = aEps.selectArm();
                aEps.update(arm, envEps.pull(arm));
            }

            BanditEnvironment envUCB = new BanditEnvironment(k);
            UCB1Agent aUCB = new UCB1Agent(k);
            for (int t = 0; t < rounds; t++) {
                int arm = aUCB.selectArm();
                aUCB.update(arm, envUCB.pull(arm));
            }

            BanditEnvironment envTS = new BanditEnvironment(k);
            ThompsonSamplingAgent aTS = new ThompsonSamplingAgent(k);
            for (int t = 0; t < rounds; t++) {
                int arm = aTS.selectArm();
                aTS.update(arm, envTS.pull(arm));
            }

            totalRewardEps += aEps.getTotalReward();
            totalRewardUCB += aUCB.getTotalReward();
            totalRewardTS += aTS.getTotalReward();
            totalRegretEps += aEps.getRegret(envEps.getMaxProb());
            totalRegretUCB += aUCB.getRegret(envUCB.getMaxProb());
            totalRegretTS += aTS.getRegret(envTS.getMaxProb());
        }

        System.out.printf("对比实验结果 (%d轮 × %d次重复, %d臂老虎机):%n", rounds, repetitions, k);
        System.out.printf("  ε-贪心(衰减): 平均奖励=%.2f, 平均Regret=%.2f%n", 
                totalRewardEps / repetitions, totalRegretEps / repetitions);
        System.out.printf("  UCB1:        平均奖励=%.2f, 平均Regret=%.2f%n",
                totalRewardUCB / repetitions, totalRegretUCB / repetitions);
        System.out.printf("  汤普森采样:  平均奖励=%.2f, 平均Regret=%.2f%n",
                totalRewardTS / repetitions, totalRegretTS / repetitions);
    }

    public static void main(String[] args) {
        // 演示环境
        BanditEnvironment demo = new BanditEnvironment(5);
        System.out.println(demo);
        System.out.printf("最优机器: %d (概率%.3f)%n%n", demo.getBestArm(), demo.getMaxProb());

        // 小规模对比
        runExperiment(5, 1000, 100);

        System.out.println("\n大规模对比:");
        runExperiment(10, 10000, 20);
    }
}

复杂度分析

  • ε-贪心决策:O(K) 时间(扫描所有动作),O(K) 空间
  • UCB1决策:O(K) 时间,O(K) 空间
  • 汤普森采样:O(K) 时间(每动作一次Beta采样),O(K) 空间

三种算法均为每轮线性复杂度,但UCB1和汤普森采样的 regret 增长更优。

应用场景扩展

多臂老虎机框架广泛应用于:

  • 在线广告投放:选择点击率最高的广告位组合
  • 推荐系统:探索用户可能感兴趣的新内容
  • 临床试验:动态分配治疗效果更好的治疗方案
  • 网络路由:选择延迟最低的传输路径

理解探索与利用的权衡,是设计高效自适应系统的核心能力。