Efficient Reinforcement Learning for Video Models

# Illustrative sketch (not from the paper)
import torch

# placeholder model that can sample a rollout given a video query
def sample_rollout(model, query):
    # returns a list of tokens (or actions) representing a generated answer
    return model.generate(query)

# placeholder oracle rollout constructed from ground‑truth annotation
def oracle_rollout(annotation):
    return annotation  # already a valid rollout

# scalar task reward for any rollout
def compute_reward(rollout):
    return torch.tensor([reward_fn(rollout)])  # reward_fn defined elsewhere

# ---- Core OraRL update for a single query ----
query, annotation = get_batch()  # obtain video query and its annotation
policy = model

# 1. generate on‑policy rollouts (GRPO style)
on_policy = [sample_rollout(policy, query) for _ in range(N)]
on_rewards = torch.stack([compute_reward(r) for r in on_policy])

# 2. add oracle rollout to the group
oracle = oracle_rollout(annotation)
oracle_reward = compute_reward(oracle)

# 3. decoupled advantage: baseline = mean(on‑policy rewards)
baseline = on_rewards.mean()
# directional gain = oracle_reward - baseline
directional_gain = oracle_reward - baseline

# 4. detached oracle advantage (bounded by strongest on‑policy reward)
max_on = on_rewards.max()
oracle_adv = torch.clamp(directional_gain, max=max_on)

# 5. sign‑balanced pruning: keep oracle + equal #positive / #negative on‑policy rollouts
pos = [r for r in on_rewards if r > baseline]
neg = [r for r in on_rewards if r < baseline]
balanced = pos[:len(pos)//2] + neg[:len(neg)//2]
# combine with oracle
selected_rewards = torch.stack([oracle_reward] + balanced)

# 6. moment correction: enforce zero‑mean advantage
advantages = selected_rewards - selected_rewards.mean()

# 7. policy gradient update (GRPO style) using detached oracle advantage
loss = -(advantages * torch.log(policy.probability(selected_rewards))).mean()
loss.backward()
optimizer.step()

// Illustrative sketch (not from the paper)
const torch = require('torch-js'); // placeholder for tensor ops

// placeholder model that can sample a rollout given a video query
function sampleRollout(model, query) {
  // returns an array representing a generated answer
  return model.generate(query);
}

// oracle rollout constructed from ground‑truth annotation
function oracleRollout(annotation) {
  return annotation; // already a valid rollout
}

// scalar task reward for any rollout
function computeReward(rollout) {
  return torch.tensor([rewardFn(rollout)]); // rewardFn defined elsewhere
}

// ---- Core OraRL update for a single query ----
const { query, annotation } = getBatch(); // obtain video query and its annotation
const policy = model;

// 1. generate on‑policy rollouts (GRPO style)
const onPolicy = [];
for (let i = 0; i < N; i++) {
  onPolicy.push(sampleRollout(policy, query));
}
const onRewards = torch.stack(onPolicy.map(r => computeReward(r)));

// 2. add oracle rollout to the group
const oracle = oracleRollout(annotation);
const oracleReward = computeReward(oracle);

// 3. decoupled advantage: baseline = mean(on‑policy rewards)
const baseline = onRewards.mean();
// directional gain = oracleReward - baseline
const directionalGain = oracleReward.sub(baseline);

// 4. detached oracle advantage (bounded by strongest on‑policy reward)
const maxOn = onRewards.max();
const oracleAdv = torch.clamp(directionalGain, { max: maxOn });

// 5. sign‑balanced pruning: keep oracle + equal #positive / #negative on‑policy rollouts
const pos = onRewards.filter(r => r.gt(baseline));
const neg = onRewards.filter(r => r.lt(baseline));
const balanced = pos.slice(0, Math.floor(pos.length / 2)).concat(
  neg.slice(0, Math.floor(neg.length / 2))
);
const selectedRewards = torch.stack([oracleReward].concat(balanced));

// 6. moment correction: enforce zero‑mean advantage
const advantages = selectedRewards.sub(selectedRewards.mean());

// 7. policy gradient update (GRPO style) using detached oracle advantage
let loss = advantages.mul(policy.logProb(selectedRewards)).neg().mean();
loss.backward();
optimizer.step();