Omar L

Add Flora inference examples

52f830a verified 6 days ago

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`).

```bash
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.

```json
{"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:

```bash
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

```bash
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

```bash
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

```bash
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:

```json
{"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.

	# 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`).

	```bash
	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.

	```json
	{"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:

	```bash
	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

	```bash
	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

	```bash
	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

	```bash
	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:

	```json
	{"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.