Merge pull request #7 from omereliy/MCTS

Mcts

Author: omereliy

.github/workflows/ci.yml

@@ -1,3 +1,4 @@
+# CI workflow - runs tests on pull requests to master
 name: CI
 
 on:

CHANGELOG.md

@@ -7,7 +7,21 @@ Versioning follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ---
 
-## [Unreleased]
+## [0.3.0] — 2026-03-15
+
+### Added
+- `BoundedLookaheadSelector` — depth-limited lookahead action selection with discounted future gain estimation
+- `IGMCTSSelector` — full UCT-based MCTS action selection (selection, expansion, greedy rollout, backpropagation)
+- `simulate_action()` — transition simulation using partial model (confirmed effects + stochastic applicability)
+- `selection_strategy='lookahead'` option with configurable depth, top_k, discount
+- `selection_strategy='mcts'` option with configurable iterations and rollout depth
+- Selection strategy and MCTS/lookahead CLI flags for AMLGym experiment script
+
+### Changed
+- Extract `_apply_injective_binding_filter()` to deduplicate binding filter logic in sequential/parallel gain computation
+
+### Notes
+- Full UCT MCTS (`selection_strategy='mcts'`) is significantly slower than other strategies due to SAT solving during rollouts. Performance improvements planned for a future release. Use `lookahead` for a balance between exploration depth and speed.
 
 ---
 
@@ -98,17 +112,26 @@ Planned features and improvements for upcoming releases.
 | Refactor | Extract literal conversion helper and deduplicate model counting | Done |
 | Quality | Fix mypy errors and modernize type annotations | Done |
 
-### 0.3.0-beta.N — MCTS Pre-releases
+### 0.3.0-beta.1 — MCTS Phase 1: Bounded Lookahead
+
+| Type | Description | Status |
+|------|-------------|--------|
+| Feature | `BoundedLookaheadSelector` — depth-limited lookahead action selection using discounted information gain | Done |
+| Feature | `simulate_action()` — transition simulation using partial model (confirmed effects + stochastic applicability) | Done |
+| Feature | `selection_strategy='lookahead'` option with configurable depth, top_k, discount | Done |
+
+### 0.3.0-beta.2 — MCTS Phase 2: Full UCT
 
 | Type | Description | Status |
 |------|-------------|--------|
-| Feature | MCTS-based action selection (unstable, iterating toward 0.3.0) | Planned |
+| Feature | `IGMCTSSelector` with UCT selection, expansion, greedy rollout, backpropagation | Done |
+| Feature | `selection_strategy='mcts'` with configurable iterations and rollout depth | Done |
 
-### 0.3.0 — MCTS
+### 0.3.0 — MCTS (Stable)
 
 | Type | Description | Status |
 |------|-------------|--------|
-| Feature | Full MCTS-based action selection (stable) | Planned |
+| Feature | Stable MCTS-based action selection with lookahead and full UCT variants | Done |
 
 ### 0.4.0 — Planned Action Sequences
 
@@ -129,7 +152,8 @@ Planned features and improvements for upcoming releases.
 
 ---
 
-[Unreleased]: https://github.com/omereliy/online_model_learning/compare/v0.2.1...HEAD
+[Unreleased]: https://github.com/omereliy/online_model_learning/compare/v0.3.0...HEAD
+[0.3.0]: https://github.com/omereliy/online_model_learning/compare/v0.2.1...v0.3.0
 [0.2.1]: https://github.com/omereliy/online_model_learning/compare/v0.2.0...v0.2.1
 [0.2.0]: https://github.com/omereliy/online_model_learning/compare/v0.1.1...v0.2.0
 [0.1.1]: https://github.com/omereliy/online_model_learning/compare/v0.1.0...v0.1.1

CLAUDE.md

@@ -38,7 +38,17 @@ Current: multi-module with Unified Planning. Target: 3-class design importable b
 - `docs/releasing.md` — PyPI release procedure
 - `.claude/rules/` — Simplification, migration, testing, experiment, releasing rules
 
+## Code Conventions
+- **Lazy imports for optional deps**: When adding integrations with external packages, import inside methods — not at module top level. Top-level imports break the AMLGym registry when the dep isn't installed.
+- **Wire through all entry points**: When adding CLI flags or features, check ALL scripts in `scripts/` and `experiments/` that use the same functionality. Don't wire a flag into only one script.
+
 ## Integration
 - Target: importable from AMLGym via `OnlineAlgorithmAdapter`
 - Currently: standalone with Unified Planning environment
 - Never write custom PDDL parsers — use library APIs
+
+## Prompting Tips
+When requesting changes, front-load constraints to avoid wrong approaches:
+- **Algorithm/SAT work**: "Reference `docs/information_gain_algorithm/INFORMATION_GAIN_ALGORITHM.md` for formulas. Relevant code is in [file]. Run pytest after changes."
+- **Experiment work**: "Check ALL scripts in `scripts/` and `experiments/` that invoke this. Start from the error trace, don't explore broadly."
+- **MCTS/optimization**: "Profile first before proposing changes. Run the experiment script to verify no regression."

README.md

@@ -34,6 +34,17 @@ learner.update_model()
 - **Lifted learning** -- learns at the operator level, generalizing across object instances
 - **Object subset selection** -- scales to large domains by focusing on relevant object subsets
 - **Parallel gain computation** -- optional multiprocessing for large action spaces
+- **MCTS-based action selection** -- lookahead and full UCT strategies for deeper exploration
+
+## Action Selection Strategies
+
+| Strategy | Description | Speed |
+|----------|-------------|-------|
+| `greedy` | Selects the action with the highest immediate information gain | Fast (default) |
+| `lookahead` | Bounded depth-limited lookahead with discounted future gain | Moderate |
+| `mcts` | Full UCT-based Monte Carlo Tree Search | Slow (see note below) |
+
+> **Performance note:** The `mcts` strategy performs SAT solving during rollouts, which makes it significantly slower than other strategies. For large domains it may be impractical. Performance improvements are planned for a future release. Use `lookahead` for a balance between exploration depth and speed.
 
 ## Configuration
 
@@ -46,6 +57,10 @@ learner = InformationGainLearner(
     spare_objects_per_type=2,             # extra objects per type beyond minimum
     num_workers=None,                      # parallel workers (None=auto, 0=sequential)
     learn_negative_preconditions=True,    # include negative precondition candidates
+    selection_strategy="greedy",          # "greedy", "lookahead", or "mcts"
+    lookahead_depth=2,                    # depth for lookahead strategy
+    mcts_iterations=50,                   # iterations for mcts strategy
+    mcts_rollout_depth=5,                # rollout depth for mcts strategy
 )
 ```
 

information_gain_aml/__init__.py

@@ -1,3 +1,3 @@
 """Information Gain-based Online Action Model Learning."""
 
-__version__ = "0.2.1"
+__version__ = "0.3.0"

information_gain_aml/algorithms/information_gain.py

@@ -126,10 +126,19 @@ def __init__(self,
         self._base_cnf_count_cache: Dict[str, int] = {}
 
         # Phase 3: Action selection strategy
-        self.selection_strategy = kwargs.get('selection_strategy', 'greedy')  # 'greedy', 'epsilon_greedy', 'boltzmann'
+        self.selection_strategy = kwargs.get('selection_strategy', 'greedy')  # 'greedy', 'epsilon_greedy', 'boltzmann', 'lookahead', 'mcts'
         self.epsilon = kwargs.get('epsilon', 0.1)  # For epsilon-greedy
         self.temperature = kwargs.get('temperature', 1.0)  # For Boltzmann
 
+        # Lookahead parameters (for selection_strategy='lookahead')
+        self.lookahead_depth = kwargs.get('lookahead_depth', 2)
+        self.lookahead_top_k = kwargs.get('lookahead_top_k', 5)
+        self.lookahead_discount = kwargs.get('lookahead_discount', 0.9)
+
+        # MCTS parameters (for selection_strategy='mcts')
+        self.mcts_iterations = kwargs.get('mcts_iterations', 50)
+        self.mcts_rollout_depth = kwargs.get('mcts_rollout_depth', 5)
+
         # Convergence tracking
         self._last_max_gain: float = float('inf')  # Track maximum information gain
 
@@ -467,6 +476,9 @@ def select_action(self, state: Any) -> Tuple[str, List[str]]:
         # Ensure state is in set format
         state = self._ensure_state_is_set(state)
 
+        # Store for use by lookahead selector (needs state for simulation)
+        self._current_state = state
+
         # Handle object subset selection with state awareness
         if self.use_object_subset and self.subset_manager:
             # First iteration: select initial subset using state information
@@ -683,42 +695,7 @@ def _calculate_all_action_gains_sequential(self, grounded_actions: List, state:
 
         # Sort by gain (highest first)
         action_gains.sort(key=lambda x: x[2], reverse=True)
-
-        # Filter/prioritize injective bindings based on action execution history
-        # Actions without successful observations: filter out non-injective (strict)
-        # Actions with successful observations: prioritize injective (lenient)
-        filtered_gains = []
-        for action_name, objects, gain in action_gains:
-            is_injective = is_injective_binding(objects)
-            has_success = self._has_successful_observation(action_name)
-
-            if has_success:
-                # Action already executed successfully - prioritize injective but allow non-injective
-                filtered_gains.append((action_name, objects, gain, is_injective))
-            else:
-                # Action not yet executed - filter out non-injective (unless no alternative)
-                if is_injective:
-                    filtered_gains.append((action_name, objects, gain, is_injective))
-                # Non-injective for unexecuted action: defer (handled in fallback below)
-
-        # Separate by injectivity, preserving gain ordering within each group
-        injective_actions = [(a, o, g) for a, o, g, inj in filtered_gains if inj]
-        non_injective_actions = [(a, o, g) for a, o, g, inj in filtered_gains if not inj]
-
-        # Fallback: For action types with NO injective bindings (unexecuted), add non-injective
-        actions_with_injective = {a for a, _, _, inj in filtered_gains if inj}
-        for action_name, objects, gain in action_gains:
-            if not is_injective_binding(objects) and not self._has_successful_observation(action_name):
-                # Check if this action type has any injective options
-                if action_name not in actions_with_injective:
-                    non_injective_actions.append((action_name, objects, gain))
-                    # TODO: use ESAM effects logic on non-injective bindings
-
-        if non_injective_actions:
-            logger.debug(f"Injective binding filter: {len(injective_actions)} injective, "
-                         f"{len(non_injective_actions)} non-injective actions")
-
-        action_gains = injective_actions + non_injective_actions
+        action_gains = self._apply_injective_binding_filter(action_gains)
 
         # Log top action after sorting
         if action_gains:
@@ -823,49 +800,51 @@ def _calculate_all_action_gains_parallel(self, grounded_actions: List, state: Se
 
         # Sort by gain (highest first)
         action_gains.sort(key=lambda x: x[2], reverse=True)
+        action_gains = self._apply_injective_binding_filter(action_gains)
+
+        # Log top action after sorting
+        if action_gains:
+            top_action, top_objects, top_gain = action_gains[0]
+            logger.debug(f"\nTop action after sorting (parallel): {top_action}({','.join(top_objects)}) with E[gain]={top_gain:.6f}")
 
-        # Filter/prioritize injective bindings based on action execution history
-        # Actions without successful observations: filter out non-injective (strict)
-        # Actions with successful observations: prioritize injective (lenient)
+        return action_gains
+
+    def _apply_injective_binding_filter(
+        self, action_gains: List[Tuple[str, List[str], float]]
+    ) -> List[Tuple[str, List[str], float]]:
+        """
+        Filter/prioritize injective bindings based on action execution history.
+
+        Actions without successful observations: filter out non-injective (strict).
+        Actions with successful observations: prioritize injective (lenient).
+        """
         filtered_gains = []
         for action_name, objects, gain in action_gains:
             is_injective = is_injective_binding(objects)
             has_success = self._has_successful_observation(action_name)
 
             if has_success:
-                # Action already executed successfully - prioritize injective but allow non-injective
                 filtered_gains.append((action_name, objects, gain, is_injective))
             else:
-                # Action not yet executed - filter out non-injective (unless no alternative)
                 if is_injective:
                     filtered_gains.append((action_name, objects, gain, is_injective))
-                # Non-injective for unexecuted action: defer (handled in fallback below)
 
-        # Separate by injectivity, preserving gain ordering within each group
         injective_actions = [(a, o, g) for a, o, g, inj in filtered_gains if inj]
         non_injective_actions = [(a, o, g) for a, o, g, inj in filtered_gains if not inj]
 
         # Fallback: For action types with NO injective bindings (unexecuted), add non-injective
         actions_with_injective = {a for a, _, _, inj in filtered_gains if inj}
         for action_name, objects, gain in action_gains:
             if not is_injective_binding(objects) and not self._has_successful_observation(action_name):
-                # Check if this action type has any injective options
                 if action_name not in actions_with_injective:
                     non_injective_actions.append((action_name, objects, gain))
                     # TODO: use ESAM effects logic on non-injective bindings
 
         if non_injective_actions:
-            logger.debug(f"Injective binding filter (parallel): {len(injective_actions)} injective, "
+            logger.debug(f"Injective binding filter: {len(injective_actions)} injective, "
                          f"{len(non_injective_actions)} non-injective actions")
 
-        action_gains = injective_actions + non_injective_actions
-
-        # Log top action after sorting
-        if action_gains:
-            top_action, top_objects, top_gain = action_gains[0]
-            logger.debug(f"\nTop action after sorting (parallel): {top_action}({','.join(top_objects)}) with E[gain]={top_gain:.6f}")
-
-        return action_gains
+        return injective_actions + non_injective_actions
 
     def _create_parallel_context(self, state: Set[str]) -> "ActionGainContext":
         """
@@ -1072,6 +1051,36 @@ def _select_by_strategy(self, action_gains: List[Tuple[str, List[str], float]])
             # Fallback to last action
             return action_gains[-1][0], action_gains[-1][1]
 
+        elif self.selection_strategy == 'lookahead':
+            # Bounded lookahead: evaluate top-k actions with depth-limited future gain
+            from information_gain_aml.algorithms.mcts_selector import BoundedLookaheadSelector
+            selector = BoundedLookaheadSelector(
+                self,
+                depth=self.lookahead_depth,
+                top_k=self.lookahead_top_k,
+                discount=self.lookahead_discount,
+            )
+            # _select_by_strategy receives action_gains but not state;
+            # we need the current state for simulation — stored from select_action()
+            best = selector.select_action(self._current_state, action_gains)
+            logger.debug(f"Lookahead: Selected {best[0]}({','.join(best[1])})")
+            return best[0], best[1]
+
+        elif self.selection_strategy == 'mcts':
+            # Full UCT-based MCTS action selection
+            from information_gain_aml.algorithms.mcts_selector import IGMCTSSelector
+            mcts_selector = IGMCTSSelector(
+                self,
+                iterations=self.mcts_iterations,
+                rollout_depth=self.mcts_rollout_depth,
+            )
+            best = mcts_selector.select_action(self._current_state, action_gains)
+            logger.debug(
+                f"MCTS: Selected {best[0]}({','.join(best[1])}) "
+                f"after {self.mcts_iterations} iterations"
+            )
+            return best[0], best[1]
+
         else:
             # Unknown strategy, default to greedy
             logger.warning(f"Unknown selection strategy '{self.selection_strategy}', defaulting to greedy")
@@ -2107,6 +2116,15 @@ def export_model_snapshot(self, iteration: int, output_dir: Path) -> None:
                 "constraint_sets": [sorted(list(cs)) for cs in self.pre_constraints.get(action_name, set())]
             }
 
+        # Add action model statistics for intermediate analysis
+        snapshot["statistics"] = {
+            "iterations": self.iteration_count,
+            "observations": self.observation_count,
+            "converged": self._converged,
+            "max_information_gain": self._last_max_gain,
+            "action_model_metrics": self.get_action_model_metrics()
+        }
+
         # Write to file
         output_path = models_dir / f"model_iter_{iteration:03d}.json"
         with open(output_path, 'w') as f: