You can deploy this model and start earning money today!

chomeed/
dice_rl_mimicgen_coffee_d0_armF2_critic

Name: dice_rl_mimicgen_coffee_d0_armF2_critic
Availability: InStock

pytorch

robotics

reinforcement-learning

Views

No views yet

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

python

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 V is 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.

chomeed/dice_rl_mimicgen_coffee_d0_armF2_critic

DICE-RL residual critic — MimicGen coffee_d0 (arm F2)