DICE-RL residual critic — MimicGen coffee_d0 (arm F2)
The learned action-value function Q(s, a) from a residual reinforcement-learning run on MimicGen
coffee_d0. A frozen flow-matching multi-task DiT policy proposes an action chunk; a small residual
actor corrects it; this ensemble of 10 critics scores the result. Only the critic is documented here
— the actor is present in the checkpoint but is not the point of this upload.
Trained for 300,003 environment steps / 38,019 gradient updates.
What it takes and what it returns
| state | 66-dim observation.state, normalised (see below) |
| action | one 8-step chunk of 7-dim actions, flattened row-major to 56 dims (flat[t*7 + d]) |
| output | 10 scalar values; the conservative reduction is min over the ensemble |
| architecture | 10 × Linear(122,256) → LayerNorm → Mish ×3 → Linear(256,1) |
State normalisation uses q01/q99 from the coffee_d0 demonstration dataset, shipped here as
state_scaler.json:
scaled = clip(2 * (s - q01) / max(q99 - q01, 1e-6) - 1, -5, 5)
Feeding raw, unnormalised states silently produces meaningless values — nothing errors.
Reward convention (this is what Q means)
gamma = 0.99, step_penalty = 0.0, tau = 0.01. The only reward is +1 on reaching the
success frame. So the quantity Q estimates has a closed form on any recorded trajectory:
G_t = 0.99 ** (frames from t to the first success frame) if the episode succeeds
G_t = 0 if it does not
Q is therefore a discounted time-to-success estimate, bounded in roughly [0, 1]. Values are not
comparable to a critic trained with a nonzero step penalty.
Usage
import json, torch
ckpt = torch.load("critic_300k.pt", map_location="cpu", weights_only=False)
print(ckpt["config"]["gamma"], ckpt["env_steps"]) # 0.99 300003
# Rebuild the ensemble (see CriticEnsemble / MLP in the source repo) and load strictly:
critics.load_state_dict(ckpt["critics"]) # 140 tensors, 10 members
critics.eval()
scaler = json.load(open("state_scaler.json"))
q = critics(state, chunk).min(dim=0).values # conservative reduction
Checkpoint keys: config, actor, critics, target_critics, env_steps, updates, extra.
config predates the critic_hidden field, so it reads None; the trained width is 256, and a
strict load_state_dict is the real check that your architecture matches.
Measured behaviour
Evaluated on 200 held-out base-policy rollouts (100 success / 100 failure) that the critic never trained on, against the closed-form return above:
| pearson | spearman | MAE | |
|---|---|---|---|
| all chunks | 0.759 | 0.607 | 0.1250 |
| success chunks only | 0.902 | 0.889 | 0.0458 |
Three properties worth knowing before building on it:
- Q is ~96% a state-value function. Perturbing the action at the prior's own draw spread (0.032/dim) moves Q by 4.1% of what changing the state does; at the learned residual's operating scale (0.011/dim) it is 1.5%. A state-only network fitted on 140 of these episodes reaches pearson 0.834 against this critic.
- It is well calibrated on successes and weaker on failures. Bellman consistency against an
independently fitted
Vis 0.908 on success chunks and 0.567 on failure chunks. - Success/failure level separation is 0.309 (means 0.363 vs 0.054, against a ground truth of 0.394 vs 0.000) — better separation than any reward model or fitted value baseline compared against it.
Caveats
Single training run. Same-configuration, same-seed replicates of arms in this study differ by
~3.6 pp in downstream success rate, so treat any comparison against another checkpoint at one
run each as unresolved. Trained and evaluated only on coffee_d0; the state layout is
task-specific and the weights will load but mean nothing on another task.