Refactor exploration to use new strategy approach and deprecate prefix appraoch

Author: alatimer

experiments/alaro/trap_bandit/run.py

@@ -188,9 +188,7 @@ def make_agent(
         prior=np.array([1.0]),
         reward_function=REWARD_FUNCTION,
         policy_discretisation=0,
-        exploration_prefix=0,
         exploration_strategy=strategy,
-        epsilon=0.0,
     )
 
 

experiments/fllor2/ib_validate_classical.ipynb

@@ -22,6 +22,7 @@
     "from ibrl.utils import construct_environment,dump_array,sample_action\n",
     "from ibrl.environments import BernoulliBanditEnvironment\n",
     "from ibrl.agents import InfraBayesianAgent, DiscreteBayesianAgent\n",
+    "from ibrl.exploration import EpsilonGreedy\n",
     "from ibrl.infrabayesian import AMeasure,Infradistribution,MultiBernoulliWorldModel"
    ]
   },
@@ -58,7 +59,8 @@
     "    # per-arm Bayesian inference — no joint enumeration needed.\n",
     "    params = wm.make_params([grid] * num_actions)\n",
     "    hypotheses = [Infradistribution([AMeasure(params)], world_model=wm)]\n",
-    "    return InfraBayesianAgent(num_actions=num_actions, hypotheses=hypotheses, exploration_prefix=None, **kwargs)"
+    "    epsilon = kwargs.pop(\"epsilon\", 0.1)\n",
+    "    return InfraBayesianAgent(num_actions=num_actions, hypotheses=hypotheses, exploration_strategy=EpsilonGreedy(epsilon), **kwargs)"
    ]
   },
   {

experiments/fllor2/ibtest.py

@@ -21,7 +21,7 @@
     "seed":        42,
     "verbose":     2,
 }
-shared = dict(num_actions=n, seed=options["seed"] + 0x01234567, verbose=options["verbose"], epsilon=0.1)
+shared = dict(num_actions=n, seed=options["seed"] + 0x01234567, verbose=options["verbose"])
 
 
 def make_bernoulli(alpha, beta):

experiments/fllor2/newcomb.ipynb

@@ -40,7 +40,6 @@
     "        dist = Infradistribution([m],wm)\n",
     "        agent = InfraBayesianAgent(\n",
     "            num_actions=2,\n",
-    "            epsilon=0.,\n",
     "            hypotheses=[dist],\n",
     "            reward_function=wm.agent_reward_matrix(),\n",
     "            policy_discretisation=5,\n",
@@ -161,7 +160,6 @@
     "    ], np.ones(6)/6)\n",
     "    agent = InfraBayesianAgent(\n",
     "        num_actions=2,\n",
-    "        epsilon=0.,\n",
     "        hypotheses=[dist],\n",
     "        reward_function=wm.agent_reward_matrix(),\n",
     "        policy_discretisation=5)\n",

experiments/fllor2/validate_classical.py

@@ -14,6 +14,7 @@
 
 from ibrl.agents import DiscreteBayesianAgent, InfraBayesianAgent
 from ibrl.environments import BernoulliBanditEnvironment
+from ibrl.exploration import EpsilonGreedy
 from ibrl.infrabayesian import AMeasure, Infradistribution, MultiBernoulliWorldModel
 from ibrl.simulators import simulate
 
