Omar L
Add Flora inference examples
52f830a verified
|
Raw History Blame
5.94 kB

FLORA Rewarder: example inference

Example scripts for the three reward prompts of the FLORA Rewarder. One set of weights answers all three; the prompt selects the reward.

Script Reward Model input Output
terminal_reward.py Episode success, 1–5 Whole episode, up to 64 frames at up to 10 FPS, camera collage Score, P(1..5), expected score
dense_reward.py Action quality, 1–3 per 0.5 s Overlapping 16 s clips, 64 frames at 4 FPS, camera collage Score per frame, segments, rationales for bad spans
rubric_reward.py Rubric step scores One video per camera, 2 FPS (up to 166 frames each), rubric Thinking trace, one score per step

Scores: terminal 1 = complete failure, 5 = complete success; dense 1 = failure, 2 = suboptimal, 3 = normal.

Setup

Validated with Python 3.12 and NVIDIA H100 80 GB GPUs: one GPU for terminal and dense rewards, two for rubric rewards with the default settings (--tensor-parallel-size 2).

pip install -r batch_utils/requirements.txt

Run the scripts from a downloaded copy of this repository, the directory that contains batch_utils/; they load the weights stored next to it. To load the weights from elsewhere, pass --model with a Hub repository id or a local directory, or set FLORA_REWARDER_REPO. The first run compiles GPU kernels and can take several minutes before scoring starts. --dry-run reads the videos and prints the first prompt without loading the model.

Inputs

Each episode is one line of a JSONL file; video paths are relative to that file. Camera videos of one episode must be synchronized: same frame rate and frame count, starting together.

{"id": "ep001", "task": "Put the cup on the tray.", "videos": ["left_wrist.mp4", "top.mp4", "right_wrist.mp4"]}

task is one plain-language instruction for the whole episode.

Camera order (terminal and dense). The views of each time step are tiled into one frame: one view as is, two or three side by side, four as a 2 × 2 grid, in the order given. Use the order the model was trained with:

Views Order
2 left camera, right camera (for example a left exterior camera and the wrist camera)
3 left, center, right (for example left wrist, top, right wrist)
4 top, front or bottom, left, right

A single episode can also be passed on the command line:

python -m batch_utils.terminal_reward --video left_wrist.mp4 top.mp4 right_wrist.mp4 \
  --task "Put the cup on the tray." --output terminal.jsonl

Terminal reward

python -m batch_utils.terminal_reward --jobs batch_utils/examples/jobs.jsonl --output terminal.jsonl

Frames are sampled uniformly over the whole episode: min(64, floor(10 × duration in s)), rounded down to an even count. Add --with-rationale for a one-line explanation before the score. Each output line has:

  • score: the generated score, 1–5;
  • probabilities: P(1..5) at the score token, and expected_score, their mean;
  • progress: (expected_score − 1) / 4, a value in [0, 1] for ranking, not a calibrated success probability.

Dense reward

python -m batch_utils.dense_reward --jobs batch_utils/examples/jobs.jsonl --output dense.jsonl

The episode is split into 16 s clips (64 frames at 4 FPS) that start every 10 s; the last clip ends at the episode end. The model scores the 32 frame pairs (0.5 s each) of every clip and writes a rationale for every span scored 1 or 2. Clips are stitched by dropping the last 4 s of every clip except the final one and giving each frame the score of the nearest kept pair.

Each output line has segments (start and end time, score, rationale), per_second (the score at the middle of each second), fraction of time per score, and the raw answer of every clip. --frame-scores adds one score per source frame. Clips whose answer cannot be parsed are scored normal and counted in invalid_clips.

Dense scores are retrospective: a clip also shows the seconds after each scored moment. Use them for offline labeling and data filtering, not as an online signal.

Rubric reward

python -m batch_utils.rubric_reward --jobs batch_utils/examples/rubric_jobs.jsonl --output rubric.jsonl

A rubric job also names the robot and describes every camera, in video order:

{"id": "ep001", "task": "Put the cup on the tray.", "robot": "bimanual robot",
 "cameras": ["a top (scene) camera", "a left-wrist camera", "a right-wrist camera"],
 "videos": ["top.mp4", "left_wrist.mp4", "right_wrist.mp4"],
 "rubric": [
   {"action": "Grasp the cup", "grading": [
     {"label": "Not attempted", "points": 0.0},
     {"label": "Grasped after re-tries", "points": 0.1},
     {"label": "Clean grasp", "points": 0.2}]},
   {"action": "Set the cup on the tray", "grading": [
     {"label": "Not attempted", "points": 0.0},
     {"label": "Placed but tipped over", "points": 0.4},
     {"label": "Placed upright", "points": 0.8}]}]}

Each step lists the scores that may be awarded; the step maxima should sum to 1. Instead of robot and cameras, a job may give header, the full opening paragraph of the prompt.

The model thinks before it answers ANSWER: [...]. The script samples --num-samples answers (default 9) and reports the medoid: the answer with the smallest summed L1 distance to the others. Use --num-samples 1 for a single answer. Traces are often several thousand tokens; --keep-thinking writes them to the output.

Notes

  • Terminal and dense answers are decoded greedily with the answer format enforced; rubric answers are sampled at temperature 1.
  • Videos are read at their native frame rate, and the model receives absolute episode times. Keep the native videos; do not resample them.
  • Validate reward mappings on held-out episodes from your own robot before using them for training or data selection.