fix aviary readme

environments/aviary_bbh/README.md

@@ -1,46 +1,166 @@
 # BixBench-Hypothesis (BBH)
+BixBench-Hypothesis is a dataset proposed by Edison Scientific to measure LLM capabilities for testing hypotheses in bioinformatics contexts. Edison Scientific and NVIDIA have also collaborated to release BBH-Train, an RL training dataset meant to improve model capabilities on bioinformatics-related data analysis.
 
-[BixBench-Hypothesis](https://huggingface.co/datasets/nvidia/Nemotron-RL-bixbench_hypothesis) measures hypothesis testing over bioinformatics data. The remote config connects Gym to an external [Hypotest](https://github.com/EdisonScientific/hypotest) dataset server; the bundled config runs Hypotest and its Enroot sandbox alongside Gym.
+There are two methods to running BBH with NeMo-Gym: the remote approach and the bundled approach. The remote approach hosts the environment sandboxes as an external service that NeMo-Gym can communicate with, serving as a modular method for running train/inference jobs. The bundled approach colocates the environment sandboxes on the Gym/RL nodes, serving as an efficient and fully packaged method for running train/inference jobs.
 
-The commands below assume that a model endpoint is configured with `policy_base_url`, `policy_model_name`, and `policy_api_key` in `env.yaml`. See the [local configuration documentation](https://docs.nvidia.com/nemo/gym/reference/configuration#local-configuration-envyaml).
+## Remote Approach:
+To run the remote approach, first launch the dataset server implemented in the [hypotest](https://github.com/EdisonScientific/hypotest) repository. Documentation on how to run the dataset server can be found [here](https://github.com/EdisonScientific/hypotest/blob/main/README.md).
 
-## Remote server
+Then, prepare your Gym data with the task_idx values of the problems you would like to train/evaluate on. An example dataset is provided for reference in [data/example.jsonl](data/example.jsonl).
 
-Start the Hypotest dataset server, then set `server_url` and `api_key` in `config_remote.yaml`. Task indices in the Gym dataset must refer to the same problems served by Hypotest. The checked-in [example dataset](data/example.jsonl) shows the expected format.
+Once the dataset server is running and is accessible at a specific URL, update your config based on [config_remote.yaml](config_remote.yaml) with the server URL and api key, and launch NeMo-Gym as follows:
 
 ```bash
-gym env start --environment aviary_bbh/config_remote --model-type vllm_model
+gym env start \
+    --environment aviary_bbh/config_remote \
+    --model-type vllm_model
 ```
 
-Keep that terminal running, then in another terminal run:
+Then collect rollouts on your data as follows (updating the input file to your Gym data if needed):
 
 ```bash
 gym eval run --no-serve \
-  --agent bbh_aviary_agent \
-  --input environments/aviary_bbh/data/example.jsonl \
-  --output environments/aviary_bbh/data/example_rollouts_remote.jsonl
+    --agent bbh_aviary_agent \
+    --input environments/aviary_bbh/data/example.jsonl \
+    --output environments/aviary_bbh/data/example_bbh_rollouts.jsonl
 ```
 
-## Bundled server
+To run training with NeMo-RL, set the following fields in your NeMo-RL container (where train_data.jsonl and validation_data.jsonl are set to your train/val Gym data respectively, and bbh_remote.yaml is updated with your dataset server URL/api-key):
+```yaml
+data:
+  train_jsonl_fpath: 3rdparty/Gym-workspace/Gym/environments/aviary_bbh/data/train_data.jsonl
+  validation_jsonl_fpath: 3rdparty/Gym-workspace/Gym/environments/aviary_bbh/data/validation_data.jsonl
+  shuffle: False
+  num_workers: 1
 
-Set the dataset paths, rubric model, and `container_sqsh_path` in `config_bundled.yaml`. The bundled environment requires Enroot and a Hypotest container image.
+env:
+  should_use_nemo_gym: true
+  nemo_gym:  # This is passed into NeMo-Gym as the initial_global_config_dict
+    is_trajectory_collection: false  # Set this to true to enable trajectory collection (no training). You may also want to increase `policy.generation.vllm_cfg.gpu_memory_utilization`
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml  # Required! And it must be *for_training
+    - environments/aviary_bbh/config_remote.yaml
+```
+
+Note that task_idx values in your Gym data must align with the data in [hypotest](https://github.com/EdisonScientific/hypotest); this means that both your train and val set problem data must be provided to [hypotest](https://github.com/EdisonScientific/hypotest)'s dataset server.
+
+## Bundled Approach:
+To run the bundled approach, first update the config in [config_bundled.yaml](config_bundled.yaml) with your desired configuration. The config fields in `dataset` closely match the config fields in [hypotest](https://github.com/EdisonScientific/hypotest). You'll also have to set [container_sqsh_path], which will be a path to a .sqsh file built from the [hypotest](https://github.com/EdisonScientific/hypotest) Docker container using [enroot](https://github.com/NVIDIA/enroot).
 
+You will also need to add your BBH problem data to [data/](data/), and update `capsule_dir` and `work_dir` with paths to your BBH capsule data and working directory. Note that the working directory must be set to a directory accessible to all nodes if you are running with multi-node jobs, in order for the environment to properly parallelize sandboxes across all available nodes (e.g. the working dir could be made be available on a network-filesystem like lustre).
+
+Once you have your environment properly configured, you'll need to make sure [enroot](https://github.com/NVIDIA/enroot) is installed into your Gym environment. In order to do this, make sure to run the following snippet before bringing up NeMo-Gym:
 ```bash
-gym env start --environment aviary_bbh/config_bundled --model-type vllm_model
+cd /tmp &&
+apt-get update &&
+arch=$(dpkg --print-architecture) &&
+curl -fSsL -O https://github.com/NVIDIA/enroot/releases/download/v4.1.1/enroot_4.1.1-1_\${arch}.deb &&
+curl -fSsL -O https://github.com/NVIDIA/enroot/releases/download/v4.1.1/enroot+caps_4.1.1-1_\${arch}.deb &&
+apt install -y ./*.deb &&
+apt-get install -y squashfuse &&
+cd /path/to/gym/directory
+```
+And then bring up NeMo-Gym:
+```bash
+gym env start \
+    --environment aviary_bbh/config_bundled \
+    --model-type vllm_model
+```
+```bash
+gym eval run --no-serve \
+    --agent bbh_aviary_agent \
+    --input environments/aviary_bbh/data/example.jsonl \
+    --output environments/aviary_bbh/data/example_bbh_rollouts.jsonl
 ```
 
-Keep that terminal running, then in another terminal run:
+If you are running training with NeMo-RL and NeMo-Gym, add the following modification to your `ray.sub` file in NeMo-RL to support adding a setup command:
+```diff
+diff --git a/ray.sub b/ray.sub
+index 9b4feb11..f765a609 100644
+--- a/ray.sub
++++ b/ray.sub
+@@ -50,6 +50,7 @@ maybe_gres_arg() {
+ CONTAINER=$CONTAINER
+ MOUNTS=$MOUNTS
+ COMMAND=${COMMAND:-}  # This is a script relative to the SLURM_SUBMIT_DIR. If left empty, it will leave the cluster idle after it's brought up.
++SETUP_COMMAND=${SETUP_COMMAND:-}  # Setup commands to run on all nodes before starting Ray
+ ########################################################
+ # Ports for all nodes (should be odd numbers since we place head/worker[0] on the same node) so all workers get the odd ports, but the head will get +1 the ports
+ NODE_MANAGER_PORT=${NODE_MANAGER_PORT:-53001}
+@@ -293,6 +294,7 @@ chmod +x /launch-head.sh
 
+ count=0
+ while [[ \$count -lt $num_retries ]]; do
++  $SETUP_COMMAND
+   bash /launch-head.sh
+   count=\$((count+1))
+   echo "Head node failed \$count/$num_retries times, restarting in 5 seconds..."
+@@ -305,6 +307,7 @@ EOF
+ srun $COMMON_SRUN_ARGS --container-name=ray-head --nodes=1 --ntasks=1 --cpus-per-task=$CPUS_PER_WORKER -w "$head_node" -o $LOG_DIR/ray-head.log bash -x -c "$head_cmd" &
+ SRUN_PIDS["ray-head"]=$!
+
++sleep 100s
+ NUM_ACTORS=$((GPUS_PER_NODE * SLURM_JOB_NUM_NODES))
+
+ # Start Ray worker nodes
+@@ -392,6 +395,7 @@ EOFINNER
+
+ count=0
+ while [[ \$count -lt $num_retries ]]; do
++  $SETUP_COMMAND
+   bash /launch-worker.sh
+   count=\$((count+1))
+   echo "Worker failed \$count/$num_retries times, restarting in 5 seconds..."
+```
+Then set up your NeMo-RL config:
+```yaml
+data:
+  train_jsonl_fpath: 3rdparty/Gym-workspace/Gym/environments/aviary_bbh/data/train_data.jsonl
+  validation_jsonl_fpath: 3rdparty/Gym-workspace/Gym/environments/aviary_bbh/data/validation_data.jsonl
+  shuffle: False
+  num_workers: 1
+
+env:
+  should_use_nemo_gym: true
+  nemo_gym:  # This is passed into NeMo-Gym as the initial_global_config_dict
+    is_trajectory_collection: false  # Set this to true to enable trajectory collection (no training). You may also want to increase `policy.generation.vllm_cfg.gpu_memory_utilization`
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml  # Required! And it must be *for_training
+    - environments/aviary_bbh/config_bundled.yaml
+```
+Following this, you can export `SETUP_COMMAND` to a snippet installing enroot prior to launching your `ray.sub` command:
 ```bash
-gym eval run --no-serve \
-  --agent bbh_aviary_agent \
-  --input environments/aviary_bbh/data/example.jsonl \
-  --output environments/aviary_bbh/data/example_rollouts_bundled.jsonl
+read -r -d '' SETUP_COMMAND <<EOF
+cd /tmp &&
+apt-get update &&
+arch=$(dpkg --print-architecture) &&
+curl -fSsL -O https://github.com/NVIDIA/enroot/releases/download/v4.1.1/enroot_4.1.1-1_\${arch}.deb &&
+curl -fSsL -O https://github.com/NVIDIA/enroot/releases/download/v4.1.1/enroot+caps_4.1.1-1_\${arch}.deb &&
+apt install -y ./*.deb &&
+apt-get install -y squashfuse &&
+cd $PWD
+EOF
+export SETUP_COMMAND
+
+COMMAND="[insert nemo rl launch command with gym config]" \
+CONTAINER=YOUR_CONTAINER \
+MOUNTS="$PWD:$PWD" \
+sbatch \
+    --nodes=${NUM_ACTOR_NODES} \
+    --account=YOUR_ACCOUNT \
+    --job-name=YOUR_JOBNAME \
+    --partition=YOUR_PARTITION \
+    --time=1:0:0 \
+    --gres=gpu:8 \
+    ray.sub
 ```
 
-Generate a task-index dataset with `prepare.py --size SIZE --output PATH`.
+# Licensing information
+Code: Apache 2.0
 
-## Licensing
+Data: MIT (GSM8k),  Apache 2.0 (BixBench), CC BY 4.0 (BixBench-Hypothesis)
 
-- Code and Hypotest dependency: Apache 2.0
-- BixBench-Hypothesis data: CC BY 4.0
+Dependencies
+- nemo_gym: Apache 2.0
+- aviary:  Apache 2.0
+- hypotest: Apache 2.0

environments/aviary_bixbench/README.md

@@ -1,6 +1,6 @@
 # BixBench (Aviary)
 
-This environment exposes [BixBench](https://arxiv.org/abs/2503.00096) through an [Aviary](https://github.com/Future-House/aviary) Jupyter notebook. Docker must be available, and the first run downloads and extracts the BixBench capsules.
+This environment runs [BixBench](https://arxiv.org/abs/2503.00096) through [Aviary](https://github.com/Future-House/aviary). It is a scientific question-answering environment with a Jupyter notebook. Docker must be available, and the first run downloads and extracts the BixBench capsules.
 
 The commands below assume that a model endpoint is configured with `policy_base_url`, `policy_model_name`, and `policy_api_key` in `env.yaml`. See the [local configuration documentation](https://docs.nvidia.com/nemo/gym/reference/configuration#local-configuration-envyaml).
 

environments/aviary_gsm8k/README.md

@@ -1,6 +1,6 @@
 # GSM8K (Aviary)
 
-This environment exposes [GSM8K](https://arxiv.org/abs/2110.14168) through the [Aviary](https://github.com/Future-House/aviary) calculator environment.
+This environment runs [GSM8K](https://arxiv.org/abs/2110.14168) through [Aviary](https://github.com/Future-House/aviary). It is a math question-answering environment with a calculator tool.
 
 The commands below assume that a model endpoint is configured with `policy_base_url`, `policy_model_name`, and `policy_api_key` in `env.yaml`. See the [local configuration documentation](https://docs.nvidia.com/nemo/gym/reference/configuration#local-configuration-envyaml).
 

environments/aviary_hotpotqa/README.md

@@ -1,6 +1,6 @@
 # HotPotQA (Aviary)
 
-This environment exposes [HotPotQA](https://aclanthology.org/D18-1259/) through the [Aviary](https://github.com/Future-House/aviary) Wikipedia-search environment.
+This environment runs [HotPotQA](https://aclanthology.org/D18-1259/) through [Aviary](https://github.com/Future-House/aviary). It is a multi-hop question-answering environment with Wikipedia search.
 
 The commands below assume that a model endpoint is configured with `policy_base_url`, `policy_model_name`, and `policy_api_key` in `env.yaml`. See the [local configuration documentation](https://docs.nvidia.com/nemo/gym/reference/configuration#local-configuration-envyaml).