Merge pull request #28 from TRI-ML/update_readme

Update README.md and Incorporate Quickstart Resources

Author: DavidSnyder-TRI

README.md

@@ -1,9 +1,39 @@
 # sequentialized_barnard_tests
 A collection of sequential statistical hypothesis testing methods for two-by-two contingency tables.
 
-## Development Instructions
+[[Paper](https://www.arxiv.org/abs/2503.10966)]  [[Website](https://tri-ml.github.io/step/)]
+<!-- ## Development Plan
+This codebase will be developed into a standalone pip package, with planned release date June 2025.
 
-An example of the enrivonment setup is shown below.
+Current features:
+- Fully automated STEP policy synthesis
+- Baseline sequential method implementations of [SAVI](https://www.sciencedirect.com/science/article/pii/S0167715223000597?via%3Dihub) and [Lai](https://projecteuclid.org/journals/annals-of-statistics/volume-16/issue-2/Nearly-Optimal-Sequential-Tests-of-Composite-Hypotheses/10.1214/aos/1176350840.full) procedures.
+- Validation tools
+    - STEP policy visualization
+    - Verification of Type-1 Error control
+- Unit tests
+    - Method functionality
+    - Recreate results from our [paper](https://www.arxiv.org/abs/2503.10966).
+
+Features in development:
+- Demonstration scripts for each stage of the STEP evaluation pipeline
+- Approximately optimal risk budget estimation tool based on evaluator priors on $$(p_0, p_1)$$
+- Fundamental limit estimator for guiding evaluation effort
+    - Determine if a particular effect size is plausibly discoverable given the evaluation budget
+- Baseline implementation of a sequential Barnard procedure which controls Type-1 error -->
+
+## Installation Instructions \[Standard\]
+The basic environmental setup is shown below. A virtual / conda environment may be constructed; however, the requirements are quite lightweight and this is probably not needed.
+```bash
+$ cd <some_directory>
+$ git clone [email protected]:TRI-ML/sequentialized_barnard_tests.git
+$ cd sequentialized_barnard_tests
+$ pip install -r requirements.txt
+$ pip install -e .
+```
+
+## Installation Instructions \[Dev\]
+For potential contributors and developers, we recommend a virtualenv:
 ```bash
 $ cd <some_directory>
 $ git clone [email protected]:TRI-ML/sequentialized_barnard_tests.git
@@ -14,3 +44,129 @@ $ pip install -r requirements.txt
 $ pip install -e .
 $ pre-commit install
 ```
+
+We assume that any specified virtual / conda environment has been activated for all subsequent code snippets.
+
+# Quick Start Guides
+We include key notes for understanding the core ideas of the STEP code. Quick-start resources are included in both shell script and notebook form.
+
+## Quick Start Guide: Making a STEP Policy for Specific \{n_max, alpha\}
+
+### (1A) Understanding the Accepted Shape Parameters
+In order to synthesize a STEP Policy for specific values of n_max and alpha, one additional set of parametric decisions will be required. The user will need to set the risk budget shape, which is specified by choice of function family (p-norm vs zeta-function) and particular shape parameter. The shape parameter is real-valued; it is used directly for zeta functions and is exponentiated for p-norms.
+
+**For p-norms**
+- $$\text{Shape Parameter: } \lambda \in \mathbb{R}$$
+- $$\text{Accumulated Risk Budget}(n) = \alpha \cdot (\frac{n}{n_{max}})^{\exp{(\lambda)}}$$
+
+**For zeta function**
+- $$\text{Shape Parameter: } \lambda \in \mathbb{R}$$
+- $$\text{Accumulated Risk Budget}(n) = \frac{\alpha}{Z(n_{max})} \cdot \sum_{i=1}^n (\frac{1}{i})^{\lambda}$$
+- $$Z(n_{max}) = \sum_{i=1}^{n_{max}} (\frac{1}{i})^{\lambda}$$
+
+The user may confirm that in each case, evaluating the accumulated risk budget at $`n=n_{max}`$ returns precisely $\alpha$.
+
+
+### (1B) Shape Parameters in Code
+In the codebase, the value of $\lambda$ is set in the \{shape_parameter\} variable. This variable is real-valued.\
+\
+The family shape is set by the \{use_p_norm\} variable. This variable is Boolean.
+- If it is True, then p-norm family is used.
+- If it is False, the zeta-function family is used.
+
+
+### (1C) Arbitrary Risk Budgets
+Generalizing the accepted risk budgets to arbitrary monotonic sequences $`\{0, \epsilon_1 > 0, \epsilon_2 > \epsilon_1, ..., \epsilon_{n_{max}} = \alpha\}`$ is in the development pipeline, but is **not handled at present in the code**.
+
+
+### (2A) Running STEP Policy Synthesis
+Having decided an appropriate form for the risk budget shape, policy synthesis is straightforward to run. From the base directory, the general command would be:
+
+```bash
+$ python scripts/synthesize_general_step_policy.py -n {n_max} -a {alpha} -pz {shape_parameter} -up {use_p_norm}
+```
+
+### (2B) What If I Don't Know the Right Risk Budget?
+We recommend using the default linear risk budget, which is the shape *used in the paper*. This corresponds to \{shape_parameter\}$`= 0.0`$ for each shape family. Thus, *each of the following commands constructs the same policy*:
+
+```bash
+$ python scripts/synthesize_general_step_policy.py -n {n_max} -a {alpha}
+```
+```bash
+$ python scripts/synthesize_general_step_policy.py -n {n_max} -a {alpha} -pz {0.0} -up "True"
+```
+```bash
+$ python scripts/synthesize_general_step_policy.py -n {n_max} -a {alpha} -pz {0.0} -up "False"
+```
+
+Note: For \{shape_parameter\} $`\neq 0`$, the shape families differ. Therefore, the choice of \{use_p_norm\} *will affect the STEP policy*.
+
+### (2C) Disclaimers
+- Running the policy synthesis will save a durable policy to the user's local machine. This policy can be reused for all future settings requiring the same \{n_max, alpha\} combination. For \{n_max\} $`< 500`$, the amount of required memory is < 5Mb. The policy is saved under:
+```bash
+$ sequentialized_barnard_tests/policies/
+```
+
+- At present, we have not tested extensively beyond \{n_max\}$`=500`$. Going beyond this limit may lead to issues, and the likelihood will grow the larger \{n_max\} is set to be. The code will also require increasing amounts of RAM as \{n_max\} is increased.
+
+## Quick Start Guide: Evaluation on Real Data
+
+We now assume that a STEP policy has been constructed for the target problem. This can either be one of the default policies, or a newly constructed one following the recipe in the preceding section.
+
+### (1) Formatting the Real Data
+The data should be formatted into a numpy array of shape $(N, 2)$. The user should create a new project directory and save the data within it:
+```bash
+$ mkdir data/{new_project_dir}
+$ cp path/to/{my_data_file.npy} data/{new_project_dir}/{my_data_file.npy}
+```
+
+We give an example that *would have generated* the included data:
+```bash
+$ mkdir data/example_clean_spill
+$ cp some/path/to/TRI_CLEAN_SPILL_v4.npy data/example_clean_spill/TRI_CLEAN_SPILL_v4.npy
+```
+
+### (2) Running the Evaluation
+Then, the user need simply run the evaluation script, which requires the project directory and file in addition to the policy synthesis arguments:
+```bash
+$ python evaluation/run_step_on_evaluation_data.py -p "{new_project_dir}" -f "{my_data_file.npy}" -n {n_max} -a {alpha} -pz {shape_parameter} -up "{use_p_norm}"
+```
+
+This will print the evaluation result to the terminal, as well as save key information in a timestamped json file.
+
+We illustrate this via an evaluation on the default data:
+```bash
+$ python evaluation/run_step_on_evaluation_data.py -p "example_clean_spill" -f "TRI_CLEAN_SPILL_v4.npy" -n {200} -a {0.05} -pz {0.0} -up "False"
+```
+
+# Code Overview
+Scripts to generate and visualize STEP policies are included under:
+```bash
+$ scripts/
+```
+
+Any resulting visualizations are stored in:
+```bash
+$ media/
+```
+
+The evaluation environment for real data is included in:
+```bash
+$ evaluation/
+```
+
+and the associated evaluation data is stored in:
+```bash
+$ data/
+```
+
+# Citation
+
+```
+@inproceedings{snyder2025step,
+title = {Is Your Imitation Learning Policy Better Than Mine? Policy Comparison with Near-Optimal Stopping},
+author = {Snyder, David and Hancock, Asher James and Badithela, Apurva and Dixon, Emma and Miller, Patrick and Ambrus, Rares Andrei and Majumdar, Anirudha and Itkina, Masha and Nishimura, Haruki},
+booktitle={Proceedings of the Robotics: Science and Systems Conference (RSS)},
+year = {2025},
+}
+```

data/example_clean_spill/TRI_CLEAN_SPILL_v4.npy

@@ -0,0 +1 @@
+{"result": [{"Hypothesis": "P0 < P1", "Decision": "P0 >= P1", "Time": 9}], "method": "STEP", "params": [{"n_max": 200, "alpha": 0.05}], "p0_hat": 0.8, "p1_hat": 0.28, "N": 50}

data/example_clean_spill/evaluation_result_2025-05-19_15:59:00.json

@@ -0,0 +1,195 @@
+import argparse
+import json
+import os
+from datetime import datetime
+
+import numpy as np
+
+from sequentialized_barnard_tests.base import Decision, Hypothesis
+from sequentialized_barnard_tests.step import MirroredStepTest
+
+# from ..sequentialized_barnard_tests.lai import MirroredLaiTest
+# from ..sequentialized_barnard_tests.savi import MirroredOracleSaviTest
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description=(
+            "This script runs a pre-computed STEP policy on real data. The data "
+            "is assumed to be stored as '<base_path>/data/<project_folder>/<file_name>.npy.' "
+            "Further, the data is assumed to be of shape (N, 2) corresponding to the success-"
+            "failure values observed in evaluation trials. Input parameters determine the nature "
+            "of the STEP test, certain save options, and the data project to access. "
+        )
+    )
+    parser.add_argument(
+        "-p",
+        "--data_folder",
+        type=str,
+        default="example_clean_spill",
+        help=(
+            "Relative path added to <base_path>/data/ which specifies the desired "
+            "evaluation data folder, from which to run the STEP test. Defaults to ''."
+        ),
+    )
+    parser.add_argument(
+        "-f",
+        "--data_file",
+        type=str,
+        default="TRI_CLEAN_SPILL_v4.npy",
+        help=(
+            "Relative path added to <base_path>/data/<project_folder>/ which specifies the desired "
+            "evaluation data on which to run the STEP test. Defaults to ''."
+        ),
+    )
+    parser.add_argument(
+        "-n",
+        "--n_max",
+        type=int,
+        default=200,
+        help=(
+            "Maximum number of robot policy evals (per policy) in the evaluation procedure. "
+            "Defaults to 200."
+        ),
+    )
+    parser.add_argument(
+        "-a",
+        "--alpha",
+        type=float,
+        default=0.05,
+        help=(
+            "Maximal allowed Type-1 error rate of the statistical testing procedure. "
+            "Defaults to 0.05."
+        ),
+    )
+    parser.add_argument(
+        "-pz",
+        "--log_p_norm",
+        type=float,
+        default=0.0,
+        help=(
+            "Rate at which risk is accumulated, reflecting user's belief about underlying "
+            "likelihood of different alternatives and nulls being true. If using a p_norm "
+            ", this variable is equivalent to log(p). If not using a p_norm, this is the "
+            "argument to the zeta function, partial sums of which give the shape of the risk budget."
+            "Defaults to 0.0."
+        ),
+    )
+    parser.add_argument(
+        "-up",
+        "--use_p_norm",
+        type=str,
+        default="False",
+        help=(
+            "Toggle whether to use p_norm or zeta function shape family for the risk budget. "
+            "If True, uses p_norm shape; else, uses zeta function shape family. "
+            "Defaults to False (zeta function partial sum family)."
+        ),
+    )
+    parser.add_argument(
+        "-so",
+        "--save_output",
+        type=bool,
+        default=True,
+        help=(
+            "Toggle whether save the evaluation output in a timestamped config file, as opposed to only printing to terminal. "
+            "If True, the evaluation result is saved in a json file in the same directory as the data. "
+            "Defaults to True."
+        ),
+    )
+    parser.add_argument(
+        "-uda",
+        "--use_default_alternative",
+        type=bool,
+        default=True,
+        help=(
+            "Determines the alternative to use in the testing procedure. If True, uses the alternative p0 < p1. "
+            "Otherwise, uses the alternative p0 > p1. Defaults to True."
+        ),
+    )
+
+    # Parse the input arguments
+    args = parser.parse_args()
+    if args.use_p_norm[0] == "F" or args.use_p_norm[0] == "f":
+        args.use_p_norm = False
+    elif args.use_p_norm[0] == "T" or args.use_p_norm[0] == "t":
+        args.use_p_norm = True
+    else:
+        raise ValueError(
+            "Invalid argument {for use_p_norm}; must be a string beginning with 'T' or 'F'"
+        )
+
+    # Define the base path so that the file can be run from any directory
+    # base_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")
+    base_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+
+    # Define the data folder path and data file path
+    data_folder_path = os.path.join(base_path, "data", args.data_folder)
+    data_file_path = os.path.join(data_folder_path, args.data_file)
+
+    # Assign step test alternative hypothesis
+    if args.use_default_alternative:
+        alt = Hypothesis.P0LessThanP1
+        alt_str = "P0 < P1"
+        null_str = "P0 >= P1"
+    else:
+        alt = Hypothesis.P0MoreThanP1
+        alt_str = "P0 > P1"
+        null_str = "P0 <= P1"
+    # Initialize the STEP test
+    mirrored_step_test = MirroredStepTest(
+        alternative=alt,
+        n_max=args.n_max,
+        alpha=args.alpha,
+        shape_parameter=args.log_p_norm,
+        use_p_norm=args.use_p_norm,
+    )
+
+    # Load data and evaluate
+    data = np.load(data_file_path)
+
+    result = mirrored_step_test.run_on_sequence(data[:, 0], data[:, 1])
+
+    decision_str = "No decision"
+    if result.decision == Decision.AcceptAlternative:
+        decision_str = alt_str
+    elif result.decision == Decision.AcceptNull:
+        decision_str = null_str
+    else:
+        pass
+
+    # Print to terminal
+    print(
+        f"Evaluation result for data stored at: {os.path.join(data_folder_path, args.data_file)}"
+    )
+    print(f"Decision: {result.decision} --> we conclude that {decision_str}")
+    print(f"Time of decision: {result.info['Time']}")
+
+    # Save to json file
+    if args.save_output:
+        now = datetime.now()
+        formatted_time = now.strftime("%Y-%m-%d_%H:%M:%S")
+        evaluation_dict = {
+            "result": [
+                {
+                    "Hypothesis": alt_str,
+                    "Decision": decision_str,
+                    "Time": result.info["Time"],
+                }
+            ],
+            "method": "STEP",
+            "params": [
+                {
+                    "n_max": args.n_max,
+                    "alpha": args.alpha,
+                }
+            ],
+            "p0_hat": np.mean(data[:, 0]),
+            "p1_hat": np.mean(data[:, 1]),
+            "N": data.shape[0],
+        }
+
+        with open(
+            data_folder_path + "/" + f"evaluation_result_{formatted_time}.json", "w"
+        ) as fp:
+            json.dump(evaluation_dict, fp)