# Atropos - Nous Research's LLM RL Gym

![newatr-02](banner-image.jpg)

<div align="center">

*In Greek mythology, Atropos was the eldest of the three Fates.  While her sisters spun and measured the threads of mortal lives, Atropos alone held the shears that would cut these threads, determining the final destiny of each soul. Just as Atropos guided souls to their ultimate fate, this system guides language models toward their optimal potential through reinforcement learning.*

</div>

<div align="center">
</div>
<div id="badges" align="center">
  <a href="https://huggingface.co/NousResearch">
    <img src="https://img.shields.io/badge/NousResearch-orange?style=for-the-badge&logo=huggingface&logoColor=white" alt="HuggingFace"/>
  </a>
  <a href="https://nousresearch.com">
    <img src="https://img.shields.io/badge/NousResearch.com-white?style=for-the-badge&logo=data:image/png;base64, " alt="Website"/>
  </a>
  <a href="https://x.com/NousResearch">
    <img src="https://img.shields.io/badge/@NousResearch-black?style=for-the-badge&logo=X&logoColor=white" alt="@NousResearch"/>
  </a>
</div>

---

## What is Atropos?
Atropos is an environment microservice framework for async RL with LLMs.

Atropos encompasses both environments, which are set up as services, and a trajectory API for the environments to send data to and for the trainer to pull batches from.

![image](https://github.com/user-attachments/assets/8ce52994-b219-49d6-970c-58a477f36151)
<div align="center">

  *Here is a diagram of how Atropos' components can interact with a trainer & inference server to complete the RL loop (trainer & inference engine not included with the atropos package)*

</div>

Atropos is a robust, scalable framework for **Reinforcement Learning Environments with LLMs**.

The goal: provide a flexible, scalable, and standardized platform to accelerate LLM-based RL research across diverse, interactive settings.

The framework supports collecting, distributing and evaluating LLM trajectories through diverse environments including:

<div align="center">

| Environment Type          | Examples                                   | Purpose                                            |
|---------------------------|--------------------------------------------|----------------------------------------------------|
| 📚 Dataset environments   | GSM8K, MMLU, Custom HF Datasets            | Evaluate and improve LLM performance on static data|
| 🎮 Online environments    | Blackjack, Taxi, Text-based games          | Train LLMs through interactive game-based learning |
| 🤖 RLAIF and RLHF         | LLM Judge/Reward Models                    | Fine-tune LLMs using human feedback and alignment  |
| 🔄 Multi-Turn RL          | deepresearch, internal tool calling        | Train LLMs on complex multi-step interactions      |
| 💻 Code Execution         | MBPP, HumanEval (via `coding_server.py`)   | Train LLMs to generate and execute code            |
| 🖼️ Multimodal             | OCR VQA, Clevr (via `multimodal_dpo/`)     | Train LLMs on tasks involving vision and language  |

</div>

---

## Experimental results from models trained using Atropos' environments

We have been able to achieve significant improvements on specific domains or tasks with Atropos - Below are some of the results.

**Tool Calling Environment Results:**

<div align="center">

| Berkeley Function Calling Benchmark Type | Base Model | With Atropos RL | Improvement |
|---------------|------------|-----------------|-------------|
| Parallel Tasks| 10%        | 46%            | **4.6x** ⬆️ |
| Simple Tasks  | 21%        | 51.75%         | **2.5x** ⬆️ |

</div>

Model Artifact:
https://huggingface.co/NousResearch/DeepHermes-ToolCalling-Specialist-Atropos


Environment Used:
[https://github.com/NousResearch/atropos/blob/main/environments/tool_calling_server.py](https://github.com/NousResearch/atropos/blob/main/environments/tool_calling_server.py)

---

**Financial Fundamentals Prediction Environment Results**:

<div align="center">

| Metric | Initial Accuracy | With Atropos RL | Improvement |
|--------|-----------------|-----------------|-------------|
| Directional Prediction Eval Accuracy | 20% | 50% | **2.5x** 📈 |

</div>

Model Artifact:
https://huggingface.co/NousResearch/DeepHermes-Financial-Fundamentals-Prediction-Specialist-Atropos

Environment Used:
[https://github.com/NousResearch/atropos/blob/main/environments/fundamental_prediction_environment.py](https://github.com/NousResearch/atropos/blob/main/environments/fundamental_prediction_environment.py)

---

## RLAIF Experiment Artifacts
Using the RLAIF Environment to change the personality of the model, we have produced several artifacts of interesting and weird personalities.

**DeepHermes Egregore v1 and v2 8B:**

https://huggingface.co/NousResearch/DeepHermes-Egregore-v1-RLAIF-8b-Atropos
https://huggingface.co/NousResearch/DeepHermes-Egregore-v2-RLAIF-8b-Atropos

**DeepHermes Ascension Maze 8B:**

https://huggingface.co/NousResearch/DeepHermes-AscensionMaze-RLAIF-8b-Atropos

Environment Used: [https://github.com/NousResearch/atropos/blob/main/environments/rlaif_server.py](https://github.com/NousResearch/atropos/blob/main/environments/rlaif_server.py)

---

## Navigating the Repo

| Category | Description |
|----------|------------|
| 📁 [`atroposlib/`](atroposlib/) | Core library containing base classes and utilities |
| 🎮 [`environments/`](environments/) | Collection of ready-to-use RL environments. Community contributions are typically placed in the [`environments/community/`](environments/community/) subdirectory. |
| 📚 [`example_trainer/`](example_trainer/) | Example training scripts and configurations |

Key Documents:
- [Base Environment Class](atroposlib/envs/README.md) - Documentation for creating custom environments
- [ManagedServer Guide](atroposlib/envs/server_handling/MANAGED_SERVER.md) - **Recommended approach** for automatic token and logprob tracking
- [Environments Overview and Contribution Guide](environments/community/README.md) - Documentation for existing environments and how to contribute new ones.
- [Full Environment Config Options](CONFIG.md) - Documentation for creating custom environments
- [Example Trainer](example_trainer/README.md) - Getting started with training
- [Slurm Guide](SLURM.md) - Guide for using Atropos with Slurm for distributed inference
- [Frequently Asked Questions (FAQ)](atroposlib/FAQ.md) - Answers to common questions for new users
- [Contributing Guide](CONTRIBUTING.md) - Guidelines for contributors
- [License](LICENSE) - MIT license details

---

## Prerequisites

Before installing Atropos, ensure you have the following:

- **Python 3.10+** — Required. Check with `python --version`
- **Git** — For cloning the repository
- **An OpenAI-compatible API endpoint** — Atropos environments need an inference server. Options include:
  - A local [vLLM](https://github.com/vllm-project/vllm) or [SGLang](https://github.com/sgl-project/sglang) instance
  - An [OpenAI API key](https://platform.openai.com/api-keys) (set as `OPENAI_API_KEY` environment variable)
  - Any provider with an OpenAI-compatible endpoint (e.g., [Together AI](https://together.ai), [OpenRouter](https://openrouter.ai))
- **Weights & Biases account** *(optional)* — For experiment tracking. Set `use_wandb=False` in your environment config to skip

> **Note:** You do not need a GPU to develop or test environments locally. A GPU is only required for running inference servers locally or for training.

---

## Installation

Get your Python 3.10 (or later) environment ready, then simply pip install:

```bash
pip install atroposlib
```

If you're looking to get into developing the repo or using the environments:


```bash
pip install -e .                # for using
pip install -e .[dev]           # for development
pip install -e .[examples]      # for running examples
pip install -e .[verifiers]     # for verifiers integration
pip install -e .[all]           # for everything
```

**Important:** If you're committing to the repository, please install the pre-commit hooks:
```bash
pre-commit install
```

---

### Quick Start Guide

1. **Create Your First Environment**
   - Review our [Base Class Documentation](atroposlib/envs/README.md) to understand the core concepts
   - Check out existing environments in the [`environments/`](environments) directory for examples

2. **Run an Example Environment**

  You should edit the config_init section of the environment file you want ([For example, in GSM8K Environment](https://github.com/NousResearch/atropos/blob/main/environments/gsm8k_server.py#L53)) to point to a running VLLM or SGLang inference server as well as any other [configuration changes](CONFIG.md) you'd like to make, such as the group size, then:

  > **Note:** By default, Atropos uses the OpenAI-compatible API endpoint which works with any provider. For enhanced features, use `VLLMServer` (atroposlib/envs/server_handling/vllm_server.py) or `SGLangServer` (atroposlib/envs/server_handling/sglang_server.py) for direct access to native APIs with full token and logprob tracking.

   ```bash
   # Start the API server
   run-api
   ```
   In a separate terminal, start the GSM8K environment microservice
   ```bash
   python environments/gsm8k_server.py serve --openai.model_name Qwen/Qwen2.5-1.5B-Instruct --slurm false
   # alternatively
   # python environments/gsm8k_server.py serve --config environments/configs/example.yaml
   # python environments/gsm8k_server.py serve --config environments/configs/example.yaml --env.group_size 8 # cli args override corresponding config settings
   ```
3. **Grabbing Rollouts**

  If you want to just start getting rollouts, and not use a trainer, see the [debug section](#testing-and-debugging-tools)
  for help getting started with the available tools, we recommend starting with process or view-run

4. **Training Your Model**
   - Follow our [training example guide](example_trainer/README.md) for detailed instructions
   - Monitor progress through our built-in logging and reporting system:
     - Completion lengths
     - Evaluation accuracies
     - Full rollouts and scores

You can use multiple environments at once, just point them all to the same server.

Environments come with detailed logging and reporting support, runs track completion lengths, eval accuracies, full rollouts and scores, and more:

![image](https://github.com/user-attachments/assets/153a2932-191a-42e3-8da9-25a1b05abb8e)

---

# Trainer Integrations
## Axolotl
<a href="https://github.com/axolotl-ai-cloud/plugin-atropos">
  <img
    src="https://github.com/user-attachments/assets/be629253-a8b1-4354-b6da-5e404e9c854d"
    alt="Atropos plugin logo"
    width="50%">
</a>

Axolotl is a powerful tool for fine-tuning a wide range of AI models, supporting techniques like LoRA and QLoRA through simple YAML configurations.

The [Atropos plugin for Axolotl](https://github.com/axolotl-ai-cloud/plugin-atropos) seamlessly integrates Atropos' RL environments into Axolotl's training pipelines.
This allows you to leverage Atropos for reinforcement learning while utilizing Axolotl's extensive features for model fine-tuning.

To use, follow the README on the [plugin repository](https://github.com/axolotl-ai-cloud/plugin-atropos).

## Tinker
<a href="https://github.com/NousResearch/tinker-atropos">
  <img
    src="https://github.com/user-attachments/assets/6c226187-4df8-4094-be5d-72f3f58de423"
    alt="Atropos Tinker logo"
    width="50%">
</a>

The Tinker API is a simple and flexible LoRA trainer framework for researchers and developers to quickly build out their ideas without worrying about the complexities of distributed training. Users write a simple loop that runs on their CPU, and Tinker manages the backend computation on their GPUs, while still providing full control over the training and algorithmic details.

The [Tinker-Atropos](https://github.com/NousResearch/tinker-atropos) integration layer enables all Atropos environments to leverage the power of Tinker for their RL experiments. This allows users with little or no compute to develop and build Atropos environments with minimal worry about the underlying compute behavior, as well as providing an easy environment integration point for Tinker users.

To get started, check out the README at the [project repository](https://github.com/NousResearch/tinker-atropos).

## Atropos' Example Trainer
Atropos repo contains an example trainer that should primarily be used as a reference example to show how a trainer and inference provider can be integrated with Atropos to complete the RL Training Loop.

To use the example trainer, see this page: [training example guide](example_trainer/README.md)

## On-Policy Distillation (API + ScoredDataGroup Contract)

Atropos now supports OPD at the transport layer by carrying distillation arrays
through `ScoredDataGroup` and the API queue/batch endpoints.

### Scope of this change

- No teacher fetching/orchestration in `BaseEnv`.
- Environments or external pipelines are responsible for populating distillation arrays.
- API stores and returns those arrays unchanged.

### Distillation payload fields

Each scored group may include:

- `distill_token_ids`: shape `[sequence][position][top_k]`
- `distill_logprobs`: shape `[sequence][position][top_k]`

These fields are optional, and when present are forwarded from:

- environment -> `/scored_data` or `/scored_data_list`
- API queue -> `/batch` -> trainer

### Minimal producer example (environment side)

```python
scores["distill_token_ids"] = distill_token_ids
scores["distill_logprobs"] = distill_logprobs
```

### Minimal consumer check (trainer/debug side)

```bash
curl -s http://localhost:8002/latest_example | jq '{has_ids:(.distill_token_ids!=null), has_lps:(.distill_logprobs!=null)}'
```

### Notes

- The API does not validate cross-field semantics beyond schema typing.
- Trainers should validate alignment assumptions they require (sequence length, per-position top-k, etc.).
- Teacher-side architecture and prompt/rendering strategy are intentionally out of scope for this PR.

### TeacherDistillationEnv follow-up

The follow-up teacher environment uses a dedicated teacher server config and
attaches teacher prompt logprobs before the group is sent to the API.

Teacher config shape:

```python
TeacherDistillationConfig(
    teacher_enabled=True,
    teacher_top_k=8,
)
```

Teacher server configs are passed separately at init, just like the primary
`server_configs`:

```python
env = MyTeacherEnv(
    config=env_config,
    server_configs=student_server_configs,
    teacher_server_configs=[
        APIServerConfig(
            base_url="http://localhost:9003/v1",
            model_name="Qwen/Qwen3-30B-A3B-Instruct-2507",
            api_key="",
            server_type="vllm",
            tokenizer_name="Qwen/Qwen3-30B-A3B-Instruct-2507",
        )
    ],
)
```

You can either:

- build a teacher-enabled env by mixing `TeacherDistillationEnv` into an existing
  `BaseEnv`-derived env such as `GSM8kEnv`, or
- subclass `TeacherDistillationEnv` directly and implement the usual environment
  methods yourself.

In both cases, `TeacherDistillationEnv` still assumes the normal `BaseEnv`
runtime contract: tokenized rollouts, `ScoredDataGroup` payloads, and the
standard `handle_send_to_api(...)` transport path.

CLI shape:

```bash
--env.teacher_enabled true \
--teacher.base_url "http://localhost:9003/v1" \
--teacher.model_name "Qwen/Qwen3-30B-A3B-Instruct-2507" \
--teacher.server_type vllm \
--env.teacher_top_k 8
```

If `--teacher.model_name` is a deployment alias rather than a tokenizer
identifier, also set `--teacher.tokenizer_name ...` so the env can validate
tokenizer compatibility.

Scope note:

- The teacher-aware CLI wiring currently exists for `serve`.
- If `teacher_enabled=True`, the generic `process` and `evaluate` commands will
  fail loudly at env construction time unless you instantiate the env yourself
  and pass `teacher_server_configs=...`.

Tokenizer requirement:

- Teacher distillation currently requires the teacher and student to use the same tokenizer vocabulary.
- If the tokenizers do not match, `TeacherDistillationEnv` raises an error instead of attempting token conversion.

Why same-tokenizer is required:

- `distill_token_ids` are consumed as student-vocabulary IDs by the trainer.
- If the teacher uses a different vocabulary, the same integer token ID refers to different text on the teacher and student sides.
- A decode/re-tokenize/remap pipeline is not a safe drop-in fix because it changes both token positions and token identities, which breaks the exact per-position token supervision that the current distillation loss assumes.

---

## Testing and Debugging Tools

The trajectory-handler provides several debugging tools to help environment developers test and understand their environments locally without requiring the full distributed infrastructure.

*   **Flexible Model Provider Support:** Atropos natively supports any model provider that adheres to the OpenAI API standard. Simply provide the provider's base URL and your API key, and Atropos can integrate with their models seamlessly for testing or running environments locally.

After launching the API and your selected environments (e.g. `run-api & python environments/gsm8k_server.py serve`), you are then able to view them to get a quick look, or try to prepare some datasets for some offline training:

*   **View Run (`view-run`):** Launch a Gradio UI to inspect batches of rollouts generated by your environment runs. This is useful for visually debugging the interactions and data flow.
*   **Offline Data Generation:** Use `atropos-sft-gen` and `atropos-dpo-gen` to collect rollouts from environments and convert them into formats suitable for Supervised Fine-Tuning (SFT) or Direct Preference Optimization (DPO).

### In-depth Local Environment Analysis with `process`

For developers looking to inspect and debug a single environment without the overhead of the `run-api` server or a full training loop, Atropos environments offer a `process` subcommand. This mode performs inference-only rollouts, meaning it runs your model within the environment to generate interactions, but does not perform any model training or updates.

The `process` subcommand executes the environment's full data pipeline:

1.  **Generation:** Produces model responses based on inputs from the environment.
2.  **Parsing:** Processes these raw model outputs into a structured format.
3.  **Scoring:** Applies the environment's reward logic to evaluate the quality of the generated responses.

**Outputs and Visualization:**

When you specify a path to save the generated data using the `--env.data_path_to_save_groups your_output_file.jsonl` argument (or a similar argument defined by the specific environment, check with `--help`), the `process` command provides several benefits:

*   **JSONL Output:** Saves all generated rollout groups, including prompts, responses, and scores, to the specified `.jsonl` file. This data can be useful for detailed offline analysis and debugging.
*   **Static HTML Visualization:** Automatically generates a corresponding `.html` file (e.g., `your_output_file.html`) that provides a user-friendly, browser-based view of the rollouts contained in the JSONL file. This is excellent for quickly understanding model behavior and identifying issues.
*   **WandB Logging:** If Weights & Biases (`use_wandb=True`) is enabled in your environment's configuration, the `process` subcommand will also log the run data, metrics, and generated rollouts to your WandB dashboard, allowing for persistent tracking and comparison even for these inference-only runs.

**Example Usage:**

To run the `process` subcommand for an environment like `gsm8k_server.py` and save the outputs:

```sh
python environments/gsm8k_server.py process --env.data_path_to_save_groups gsm8k_rollouts.jsonl
```

This will create `gsm8k_rollouts.jsonl` and `gsm8k_rollouts.html`.

**Customization:**

You can customize the inference endpoint and other parameters for the `process` subcommand. For example, to use a different model or API endpoint:

```sh
python environments/gsm8k_server.py process \
  --env.data_path_to_save_groups gsm8k_rollouts.jsonl \
  --env.my_custom_field "value" \
  --openai.base_url https://your-custom-api-url/v1 \
  --openai.api_key YOUR_API_KEY \
  --openai.model_name your_model_identifier
```

You can add custom fields to the `env` namespace by returning a custom subclass of BaseEnvConfig in `config_init` [[example](https://github.com/NousResearch/atropos/blob/bdb15e5d85ddcf8a6ede352977719df442e60a22/environments/math_server.py#L181)].

Always refer to the specific environment script's help for all available options:

```sh
python environments/your_environment_script.py process --help
```

### Environment Evaluation with `evaluate`

For running evaluation on environments, Atropos provides an `evaluate` subcommand that calls the environment's `evaluate` method:

```sh
python gsm8k_server.py evaluate \
  --openai.base_url https://openrouter.ai/api/v1 \
  --openai.api_key $OPENROUTER_API_KEY \
  --openai.model_name qwen/qwen3-14b
```

### Offline Data Generation Quick Start

Run the following commands in **separate terminals**, in this order:

**Terminal 1** — Start the API server first (must be running before environments connect):
```sh
run-api
```

**Terminal 2** — Start an environment:
```sh
python gsm8k_server.py serve --slurm False # or an env of your choice
```

**Terminal 3** — Generate data:
```sh
atropos-sft-gen path/to/output.jsonl --tokenizer Qwen/Qwen2.5-1.5B-Instruct # or whichever tokenizer you have in your env config
```
Rejection sampling can be controlled via `--save-top-n-per-group`, `--allow-negative-scores`, and `--minimum-score-diff-max-min`. See `atropos-sft-gen -h` for more detailed usage info.

If you would like to use OpenAI models, please edit your `config_init` to something like the following:
```python
    @classmethod
    def config_init(cls) -> Tuple[BaseEnvConfig, List[APIServerConfig]]:
        env_config = BaseEnvConfig(
            tokenizer_name="Qwen/Qwen2.5-1.5B-Instruct",
            group_size=8,
            use_wandb=True,
            rollout_server_url="http://localhost:8000",
            total_steps=1000,
            batch_size=12,
            steps_per_eval=100,
            max_token_length=2048,
            wandb_name="gsm8k",
        )
        server_configs = [
            APIServerConfig(
                model_name="gpt-4.1-nano",
                base_url=None,
                api_key=os.environ.get("OPENAI_API_KEY"),
                num_requests_for_eval=256,
            ),
        ]

        return env_config, server_configs
```

For DPO, replace `atropos-sft-gen` with `atropos-dpo-gen` and check `atropos-dpo-gen -h` for data filtering and saving options.

---

## Troubleshooting

**`Address already in use` when running `run-api`**

Port 8000 is already occupied. Either stop the existing process or specify a different port:

```bash
# Find and stop the process using port 8000
lsof -ti:8000 | xargs kill -9

# Or use a different port
run-api --port 8001
```

**`ModuleNotFoundError` or dependency conflicts**

Ensure you're using a clean virtual environment with the correct Python version:

```bash
python -m venv .venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate
pip install -e ".[dev]"
```

**`OPENAI_API_KEY` not set errors**

Set your API key as an environment variable, or configure it in the environment's `config_init`:

```bash
export OPENAI_API_KEY="your-key-here"
```

**Out of memory (OOM) when running environments locally**

Use a smaller model for local development and testing. For example, configure `model_name` to a lightweight model like `gpt-4.1-nano` with an OpenAI API key, or use a quantized local model with vLLM.

**Environment not connecting to the API server**

Ensure `run-api` is running before starting any environments. By default, environments connect to `http://localhost:8000`. If your API is on a different host or port, update `rollout_server_url` in your environment's config.

---

## Citation

If you have found the library helpful in your work, you can cite this repository as:

```latex
@misc{atropos,
  title        = {Atropos: An Async First Environment Rollout Controller},
  author       = {Mahan, Dakota and Jin, Roger and Teknium and Sands, Shannon and Yatsenko, Artem and Suphavadeeprasit, Jai and Malhotra, Karan and Guang, Chen and Li, Joe},
  howpublished = {\url{https://www.github.com/NousResearch/atropos}},
  year         = {2025},
  month        = {apr},
  note         = {Version 0.3.0},
}
```

---

## Contributing

Atropos is built by the open-source AI community, and relies on our amazing contributors! Please see our [contributing](CONTRIBUTING.md) guide for more details on our code formatting, testing, etc.
Please follow the [Code of Conduct](CODE_OF_CONDUCT.md).

---

## License
Atropos uses the MIT license, see the [LICENSE](LICENSE) file here for more information
