Browse the docs · Apsi 2.0.0

Reinforcement learning

use rl provides environments, advantage estimation, policy helpers, complete training algorithms (PPO, DQN, REINFORCE), the losses used to train language models from human preferences (RLHF, DPO, IPO, GRPO), and search algorithms (Monte Carlo tree search, beam search, best-of-n). Everything runs on the CPU without external libraries and works with any model written with use torch or use layers.

A complete example: PPO on CartPole

use torch
use rl

let l1 = torch.Linear(4, 64)
let l2 = torch.Linear(64, 64)
let pi = torch.Linear(64, 2)
let vf = torch.Linear(64, 1)
let params = record(l1=l1.init(1), l2=l2.init(2), pi=pi.init(3), vf=vf.init(4))

# the model returns [action logits, state value]
fn policy(p, obs):
    let h = tanh(l2(p.l2, tanh(l1(p.l1, obs))))
    return [pi(p.pi, h), vf(p.vf, h)]

let result = rl.ppo(rl.cartpole(8, seed=0), policy, params, total_steps=200000, epochs=10, lr=1e-3)
let score = rl.evaluate(rl.cartpole(1, seed=5), fn(obs) => argmax(policy(result.params, obs)[0], axis=1), episodes=5)
print(f"mean return: {score.mean}")        # 500 is the maximum

With these settings PPO reaches the maximum return of 500 on CartPole-v1 in about 200 000 environment steps.

Environments

Function	Task
`cartpole(n=1, seed, max_steps=500)`	Balance a pole on a cart. Dynamics identical to Gymnasium’s CartPole-v1.
`gridworld(size=5, n=1, seed)`	Walk to the goal on a grid; each step costs a little reward.
`bandit(arms=10, n=1, seed)`	Multi-armed Gaussian bandit.

All environments are vectorised: n copies run side by side. reset(seed) returns observations of shape [n, …]; step(actions) returns record(obs, reward, done, truncated, final_obs). Finished copies reset themselves, and final_obs holds the last observation before the reset. Completed episode returns are collected in env.episode_returns.

Your own environment only needs the same members: n, reset(seed), step(actions) and episode_returns.

Algorithms

`ppo(env, model, params, …)`

Proximal Policy Optimization with the clipped surrogate objective, generalised advantage estimation, a value loss, an optional entropy bonus, advantage normalisation, global gradient-norm clipping and linear learning-rate annealing.

Argument	Default	Meaning
`total_steps`	100 000	environment steps across all copies
`rollout`	128	steps per copy between updates
`epochs`	4	passes over each rollout
`minibatch`	256	samples per gradient step
`lr`	3e-4	Adam learning rate
`gamma`, `lam`	0.99, 0.95	discount and GAE parameter
`clip`	0.2	ratio clipping range
`vf_coef`, `ent_coef`	0.5, 0.0	weights of the value loss and entropy bonus
`max_grad_norm`	0.5	gradient clipping
`anneal_lr`	true	decay the learning rate linearly to 0
`opt`	Adam	any optimizer from `use torch`

model(params, obs) returns [logits, value]. The result is record(params, episode_returns, mean_returns, updates).

`dqn(env, model, params, …)`

Double DQN with a replay buffer, a target network, ε-greedy exploration and the Huber loss. model(params, obs) returns one value per action.

`reinforce(env, model, params, episodes, …)`

Monte Carlo policy gradient with a running baseline.

`evaluate(env, policy, episodes=10, seed=0)`

Runs a deterministic policy (policy(obs) returns actions) and returns record(mean, returns).

Building blocks

Function	Purpose
`discounted_returns(rewards, dones, gamma, last_value)`	discounted returns that stop at episode ends
`gae(rewards, values, dones, last_values, gamma, lam)`	generalised advantage estimation; returns `record(advantages, returns)`
`categorical_sample(logits, seed)` · `categorical_log_prob(logits, actions)` · `categorical_entropy(logits)`	discrete policies
`gaussian_log_prob(mean, log_std, actions)` · `gaussian_entropy(log_std)`	continuous policies
`sample_token(logits, temperature, top_k, top_p, seed)`	sampling for language models

Learning from preferences (RLHF, DPO, GRPO)

Function	Purpose
`preference_loss(r_chosen, r_rejected, margin)`	Bradley–Terry loss for training a reward model
`dpo_loss(policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta, label_smoothing)`	Direct Preference Optimization; inputs are sequence log-probabilities; returns `record(loss, chosen_rewards, rejected_rewards, accuracy)`
`ipo_loss(…, tau)`	Identity Preference Optimization
`kl_penalized_rewards(rewards, logp, ref_logp, beta)`	per-token rewards with a KL penalty against a reference model (RLHF with PPO)
`grpo_advantages(rewards, group_size)`	group-relative advantages: several answers per prompt, normalised within each group
`ppo_clip_loss(logp, old_logp, advantages, clip, mask)`	the clipped surrogate for token sequences
`sequence_log_prob(logits, tokens, mask)`	log-probability of each sequence, differentiable

A DPO step for a causal language model:

fn dpo_step(p):
    let pc = rl.sequence_log_prob(model(p, chosen), chosen_targets, chosen_mask)
    let pr = rl.sequence_log_prob(model(p, rejected), rejected_targets, rejected_mask)
    return rl.dpo_loss(pc, pr, ref_chosen, ref_rejected, beta=0.1).loss

let loss, grads = value_and_grad(dpo_step)(params)
params = opt.step(params, grads)

ref_chosen and ref_rejected are the same log-probabilities computed once with the frozen reference model.

Search

Function	Purpose
`mcts(state, legal_actions, apply, is_terminal, reward, simulations, c, rollout_depth, seed, two_player)`	Monte Carlo tree search with UCT; returns `record(action, visits, values)`
`beam_search(step, start, beam_width, max_len, eos)`	keeps the best partial sequences; `step(seq)` returns log-probabilities of the next token; returns a list of `record(tokens, score)`, best first
`best_of_n(candidates, score)`	the highest-scoring candidate (rejection sampling); returns `record(best, index, scores)`

Memory and speed

Train in float32 with fit(…, precision="fp32") or by creating parameters with dtype=f32; gradients and optimizer state stay float32, which halves memory.
Wrap expensive blocks with torch.checkpoint_layer(layer, params, x) to recompute activations in the backward pass instead of storing them: roughly a third of the peak memory for about twice the compute.
Run many environment copies (n=8 or more) so each model call processes a batch.

Found a mistake? Report it · Markdown source