@@ -32,10 +33,11 @@ def make_classical_ib_agent(
     grid = [np.array([1.0 - p, p]) for p in np.linspace(0.0, 1.0, num_hypotheses)]
     params = wm.make_params([grid] * num_actions)
     hypothesis = Infradistribution([AMeasure(params)], world_model=wm)
+    epsilon = kwargs.pop("epsilon", 0.1)
     return InfraBayesianAgent(
         num_actions=num_actions,
         hypotheses=[hypothesis],
-        exploration_prefix=None,
+        exploration_strategy=EpsilonGreedy(epsilon),
         **kwargs,
     )
 

ibrl/agents/base_greedy.py

@@ -1,30 +1,36 @@
 import numpy as np
 
 from . import BaseAgent
+from ..exploration import EpsilonGreedy, ExplorationStrategy, Softmax
 
 
 class BaseGreedyAgent(BaseAgent):
     """
-    Base class for agents that use either an epsilon-greedy or softmax policy to encourage exploration.
+    Base class for agents that use an exploration strategy over action values.
 
     Arguments:
+        exploration_strategy: Optional strategy object for converting values to a policy
         epsilon:        For epsilon-greedy policy
         temperature:    For softmax policy
         decay_type:     Select formula for decreasing rate (0: exponential, 1: linear)
 
     Both epsilon and temperature can be either a fixed float or a tuple (start,decay constant,min) for decreasing exploration.
     """
     def __init__(self, *,
+            exploration_strategy : ExplorationStrategy | None = None,
             epsilon     : float | tuple[float] | None = None,
             temperature : float | tuple[float] | None = None,
             decay_type  : float = 0,
             **kwargs):
         super().__init__(**kwargs)
 
-        if epsilon is not None and temperature is not None:
-            raise RuntimeError("Cannot specify both epsilon and temperature")
-        if epsilon is None and temperature is None:
-            epsilon = 0.1  # default value
+        if exploration_strategy is not None and (epsilon is not None or temperature is not None):
+            raise RuntimeError("Cannot specify exploration_strategy with epsilon or temperature")
+        if exploration_strategy is None:
+            if epsilon is not None and temperature is not None:
+                raise RuntimeError("Cannot specify both epsilon and temperature")
+            if epsilon is None and temperature is None:
+                epsilon = 0.1  # default value
 
         assert epsilon is None or isinstance(epsilon,float) or (isinstance(epsilon,tuple) and len(epsilon)==3)
         assert temperature is None or isinstance(temperature,float) or (isinstance(temperature,tuple) and len(temperature)==3)
@@ -33,34 +39,20 @@ def __init__(self, *,
         self.temperature = temperature
         self.decay_type = int(decay_type)
         assert self.decay_type in [0,1]
+        if exploration_strategy is not None:
+            self.exploration_strategy = exploration_strategy
+        elif epsilon is not None:
+            self.exploration_strategy = EpsilonGreedy(epsilon, self.decay_type)
+        elif temperature is not None:
+            self.exploration_strategy = Softmax(temperature, self.decay_type)
+        else:
+            raise RuntimeError("Invalid state")
 
     def build_greedy_policy(self, values : np.ndarray) -> np.ndarray:
         """
-        Construct probabilities based on given reward estimates and selected policy
+        Construct probabilities from reward estimates using the configured strategy.
         """
-        if self.epsilon is not None:
-            return self.build_epsilon_greedy_policy(values)
-        if self.temperature is not None:
-            return self.build_softmax_policy(values)
-        raise RuntimeError("Invalid state")
-
-    def build_epsilon_greedy_policy(self, values : np.ndarray) -> np.ndarray:
-        # Exploitation: sample uniformly across actions with highest value
-        best_actions = np.isclose(values, values.max())
-        exploit = np.ones_like(values)*best_actions / best_actions.sum()
-
-        # Exploration: sample uniformly across all actions
-        explore = np.ones_like(values) / self.num_actions
-
-        epsilon = self.parse_parameter(self.epsilon)
-        return (1 - epsilon) * exploit + epsilon * explore
-
-    def build_softmax_policy(self, values : np.ndarray) -> np.ndarray:
-        temperature = self.parse_parameter(self.temperature)
-
-        # Numerically stable softmax
-        exp = np.exp((values - values.max()) / temperature)
-        return exp / exp.sum()
+        return self.exploration_strategy.get_probabilities(self, values)
 
     def parse_parameter(self, parameter : float | tuple[float]) -> float:
         """

ibrl/agents/infrabayesian.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 from . import BaseGreedyAgent
-from ..exploration import ExplorationStrategy
+from ..exploration import ExplorationStrategy, Greedy
 from ..infrabayesian.a_measure import AMeasure
 from ..infrabayesian.infradistribution import Infradistribution
 
@@ -21,20 +21,17 @@ class InfraBayesianAgent(BaseGreedyAgent):
         prior:            distribution over hypotheses; default: uniform
         reward_function:  reward_function[a,o] is reward upon seeing outcome o from action a
         policy_discretisation:  number of mixed policies to consider per action; default: 0 (i.e. only pure policies)
-        exploration_prefix:     parameter to control exploration
-                                =0    no exploration
-                                >0    forced exploration prefix for given number of steps, then no exploration
-                                None  greedy exploration (epsilon or softmax; breaks regret bounds)
+        exploration_strategy:  optional strategy object for Bayesian single-measure exploration
     """
     def __init__(self, *,
             hypotheses : list[Infradistribution],
             prior : np.ndarray | None = None,           # shape (len(hypotheses),)
             reward_function : np.ndarray | None = None, # shape (num_actions, num_outcomes)
             policy_discretisation : int = 0,
-            exploration_prefix : int | None = 0,
             exploration_strategy : ExplorationStrategy | None = None,
             **kwargs):
-        super().__init__(**kwargs)
+        base_strategy = exploration_strategy if exploration_strategy is not None else Greedy()
+        super().__init__(exploration_strategy=base_strategy, **kwargs)
         assert len(hypotheses) > 0
         assert all(isinstance(h.world_model, type(hypotheses[0].world_model))
                    for h in hypotheses), "All hypotheses must share the same WorldModel type"
@@ -43,7 +40,6 @@ def __init__(self, *,
         # default: reward_function[a,o] = o with o ∈ {0,1}
         self.reward_function = (reward_function if reward_function is not None
                                 else np.linspace(np.zeros(self.num_actions),np.ones(self.num_actions),2).T)
-        self.exploration_prefix = exploration_prefix
         self.exploration_strategy = exploration_strategy
 
         # Discretise policy space:
@@ -75,15 +71,6 @@ def get_probabilities(self) -> np.ndarray:
                 )
             return self.exploration_strategy.get_probabilities(self, self._action_values())
 
-        # Greedy policy: reproduces classical agent, breaks regret bounds
-        if self.exploration_prefix is None:
-            return self.build_greedy_policy(self._action_values())
-
-        # Forced exploration prefix: regret bounds asymptotically preserved
-        if self.step <= self.exploration_prefix:
-            return np.ones(self.num_actions) / self.num_actions
-
-        # Use optimal policy: no exploration, almost no learning
         return self._optimal_policy()
 
     def dump_state(self) -> str:

ibrl/exploration/__init__.py

@@ -6,6 +6,7 @@
     ExplorationStrategy,
     Greedy,
     HypothesisThompsonSampling,
+    Softmax,
     ThompsonSampling,
     UniformPrefixThen,
 )
@@ -16,6 +17,7 @@
     "ExplorationStrategy",
     "Greedy",
     "HypothesisThompsonSampling",
+    "Softmax",
     "ThompsonSampling",
     "UniformPrefixThen",
 ]

