PyAres Ax Planners

This repository contains methods and examples for using Bayesian optimization and active learning planners based on Meta's Ax API with ARES OS using PyAres. It provides a compatibility layer to translate ARES OS planning requests into a format compatible with the Ax platform.

Features

• Ax v1.3.0+ Compliance: Completely updated for Meta's modern, graph-based GenerationNode architecture.
• Custom Constraints: Supports the use of implicit values within constraints and evaluates constraint strings via SymPy prior to planning.
• Implicit Values & Parameters: Supports planning with implicit variables in cases where constraints mandate that some parameters coming from ARES OS are dependent on other parameters. You can define implicit values using SymPy-compatible string expressions (e.g., "flow_4 = total_flow - (flow_1 + flow_2 + flow_3)").
• Dual-Mode Hardware Intercept: Features optional middleware capability to dynamically scale and cap search boundaries based on physical mass flow controller (MFC) hardware capacities and target total flow rates.
• Stateless-Safe Execution: Programmatically bypasses Ax amnesia issues caused by stateless network protocols, ensuring smooth transitions between random initialization and active surrogate model optimization.
• Previous Data: Imports prior experimental seed data from .csv, .xls, or .xlsx files. Column headers should match ARES OS provided parameter names, with the objective score value column labeled objective.
• Confidence Threshold Stopping: Supports autonomous campaign halting when model uncertainty (Standard Error of the Mean) drops below a user-defined threshold.

Currently Implemented Planners

• Single Objective Bayesian Optimization (SOBO): SOBO_Ax_Planner provides standard single objective optimization for continuous variables, driving parameters toward a maximum or minimum target value.
• Active Learning Boundary Hunter: SOBO_Active_Learning_Planner executes variance-based exploration. Rather than maximizing a value, it maps out unknown regions and boundary envelopes by aggressively deploying BoTorch's qNegIntegratedPosteriorVariance acquisition function to probe areas of highest uncertainty.

---

Installation & Environment Setup

This project requires Python >=3.11. A dedicated environment is highly recommended to keep module dependencies for different PyAres services from causing issues.

Option A: Using Anaconda / Miniconda / Miniforge (Recommended)

1. Create the environment:

   conda create -n pyares_ax python=3.11 pip setuptools

2. Activate the environment:

   conda activate pyares_ax

3. Install the package: Navigate to the project root directory and run:

   pip install .

   Note: This will automatically install dependencies listed in pyproject.toml.

Option B: Using Python venv

1. Ensure you have Python >=3.11 installed: Check your version:

   python3 --version

2. Create the virtual environment:

   python3 -m venv pyares_ax_venv

3. Activate the environment:
   • Windows:

     .\pyares_ax_venv\Scripts\activate

   • Linux/macOS:

     source pyares_ax_venv/bin/activate

4. Install the package:

   pip install .

---

Quick Start

To launch the planner service daemon, run the provided startup script:

python start_pyares_ax_planner.py

This script initializes the communication layer and starts the AresPlannerService on port 1337.

---

Planner Settings

The core PyAres_Ax_Planner framework exposes several base settings by default that can be configured directly from the ARES OS user interface:

• Seed Data: A string path pointing to a local file (.csv, .xls, or .xlsx) containing historical campaign data.
• Constraints: An array of string constraints evaluated by SymPy (e.g., ["Flow_Ar + Flow_O2 <= 100.0"]).
• Implicit Values: An array of string definitions for calculating implicit factors and parameters (e.g., ["Flow_N2 = 100.0 - (Flow_Ar + Flow_O2)"]).
• Parameter Value Type: A string restricted to either 'Planned' or 'Achieved' (Default: 'Planned'). Set to 'Achieved' to force the model to train on actual MFC physical readings rather than target commands.
• Verbose Output: A boolean to toggle detailed telemetry logging in the console (Default: True).

Active Learning Planner Settings

When running the SOBO_Active_Learning_Planner instance, the following additional configurations are available to control the dual-mode hardware behavior and exploration logic:

• Minimize (Boolean): Set to False when mapping peak operational windows or high-signal regions. Set to True when tracking noise floors or hardware defect transitions.
• Initial Random Points (Number): The number of initial points generated via Sobol sequence to scout the space before active learning engages (Default: 5.0).
• Confidence Threshold (SEM) (Number): The minimum confidence limit (in Standard Error of the Mean). If model uncertainty drops below this threshold, the planner returns a flag to ARES OS requesting the campaign to halt (Default: 0.5).
• RNG Seed (Number): Leave blank to initialize the pseudo-random Sobol sequence via the host clock. Enter an integer (e.g., 42) to lock in reproducible testing paths for debugging.
• Target Total Flow (Number): The total targeted flow rate of the manifold (e.g., 400.0 sccm). Leave at 0.0 to run as a Pure Black-Box.
• Active MFCs (JSON String): A serialized text dictionary specifying the maximum capacities of your live flow meters. Leave blank to run as a Pure Black-Box.
  • Example: {"Flow_Ar": 500, "Flow_O2": 500, "Flow_N2": 100}

⚠️ CRITICAL STRING CASING NOTE: Strings used within constraints, implicit values, and the Active MFCs dictionary are strictly case-sensitive and character-exact. They must perfectly match the parameter names passed in the raw network PlanRequest (e.g., if ARES registers Flow_Ar, typing flow_ar or Ar will cause a runtime failure).

---

Implementing a Custom Planner

You can easily create your own Ax-based planner for ARES OS by subclassing the PyAres_Ax_Planner base class. This base class handles all the heavy lifting of translating PyAres requests, processing historical/seed data into pandas DataFrames, and managing implicit constraints.

To implement a custom planner, follow these steps:

1. Subclass PyAres_Ax_Planner and Override Initialization Attributes

In your subclass's __init__ method, call super().__init__() and then override the placeholder attributes with your planner's specific information.

from .ax_compatibility import PyAres_Ax_Planner
from PyAres import AresDataType

class MyCustom_Ax_Planner(PyAres_Ax_Planner):
    def __init__(self):
        super().__init__()
        # Override generic planner info
        self.name = "My Custom Ax Planner"
        self.description = "A tailored Ax Bayesian Optimization planner."
        self.version_number = "1.0.0"
        
        # Link to your custom planning function
        self.plan_function = my_custom_planner_logic 
        
        # Add any planner-specific settings exposed to ARES OS
        self.add_setting('Minimize', AresDataType.BOOLEAN, False)
        self.add_setting('Custom Iterations', AresDataType.NUMBER, default_value=10)

2. Override _configure_objectives()

Override this method in your child class to properly configure the Ax ObjectiveProperties based on the active parameters provided by ARES OS.

    def _configure_objectives(self):
        minimize_flag = self.settings.get('Minimize', False)
        self.objectives = {'objective': ObjectiveProperties(minimize=minimize_flag)}

3. Define Your plan_function

The plan_function is the core routine of your planner. It is called by the compatibility layer and is provided with fully parsed and formatted Ax API inputs. It must accept specific arguments and return a tuple containing the proposed next test condition (dictionary) and a boolean indicating whether the campaign is complete.

from ax.service.ax_client import AxClient
from ax.generation_strategy.generation_strategy import GenerationStrategy
from ax.generation_strategy.generation_node import GenerationNode
from ax.generation_strategy.generator_spec import GeneratorSpec
from ax.adapter.registry import Generators
from PyAres import PlanStatusCode

def my_custom_planner_logic(parameters: list[dict], 
                            objective: dict, 
                            constraints: list[str], 
                            data: list[dict],
                            settings: dict,
                            previous_plan_code: PlanStatusCode) -> tuple[dict, bool]:
    """
    Args:
        parameters: List of Ax formatted parameters
        objective: Dict containing the ObjectiveProperties object
        constraints: List of string constraints (evaluated by SymPy)
        data: List of dicts representing historical trials ('parameters' and 'objectives')
        settings: Dict of planner settings from ARES OS
        previous_plan_code: Enum indicating the success/failure status of the last executed plan
    """
    # 1. Define a v1.3.0+ Compliant Generation Strategy using Graph Nodes
    custom_gs = GenerationStrategy(
        name="Custom_BO_Strategy",
        nodes=[
            GenerationNode(
                node_name="Optimization_Phase",
                generator_specs=[
                    GeneratorSpec(generator_enum=Generators.BOTORCH_MODULAR)
                ]
            )
        ]
    )

    # 2. Initialize your Ax Client with the strategy
    ax_client = AxClient(generation_strategy=custom_gs)
    ax_client.create_experiment(
        parameters=parameters,
        objectives=objective,
        parameter_constraints=constraints
    )

    # 3. Attach historical and seed data
    for trial in data:
        _, trial_index = ax_client.attach_trial(parameters=trial['parameters'])
        
        # Check if the previous experiment iteration crashed the physical rig
        is_crashed = previous_plan_code == PlanStatusCode.PLAN_FAILED 
        
        if is_crashed:
            ax_client.log_trial_failure(trial_index=trial_index)
        else:
            ax_client.complete_trial(trial_index=trial_index, raw_data=trial['objectives'])
        
    # 4. Generate the next trial (triggers the Gaussian Process fit)
    parameterization, trial_index = ax_client.get_next_trial()
    
    # 5. Return the predicted parameters, and False to keep the campaign running
    return parameterization, False

By adhering to this structure, your custom Ax planner automatically inherits robust data parsing, hardware exception handling, smart seeding data structures, and seamless integration with the ARES OS network environment.

License

This project is licensed under the MIT License. Copyright (c) 2026 AFRL-ARES.