Add the Cosmos workflow

[alexmillane]

Summary

Run Cosmos policies on OSMO.

Detailed description

• The Cosmos policy client was merged in Add cosmos policy #963.
• We now enable this in OSMO.
• Adds:
  • CosmosServerTask serves the checkpoint baked into the server image
  • CosmosPolicyRunnerTask runs the policy runner configured to point at the cosmos server.
  • CosmosPolicyRunnerWorkflow pairs the two.

[greptile-apps]

Greptile Summary

Adds Cosmos as an OSMO evaluation backend that co-schedules a remote policy runner and a checkpoint-serving sidecar.

• Adds Cosmos server and policy-runner task implementations.
• Registers --policy cosmos with the evaluation submission entry point.
• Adds a two-task workflow with OSMO host-token discovery and shared server port configuration.
• Adds rendering assertions covering task ordering, policy arguments, host wiring, port wiring, and server invocation.

Confidence Score: 5/5

The PR appears safe to merge, with the new Cosmos workflow consistently wired to the existing OSMO task and remote-policy abstractions.

The generated runner and server commands agree on host discovery and port configuration, the server image and checkpoint defaults match the repository's Cosmos build scripts, and the runner image contains the policy package and required client dependency.

Important Files Changed

Filename	Overview
osmo/tasks/cosmos_server_task.py	Adds a Cosmos server task whose image, baked checkpoint path, module invocation, and shared port match the repository's Cosmos image scripts.
osmo/tasks/cosmos_policy_runner_task.py	Adds a remote Cosmos policy-runner task using the workflow-resolved server host and shared policy-server port.
osmo/workflows/server_plus_policy_runner_workflow.py	Adds a gang-scheduled Cosmos runner/server workflow following the existing server-plus-runner abstraction.
osmo/submit_evaluation_workflow.py	Exposes the new Cosmos workflow through the evaluation submission policy registry.
isaaclab_arena/tests/test_osmo_workflow.py	Tests the rendered Cosmos task ordering, policy class, host token, port, and server entry point.

Sequence Diagram

sequenceDiagram
  participant User
  participant Submit as submit_evaluation_workflow
  participant OSMO
  participant Runner as CosmosPolicyRunnerTask
  participant Server as CosmosServerTask
  User->>Submit: --policy cosmos
  Submit->>OSMO: Submit CosmosPolicyRunnerWorkflow
  par Gang-scheduled tasks
    OSMO->>Server: Start baked Cosmos checkpoint server
    OSMO->>Runner: Start Isaac Lab evaluation
  end
  Runner->>Server: Connect via host token on port 8000
  loop Evaluation rollout
    Runner->>Server: Send observations and prompt
    Server-->>Runner: Return action chunk
  end

Loading

_{Reviews (1): Last reviewed commit: 446e391 | Re-trigger Greptile}

[arena-review-bot]

🤖 Isaac Lab-Arena Review Bot

Summary

Adds a Cosmos server task, a Cosmos policy-runner task, and a workflow pairing them, then registers it as --policy cosmos. The change slots cleanly into the existing ServerPlusPolicyRunnerWorkflow pattern, mirrors the pi0/gr00t tasks well, is additive/opt-in, and is covered by a rendering test. Two things worth a look before merge.

Findings

🟡 cosmos_policy_runner_task.py — CosmosRemotePolicyCfg.ping_timeout defaults to 20s and this runner does not raise it. Pi0RemotePolicyRunnerTask bumps it to 300 because on OSMO the first inference dropped the connection while the server compiled kernels. Cosmos is a heavier model, so it likely needs the same treatment — worth confirming or passing --ping_timeout.

🔵 cosmos_policy_runner_task.py — the remote_host constructor is now the third identical copy (pi0, gr00t, cosmos). A small shared RemotePolicyRunnerTask base would remove the duplication; optional.

Test Coverage

test_cosmos_workflow_renders_policy_runner_and_server renders the workflow and asserts task order, policy type, host token, and port — a good fit for this non-sim rendering path (inner/outer sim pattern not needed here). It does not assert a ping-timeout arg, which lines up with the finding above.

Verdict

Minor fixes needed