ibrl/exploration/strategies.py

@@ -22,20 +22,40 @@ def get_probabilities(self, agent, values: np.ndarray) -> np.ndarray:
 class EpsilonGreedy(ExplorationStrategy):
     """Epsilon-greedy exploration with a fixed or scheduled epsilon."""
 
-    def __init__(self, epsilon: float | tuple[float, float, float]):
+    def __init__(self, epsilon: float | tuple[float, float, float], decay_type: int = 0):
         self.epsilon = epsilon
+        self.decay_type = int(decay_type)
 
     def get_probabilities(self, agent, values: np.ndarray) -> np.ndarray:
         greedy = Greedy().get_probabilities(agent, values)
-        eps = self._epsilon(agent)
+        eps = _scheduled_value(self.epsilon, agent.step, self.decay_type)
         uniform = np.ones(agent.num_actions) / agent.num_actions
         return (1 - eps) * greedy + eps * uniform
 
-    def _epsilon(self, agent) -> float:
-        if isinstance(self.epsilon, tuple):
-            start, rate, end = self.epsilon
-            return max(start / (agent.step ** rate), end)
-        return float(self.epsilon)
+
+class Softmax(ExplorationStrategy):
+    """Softmax action selection with a fixed or scheduled temperature."""
+
+    def __init__(self, temperature: float | tuple[float, float, float], decay_type: int = 0):
+        self.temperature = temperature
+        self.decay_type = int(decay_type)
+
+    def get_probabilities(self, agent, values: np.ndarray) -> np.ndarray:
+        temperature = _scheduled_value(self.temperature, agent.step, self.decay_type)
+        exp = np.exp((values - values.max()) / temperature)
+        return exp / exp.sum()
+
+
+def _scheduled_value(parameter: float | tuple[float, float, float], step: int, decay_type: int) -> float:
+    if isinstance(parameter, float):
+        return parameter
+    if decay_type == 0:
+        start, rate, end = parameter
+        return max(start / (step ** rate), end)
+    if decay_type == 1:
+        start, last_step, end = parameter
+        return end if step >= last_step else (start + (end - start) * (step / last_step))
+    raise RuntimeError("Invalid decay_type")
 
 
 class UniformPrefixThen(ExplorationStrategy):

ibrl/infrabayesian/world_model.py

@@ -76,3 +76,15 @@ def compute_expected_reward(self, belief_state, reward_function: np.ndarray,
         Returns a scalar.
         """
         pass
+
+    def get_posterior_component_weights(self, belief_state, params) -> np.ndarray:
+        """Posterior weights for finite component-mixture exploration strategies."""
+        raise NotImplementedError(
+            f"{type(self).__name__} does not expose posterior component weights"
+        )
+
+    def get_component_expected_rewards(self, component, reward_function: np.ndarray) -> np.ndarray:
+        """Per-action expected rewards for one finite posterior component."""
+        raise NotImplementedError(
+            f"{type(self).__name__} does not expose component expected rewards"
+        )