Redemption
We built an OpenEnv reinforcement learning environment where an agent learns to act as an optimizer for training machine learning models.
Core trick
SAC agent outputs learning rate, momentum, clipping, and weight decay each step from training statistics.
Description
We built an OpenEnv reinforcement learning environment where an agent learns to act as an optimizer for training machine learning models. Instead of using fixed optimization rules like SGD or Adam, the agent observes training statistics such as gradients, loss values, and parameter norms, and decides how model weights should be updated. The environment simulates a training loop where each action modifies the update step, and rewards are based on improving loss reduction while maintaining stability during training. This allows the agent to learn adaptive optimization strategies that change depending on the model state and training phase. Our goal is to explore whether optimization itself can be learned through interaction with a training environment rather than relying on static mathematical formulas.
README excerpt
# Redemption: Meta-Optimizer RL for SLM Training Train an RL agent (SAC) to act as an optimizer: at each step it chooses learning rate, momentum, gradient clipping, and weight decay to minimize the inner task loss. The inner task is either **sinusoidal regression** (MLP) or **next-token prediction** (tiny transformer SLM). Compare the learned meta-optimizer to baselines (AdamW for SLM, Adam/SGD for sinusoid) and view loss curves on the OpenEnv dashboard (e.g. on Hugging Face Spaces). ## Quick Start ```bash # Install (from repo root) cd my_env && uv sync # Train the meta-optimizer (SAC) on 50 SLM tasks uv run --project my_env python scripts/train_sac.py # Run SGD baseline and save loss/perplexity graph uv run --project my_env python scripts/plot_adamw_baseline.py --out sgd_baseline.png # Compare SAC vs AdamW on SLM eval tasks uv run --project my_env python scripts/compare_slm_baseline.py ``` ## Project Structure | Path | Purpose | |------|--------| | `my_env/server/meta_optimizer_environment.py` | Core env: inner SLM/sinusoid step, meta-optimizer action (lr, momentum, clip, wd), reward, baselines (AdamW, SGD, Adam). | | `my_env/server/tasks.py` | Task specs: sinusoid (TaskSpec) and SLM (SLMTaskSpec), 50 train / 2 eval tasks, default corpus. | | `my_env/server/slm_model.py` | Tiny decoder-only transformer (TinyLM) and corpus sampling for SLM. | | `my_env/env_gym.py` | Gymnasium wrapper for SAC: obs (loss, step, grad_norm) → vector, action vector → MetaOptimizerAction. |