refactor(trainer): move trainer modules from experimental to areal/trainer

[garrett4wade]

Description

Move trainer-related modules to establish a cleaner architecture by separating training algorithm concerns from backend infrastructure. This refactoring makes the codebase more modular and simplifies imports - trainers can now be imported directly via from areal import PPOTrainer.

Key changes:

• Move PPOTrainer and SFTTrainer from areal/experimental/trainer/ to areal/trainer/
• Move PPO actor/critic from areal/engine/ppo/ to areal/trainer/ppo/
• Move SFT lm_engine from areal/engine/sft/ to areal/trainer/sft/
• Move RW engine from areal/engine/rw/ to areal/trainer/rw/
• Export PPOTrainer and SFTTrainer from top-level areal package

Related Issue

N/A

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Code refactoring (no functional changes)
• Performance improvement
• Test coverage improvement

Checklist

• I have read the Contributing Guide
• I have run formatting tools (pre-commit or manual)
• I have run relevant unit tests and they pass
• I have added tests for new functionality
• I have updated documentation if needed
• My branch is up to date with main
• This PR introduces breaking changes (if yes, fill out details below)
• If this PR changes documentation, I have built and previewed it locally with jb build docs
• No critical issues raised by AI reviewers (/gemini review)

Breaking Change Details (if applicable):

N/A - Old import paths from areal.experimental.trainer will no longer work, but this is expected as experimental modules are not considered stable API.

Additional Context

Files changed (44 files):

• areal/trainer/ - New trainer module location
• areal/trainer/ppo/ - PPO actor/critic moved from engine
• areal/trainer/sft/ - SFT lm_engine moved from engine
• areal/trainer/rw/ - RW engine moved from engine
• areal/__init__.py - Export PPOTrainer and SFTTrainer
• Updated all examples, tests, and documentation to use new import paths

[gemini-code-assist]

Summary of Changes

Hello @garrett4wade, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly refactors the project's module structure by centralizing all training-related components into a dedicated areal/trainer/ directory. This change aims to improve code organization and simplify the user experience by providing more intuitive and stable import paths for core trainer classes. While enhancing maintainability, this also introduces a breaking change for existing codebases that rely on the old experimental module paths.

Highlights

• Trainer Module Consolidation: All trainer-related code, including PPOTrainer and SFTTrainer, has been moved from the areal/experimental/trainer/ directory to a new, unified areal/trainer/ directory for better organization.
• Engine Component Reorganization: PPO, SFT, and RW engine components have been relocated from areal/engine/ subdirectories into their respective new homes within areal/trainer/ (e.g., areal/engine/ppo/ to areal/trainer/ppo/).
• Simplified Top-Level Imports: PPOTrainer and SFTTrainer are now directly exposed and importable from the top-level areal package, simplifying user import paths from from areal.experimental.trainer import PPOTrainer to from areal import PPOTrainer.
• Breaking Change for Old Imports: The previous import paths for PPOTrainer and SFTTrainer (e.g., from areal.experimental.trainer import PPOTrainer) are no longer valid, requiring users to update their code.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• Core Import Updates
  • areal/__init__.py: PPOTrainer and SFTTrainer are now imported from areal.trainer and exposed at the top-level areal package.
  • areal/engine/fsdp_engine.py, areal/engine/megatron_engine.py, areal/experimental/engine/archon_engine.py: Updated internal imports for PPO, SFT, and RW engine components to reference their new locations under areal/trainer/.
• Documentation and Example Updates
  • .claude/agents/algorithm-expert.md, .claude/data/pr-review-change-types.md, .claude/hooks/check-expert-update.sh, AGENTS.md: Updated references to trainer and engine paths.
  • areal/models/fsdp/ulysses.py: Updated comment reference for ppo_loss_fn.
  • areal/tests/experimental/archon/test_grpo.py, areal/tests/grpo/entrypoint.py, areal/tests/sft/entrypoint.py, areal/tests/test_prox_approx.py: Updated import statements for trainer and PPO-related modules.
  • docs/customization/algorithm.md: This documentation file was removed.
  • docs/customization/dataset.md, docs/reference/metrics_tracking.md, docs/tutorial/agentic_rl.md, docs/tutorial/gsm8k_grpo.md: Updated code snippets and links to reflect the new module structure and import paths.
  • All examples/ files (e.g., examples/camel/train.py, examples/math/gsm8k_rl.py, etc.): Updated import statements for PPOTrainer and SFTTrainer to use the new top-level from areal import ... path.
• Module Structure Changes
  • areal/experimental/trainer/__init__.py was removed.
  • areal/engine/ppo/actor.py was renamed to areal/trainer/ppo/actor.py.
  • areal/engine/ppo/critic.py was renamed to areal/trainer/ppo/critic.py.
  • areal/experimental/trainer/rl.py was renamed to areal/trainer/rl_trainer.py.
  • areal/engine/rw/rw_engine.py was renamed to areal/trainer/rw/rw_engine.py.
  • areal/engine/sft/lm_engine.py was renamed to areal/trainer/sft/lm_engine.py.
  • areal/experimental/trainer/sft.py was renamed to areal/trainer/sft_trainer.py.
  • areal/tests/experimental/trainer/test_workflow_detection.py was renamed to areal/tests/test_workflow_detection.py.
  • New package areal/trainer/ was created with __init__.py, ppo/__init__.py, rw/__init__.py, and sft/__init__.py to define the new module hierarchy and export components.

Activity

• The pull request was authored by garrett4wade.
• It is categorized as a breaking change and a code refactoring, indicating significant structural modifications.
• The author has confirmed reading the contributing guide, running formatting tools, updating documentation, and ensuring the branch is up to date with main.
• The changes involve moving and renaming numerous files, creating new __init__.py files for the new module structure, and updating import paths across the codebase, including documentation and examples.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request is a significant refactoring that moves trainer-related modules from the experimental directory to a new areal/trainer directory. This is a good move to consolidate code and promote these modules to the main package. The changes are extensive, touching many files to update import paths and documentation. The new top-level imports for PPOTrainer and SFTTrainer simplify usage.

My review found a few minor issues, mostly incorrect file paths in documentation and one incorrect import in a test file, which are easy to miss in such a large refactoring. Overall, the changes are well-executed and improve the project structure.

[rchardx]

examples/tau2/train.py:13
This file still uses the old import path from areal.experimental.trainer.rl import PPOTrainer which will break after this PR is merged. The file was NOT included in the PR changes.

[rchardx]

docs/tutorial/gsm8k_grpo.md:610
Contains link Custom Algorithms but docs/customization/algorithm.md was deleted in this PR. This will result in a 404 error in the documentation.

[rchardx]

docs/tutorial/quickstart.md:158
Contains link Custom algorithms but the file was deleted.

[rchardx]

LGTM