Skip to content

feat: support Cosmos3-Edge-Policy-DROID in robolab policy server#117

Closed
ychao-nvidia wants to merge 1 commit into
mainfrom
ychao/20260718_edge_policy_droid
Closed

feat: support Cosmos3-Edge-Policy-DROID in robolab policy server#117
ychao-nvidia wants to merge 1 commit into
mainfrom
ychao/20260718_edge_policy_droid

Conversation

@ychao-nvidia

@ychao-nvidia ychao-nvidia commented Jul 18, 2026

Copy link
Copy Markdown
Collaborator

Summary

Enables serving the new Cosmos3-Edge-Policy-DROID checkpoint through the robolab action policy server. The Edge policy was post-trained from Cosmos3-Edge (Nemotron-3-Dense-VL backbone) and differs from the existing Nano policy in its processor, model config, tokenizer export, and prompt formatting — this MR wires up each of those differences so the checkpoint runs end-to-end.

Changes

  • data/generator/processors/__init__.py — Build the Edge processor from nvidia/Cosmos3-Edge-Reasoner. The policy checkpoint bundles a custom Cosmos3EdgeProcessor with no loadable processor code (no auto_map / remote
    code), so it can't be built via AutoProcessor. Sourcing it from the reasoner repo it was post-trained from mirrors how internal training sources the Nemotron-3-Dense-VL processor.
  • inference/common/public_model_config.py — Register the Nemotron3DenseVLMoTConfig.from_json_file and Nemotron3DenseVLTextForCausalLM target aliases used by the Edge policy.
  • scripts/action_policy_server_robolab.py — Add the Edge policy HF repo mapping, and a new format_prompt_as_json server arg that serves prompts as structured JSON to match training. When the transform emits ai_caption as a dict, it's JSON-encoded before serving.
  • utils/checkpoint_db.py — Re-form the bare relative Wan2.2 VAE path exported by the Edge checkpoint into an s3://bucket/... registry key. The Edge conversion exports the VAE tokenizer with an empty bucket_name (Nano exports bucket_name="bucket"), so the raw relative vae_path reaches the downloader instead of a registry key; re-forming it lets the VAE resolve through the registry and auto-download from HF, matching Nano's path. Local files still take precedence.

@ychao-nvidia
ychao-nvidia requested review from fwd4 and lfengad July 18, 2026 21:31
@ychao-nvidia ychao-nvidia self-assigned this Jul 18, 2026
Wire up serving for the Cosmos3-Edge-Policy-DROID checkpoint, which
differs from the Nano policy in its processor, config, tokenizer export,
and prompt formatting.

- processors: build the Edge processor from nvidia/Cosmos3-Edge-Reasoner,
  since the policy checkpoint bundles a custom Cosmos3EdgeProcessor with
  no loadable code (matching how training sources the reasoner processor).
- public_model_config: register Nemotron3DenseVL MoT config/model target
  aliases used by the Edge policy.
- action_policy_server_robolab: add the Edge policy HF repo mapping and a
  format_prompt_as_json arg to serve structured JSON prompts matching
  training; JSON-encode ai_caption when produced as a dict.
- checkpoint_db: re-form the bare relative Wan2.2 VAE path exported by the
  Edge checkpoint into an s3://bucket/... registry key so it auto-downloads
  from HF, matching Nano's path.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@ychao-nvidia
ychao-nvidia force-pushed the ychao/20260718_edge_policy_droid branch from 459e3a3 to 20a26a0 Compare July 18, 2026 21:44

@rickzw rickzw left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants