Repository navigation
Conversation
Copy the 34 files the DROID reader needs from nvidia/cosmos-framework 5950e7b4, byte for byte, under examples/worldmodel/cosmos3/vendor/, with the OpenMDW-1.1 text from the same commit. The example read DROID through its own reimplementation, and a reimplementation can drift from what it reimplements with nothing failing. One such drift was measured: the instruction was drawn with the global random.choice where the reference pins it, so the prompt could describe the same window differently -- 104 tokens against 149 in one measured case. Reading through upstream's own code makes a difference against upstream a difference in upstream. The files are not edited at all, so checking against a newer commit is one `diff -r`. That only holds while nothing reformats them, hence the pre-commit exclude: autoflake, isort and black would each rewrite code this repository does not own. The dependency set this needs is pinned and hand-picked rather than resolved, because resolving lerobot's own closure would install diffusers 0.35.2 over the 0.40.0 that provides the model and av 15.1.0 over the 18.1.0 the decoder links against. vendor/README.md records what that costs: the code is upstream's, the environment around it is not.
…ract The reader returns the composite quantised to uint8, which is exactly the value droid_dataset reaches before padding, so the seam sits there: this module pads to the canvas, normalises to [-1, 1] and takes the state row back off the action. Each step undoes something the reader did rather than re-deriving it. Measured on one pinned window, with the instruction pinned on both sides and augmentation off: all 34,214,400 composite values identical, maximum absolute difference 0.000 levels, and the state and action rows identical. Through this module all thirteen keys match the example's own reader except sample_id, whose difference is explained where it is built -- upstream's index names episodes by integer, so replay manifests do not carry across the two readers. The scaling half of resize_and_reflect_pad never runs for DROID: the composite is always 540x640, the 4:3 canvas is 544x736, and the scale factor is capped at 1.0, so the call only ever pads.
Put vendor/ on sys.path from the entry point only, and append rather than insert, so a host that also has cosmos-framework installed keeps loading the installed one -- the comparison machine does, and the reference side must never silently run this copy. Which one loaded is printed rather than assumed. The example's own reader stays reachable as --dataset example. It is not a fallback: its output was measured identical to upstream's, so it is a known-equivalent reference for comparing the two again after either side changes. Two steps on one card through the vendored reader: 16338993 window(s), samples=1 AR=168 DM=3093, loss 24.177132 then 11.412529.
Merging npu/dev brought the reworked main(), the package-qualified imports and two further dataset arguments, so the switch is re-applied onto that structure rather than merged textually. - keep_ranges now reaches the vendored reader, mapped to its own use_filter_dict/filter_dict_path. The entry passes it and the adapter did not take it, which would have trained on every window instead of the curated ones with nothing in the output to say so. The two sides read the same keep_ranges_1_0_1.json and droid_dataset's filter is transcribed from the reader's _append_index_records, so handing the file over makes the same selection. filter_dict_source_prefix is passed explicitly because upstream's own factory neither declares nor forwards it; calling the reader directly steps around that. - Both readers are resolved from module-level names, so a test substitutes either the same way. The entry defaults to the vendored reader, and the existing entry test stubs the reader by name, so the stub now stands in for both -- otherwise the default path would reach for the real pipeline and its dependencies. - Imports follow the entry's package-qualified style. 256 passed, 3 skipped. Two steps on one 910B2 through the vendored reader: 16338993 windows, step 1 loss 24.555565 -- the same value the example's own reader produces on this window, which is the end-to-end form of the field-by-field comparison recorded in vendor/README.md.
for more information, see https://pre-commit.ci
The released keep_ranges_1_0_1.json is keyed by composite gs:// URIs. The vendored reader refuses such a key without the URI prefix in front of the episode id, rather than guess it and select the wrong episodes, so the adapters hard-coded empty prefix turned --keep-ranges into a ValueError on the default reader while the same flag worked on the other one. Resolve the prefix the way droid_dataset.py resolves it, from an explicit argument or DROID_FILTER_SOURCE_PREFIX, so one setting serves both readers. Measured against the release: both readers index 16,338,993 windows unfiltered and 12,157,927 with the filter, 74.41 percent, the share the upstream recipe documents for the released model.
…sode id The vendored reader's _episode_records holds one entry per span, not per episode, and episode_id identifies neither. With the curated window filter on, an episode contributes one span per kept range and its id repeats across them; the id also counts from zero again in every registered root, so on a mixed success-and-failure run the same id names two different episodes. The table this module built was keyed by the id alone, so it collapsed both cases, and the linear scan that read it then resolved against whichever start survived. This is quiet in the worst way. The samples stay correct, because the reader resolves them; only the episode and frame a sample reports come out wrong, and those are what a sample id, a replay pack and a held-out manifest are built from. On the real curated set -- 12,157,927 windows, 66,423 spans over 54,987 episodes, 8,299 of them split into more than one span -- the old index named the wrong episode for 12.0% of 500 random windows and the wrong frame for 19.8%. Resolve a window the way the reader resolves it: a bisect over the cumulative span ends, then that span's record. Count the frame from the start of the episode rather than the start of the span, so two kept ranges of one episode cannot hand two windows the same number, and put the dataset index in the identity, so a per-root episode id cannot collide across roots. An index outside the reader's range is refused instead of being reported against the last span. The bisect also replaces a scan over every episode that ran once per sample. Verified on the real curated index: 200,000 random windows give 200,000 distinct ids, and the two spans of episode 31450 report frames 0.. and 144.. rather than both starting at zero.
|
Pushed two commits for the review comment on the index. Rebased onto the pre-commit.ci commit that was on the branch; the only conflict was the
The table this module built was keyed by the id alone, so it collapsed both cases, and the linear scan that read it resolved against whichever start survived. The samples stay correct — the reader resolves them — so only the episode and frame a sample reports were wrong, which is what a sample id, a replay pack and a held-out manifest are built from. On the real curated set (12,157,927 windows, 66,423 spans over 54,987 episodes, 8,299 of them split into more than one span) the old index named the wrong episode for 12.0% of 500 random windows and the wrong frame for 19.8%. Now resolved the way the reader resolves it: a bisect over the cumulative span ends, then that span's record. The frame is counted from the start of the episode rather than the start of the span, so two kept ranges of one episode cannot hand two windows the same number; the dataset index is part of the identity, so a per-root id cannot collide across roots; an index outside the reader's range raises instead of being reported against the last span. The bisect also replaces a scan over every episode that ran once per sample. Verified on the real curated index: 200,000 random windows give 200,000 distinct ids, and the two spans of episode 31450 report frames 0.. and 144.. rather than both starting at zero. Five new tests in The curated filter prefix never reached the reader (4381024) Found while checking the above. The released Measured against the release: both readers index 16,338,993 windows unfiltered and 12,157,927 with the filter — 74.41%, the share the upstream recipe documents for the released model. Four new tests in Full example suite: 265 passed, 3 skipped. |
Background
This example read DROID through
droid_dataset.py, its own implementation of what upstream's reader does.A reimplementation can drift from what it reimplements with nothing failing, and one such drift was
measured here: the window's instruction was drawn with the global
random.choicewhere the reference pinsit, so the same window could be described by a different sentence — 104 tokens against 149 in one measured
case. Nothing errors; the loss is simply computed on another input.
This PR reads DROID through upstream's own code instead. A difference against upstream then lives in
upstream, not in a transcription of it.
This changes no numbers. The pixels, action, state, instruction and idle-frame count were compared
field by field before the switch, and the first training step produces the same loss through either
reader.
What this PR adds
vendor/cosmos_framework/— the 34 files the DROID reader needs, copied fromnvidia/cosmos-frameworkat5950e7b4byte for byte, with the OpenMDW-1.1 text from the samecommit. No edits of any kind: no added provenance headers, no changed import paths. Checking against a
newer upstream is then one
diff -r, which is also why.pre-commit-config.yamlexcludes the directory— autoflake, isort and black would each rewrite code this repository does not own.
upstream_dataset.py— the seam. The reader returns the composite quantised to uint8, which isexactly the value
droid_datasetreaches before padding, so this pads to the canvas, normalises to[-1, 1] and takes the state row back off the action. Each step undoes something the reader did rather
than re-deriving it.
--dataset {upstream,example}, defaulting toupstream. The example's own reader stays reachable:not as a fallback but as a known-equivalent reference, so the two can be compared again after either
side changes.
vendor/goes onsys.pathfrom the entry point only, and is appended rather than inserted, so a hostthat also has cosmos-framework installed keeps loading the installed one — the comparison machine does,
and the reference side must never silently run this copy. Which one loaded is printed, not assumed.
What stays this example's
Reading the data is upstream's; turning it into model input is ours. Prompt assembly, CFG dropout,
tokenisation and packing stay here. Upstream has the first three too, in its SFT wrapper, and they are
deliberately not copied: that chain cannot be taken without
sequence_packing, upstream's packing forOmniMoTModel, which this example cannot use at all — it trains diffusers'Cosmos3OmniTransformer,whose forward takes loose tensors rather than a
packed_seqobject. 9949 further lines to reach code thatthen has to be bypassed buys nothing, and the three text stages are already known to agree.
Verification
One pinned window (episode 26827, first window), instruction pinned on both sides, augmentation off.
At the reader's output:
Through the adapter — what
pack_batchis handed — all thirteen keys, none missing or extra on eitherside, every one identical except
sample_id:video(so the padding and normalisation reproduceexactly),
action,state,task,idle_frames,canvas_hw,content_hw,aspect_key,episode_index,start,draw_id,additional_view_description.End to end, two steps on one Ascend 910B2:
16338993 window(s), step 1 loss 24.555565 — the samevalue this example's own reader produces on that window.
Three things not to over-read
same bytes on the same machine. The +0.81 LSB recorded between this accelerator and x86 is a platform
difference in the scaler and this says nothing about it.
cameras always compose to 540x640, the 4:3 canvas is 544x736, and the scale is capped at 1.0, so the
call only pads; the returned content region comes back (540, 640), unchanged.
after this change, so any difference in the previous implementation only affects runs already taken.
Dependencies, and the one thing the copy does not buy
lerobot==0.4.4(--no-deps),torchcodec==0.14.0+cpu,datasets==4.8.5,loguru. Resolving lerobot'sown closure is not safe here: it would install
diffusers 0.35.2over the 0.40.0 that provides the modeland
av 15.1.0over the 18.1.0 the decoder links against, so the four are pinned individually.torchcodec loads FFmpeg by SONAME and this image has no system FFmpeg; it ships a variant for FFmpeg 8,
which PyAV already bundles under mangled names, so a directory of correctly-named symlinks satisfies it
without changing the image.
The code is upstream's; the environment around it is not. Upstream runs inside lerobot's full
dependency set and this runs inside a subset chosen by hand, so a difference that comes from a package
version is still ours to explain.
vendor/README.mdsays so in those words.Files
41 files, +8073/-1: 36 under
vendor/(34 copied modules, the licence, a README),upstream_dataset.py,train_sft.py,tests/test_training_entry.py,.pre-commit-config.yamland a.gitignorefor theruns/directory the entry writes to by default.Not validated
Multi-card, throughput at scale, and any run longer than two steps. The first sample through the vendored
reader is slow (79.6 s of loader wait on a cold index); whether that persists at
--num-workers > 0isnot measured here.