Spaces:

armnet
/

armnet-eval

Running

deploy demo space from ad13168d36ff

523d6e2 verified about 1 month ago

8.29 kB

"""armnet eval Space configuration.
A deliberately constrained public demo of the armnet infra: no install, no auth
setup, no Docker. A pre-provisioned API key is baked in (it is a throwaway demo
key, intended to be invalidated later), jobs route to any free BusyBox cell, and
the runtime image is pre-built and referenced by tag (a Space can't build Docker).
Everything that differs between the tasks on offer — the dataset to train
against, the reference policy, the leaderboard scope — hangs off :data:`TASKS`,
so adding a third task is a matter of one more entry rather than a sweep through
the UI.
"""

from __future__ import annotations

import os
from dataclasses import dataclass

# Demo API key. Provided as a Space secret (ARMNET_API_KEY) — nothing is baked
# into the repo, so no key ships in source. armnet-client reads it straight from
# the environment (armnet_client._config.api_key()).
DEMO_API_KEY = os.environ.get("ARMNET_API_KEY")

# Every task runs on the single-arm BusyBox cell. Routing is by task slug alone:
# the orchestrator resolves the slug to the `busybox` environment and sends the
# job to a cell serving it, so the Space never names a cell in a submission.
# Only tasks that auto-score off the instrumented box and reset themselves
# between rollouts belong here — one needing a human to tidy up between rollouts
# would stall an unattended public demo.
EMBODIMENT = "lerobot/so-101"
MULTITASK_DATASET = "armnet/busybox_multitask"


@dataclass(frozen=True)
class Task:
    """One evaluable task: what to train on, what to beat, what to rank."""

    slug: str
    """The armnet task slug submitted with the job."""

    label: str
    """What the selector calls it."""

    blurb: str
    """One line describing the rollout, shown under the selector."""

    dataset: str
    """HF dataset to train against, linked from the instructions."""

    reference_policy: str
    """Pre-trained policy offered as the default in the Live tab."""

    reference_label: str
    """How the reference policy is named in the dropdown."""

    @property
    def presets(self) -> dict[str, str]:
        return {self.reference_label: self.reference_policy}


TASKS: dict[str, Task] = {
    task.label: task
    for task in (
        Task(
            slug="push_green_button",
            label="Push the green button",
            blurb=(
                "Push the green button on the BusyBox panel. The button is "
                "momentary and springs back, so the scene resets itself."
            ),
            dataset="armnet/busybox_push_green_button",
            reference_policy="armnet/busybox_push_green_button_diffusion",
            reference_label="Diffusion · push green button — reference",
        ),
        Task(
            slug="busybox_left_toggle_down",
            label="Flip left switch off",
            blurb=(
                "Flip the left switch down into the off position. The arm flips it "
                "back up between rollouts, so the scene resets itself."
            ),
            dataset="armnet/busybox_flip_left_switch_off",
            reference_policy="armnet/busybox_flip_left_switch_off_diffusion",
            reference_label="Diffusion · flip left switch off — reference",
        ),
    )
}

TASK_LABELS = list(TASKS)
DEFAULT_TASK_LABEL = TASK_LABELS[0]


def task_for(label: str) -> Task:
    """The selected task, falling back to the default for an unknown label."""

    return TASKS.get(label) or TASKS[DEFAULT_TASK_LABEL]

# Fallback cell id written on leaderboard rows that predate per-job attribution
# (and on a run whose assigned cell is not yet known). Live availability no
# longer names a particular cell — the UI reports how many cells are free.
CELL_ID = os.environ.get("ARMNET_DEMO_CELL_ID", "cell-01")

# How often the background poller refreshes fleet availability. This is a
# single process-wide poll shared by every viewer (not per tab), so it can
# stay brisk without adding to the per-viewer request budget.
CELL_STATUS_POLL_S = float(os.environ.get("ARMNET_DEMO_CELL_STATUS_POLL_S", "5.0"))

# Number of rollouts per eval run. 20 to keep per-policy variance down.
N_EPISODES = int(os.environ.get("ARMNET_DEMO_N_EPISODES", "20"))
# Protect the shared public queue from unbounded submissions. The orchestrator
# Client-side abuse guard: only queued jobs submitted by this Space count.
MAX_QUEUED_JOBS = int(os.environ.get("ARMNET_DEMO_MAX_QUEUED_JOBS", "10"))
SUBMISSION_SOURCE = "armnet_eval_space"

# SDK-side wait budget. 10 rollouts x ~30s + resets + policy load; generous.
TIMEOUT_SECONDS = int(os.environ.get("ARMNET_DEMO_TIMEOUT_S", "2400"))

# Pre-built + pre-pushed runtime image ref. A Space can't build Docker, so the
# lerobot-eval runtime (which bakes in the molmoact2 plugin) is built and pushed
# elsewhere and named here. The policy itself is NOT baked in — it is loaded at
# runtime from the repo id passed as a job arg, so one image serves every policy.
#
# The ARMNET_DEMO_EVAL_IMAGE variable is what actually decides, and the
# armnet-demo-eval-image workflow keeps it pointing at the current runtime. This
# constant is only the fallback for a Space (or a local run) without it set, so
# treat it as a floor rather than the truth: it is bumped by hand and will lag.
_DEFAULT_EVAL_IMAGE = (
    "europe-north1-docker.pkg.dev/alpha-robotics-476817/customer-images/"
    "ville/armnet-lerobot-eval:armnet-01c6ad4641e6"
)
EVAL_IMAGE = os.environ.get("ARMNET_DEMO_EVAL_IMAGE") or _DEFAULT_EVAL_IMAGE

# Per-rollout wall-clock and control rate (mirror armnet-lerobot-eval defaults).
EPISODE_TIME_S = float(os.environ.get("ARMNET_DEMO_EPISODE_TIME_S", "30.0"))
POLICY_FPS = int(os.environ.get("ARMNET_DEMO_POLICY_FPS", "20"))
START_SEED = 1000

# Every leaderboard policy is scored across the same spread of rail positions,
# camera angles and lighting rather than one staged pose, so the ranking
# rewards generalisation. Not exposed to the submitter and deliberately not
# read from the environment: a seed anyone could change is a leaderboard whose
# entries are not comparable.
VARIATION_SEED = 42

# Optional HF token. Two uses: (1) the Space stores it once as a named
# orchestrator secret (EVAL_HF_SECRET_NAME) under the demo user and the eval
# references it as $HF_TOKEN, so the container can pull a gated policy; (2) it
# authenticates the leaderboard dataset writes. Set it as a Space secret; without
# it only fully-public policies load and the leaderboard is read-only.
HF_TOKEN = os.environ.get("HF_TOKEN")

# Name of the orchestrator secret (under the demo user) that holds the HF token.
# The job references it as {"HF_TOKEN": EVAL_HF_SECRET_NAME}; the orchestrator
# resolves the *name* to its value — it never accepts an inline token value.
EVAL_HF_SECRET_NAME = os.environ.get("ARMNET_DEMO_HF_SECRET_NAME", "hf_token")

# Push the recorded eval dataset to the Hub. Off for the demo: it avoids needing
# Hub write access and spamming datasets; the live results still show in the UI.
PUSH_DATASET = os.environ.get("ARMNET_DEMO_PUSH_DATASET", "").lower() in ("1", "true", "yes")

# Namespace for that dataset. The eval otherwise calls whoami() and uses the
# token owner's personal account, which is not the org the token can write to.
DATASET_NAMESPACE = os.environ.get("ARMNET_DEMO_DATASET_NAMESPACE", "armnet")

# Live-video stream size, passed to the eval as job args. Small values keep the
# NATS frames tiny so the stream stays live and doesn't drop frames (e.g. the
# success/fail moment) under backpressure. scale=0.3 -> ~9% of the pixels;
# jpeg_quality=30 -> strong compression. Together ~10% of the original size.
RERUN_IMAGE_SCALE = float(os.environ.get("ARMNET_DEMO_RERUN_IMAGE_SCALE", "0.3"))
RERUN_JPEG_QUALITY = int(os.environ.get("ARMNET_DEMO_RERUN_JPEG_QUALITY", "30"))


# Leaderboard persistence. Gradio has no built-in DB and Space disk is
# ephemeral, so completed runs are appended to this HF *dataset* repo (reloaded
# on restart). Reads are public; writes need HF_TOKEN (same one used for gated
# policy pulls). Pooled by (policy repo, revision).
LEADERBOARD_REPO = os.environ.get(
    "ARMNET_DEMO_LEADERBOARD_REPO", "armnet/armnet-demo-leaderboard"
)


def eval_image_configured() -> bool:
    return bool(EVAL_IMAGE)

	"""armnet eval Space configuration.

	A deliberately constrained public demo of the armnet infra: no install, no auth
	setup, no Docker. A pre-provisioned API key is baked in (it is a throwaway demo
	key, intended to be invalidated later), jobs route to any free BusyBox cell, and
	the runtime image is pre-built and referenced by tag (a Space can't build Docker).

	Everything that differs between the tasks on offer — the dataset to train
	against, the reference policy, the leaderboard scope — hangs off :data:`TASKS`,
	so adding a third task is a matter of one more entry rather than a sweep through
	the UI.
	"""

	from __future__ import annotations

	import os
	from dataclasses import dataclass

	# Demo API key. Provided as a Space secret (ARMNET_API_KEY) — nothing is baked
	# into the repo, so no key ships in source. armnet-client reads it straight from
	# the environment (armnet_client._config.api_key()).
	DEMO_API_KEY = os.environ.get("ARMNET_API_KEY")

	# Every task runs on the single-arm BusyBox cell. Routing is by task slug alone:
	# the orchestrator resolves the slug to the `busybox` environment and sends the
	# job to a cell serving it, so the Space never names a cell in a submission.
	# Only tasks that auto-score off the instrumented box and reset themselves
	# between rollouts belong here — one needing a human to tidy up between rollouts
	# would stall an unattended public demo.
	EMBODIMENT = "lerobot/so-101"
	MULTITASK_DATASET = "armnet/busybox_multitask"


	@dataclass(frozen=True)
	class Task:
	"""One evaluable task: what to train on, what to beat, what to rank."""

	slug: str
	"""The armnet task slug submitted with the job."""

	label: str
	"""What the selector calls it."""

	blurb: str
	"""One line describing the rollout, shown under the selector."""

	dataset: str
	"""HF dataset to train against, linked from the instructions."""

	reference_policy: str
	"""Pre-trained policy offered as the default in the Live tab."""

	reference_label: str
	"""How the reference policy is named in the dropdown."""

	@property
	def presets(self) -> dict[str, str]:
	return {self.reference_label: self.reference_policy}


	TASKS: dict[str, Task] = {
	task.label: task
	for task in (
	Task(
	slug="push_green_button",
	label="Push the green button",
	blurb=(
	"Push the green button on the BusyBox panel. The button is "
	"momentary and springs back, so the scene resets itself."
	),
	dataset="armnet/busybox_push_green_button",
	reference_policy="armnet/busybox_push_green_button_diffusion",
	reference_label="Diffusion · push green button — reference",
	),
	Task(
	slug="busybox_left_toggle_down",
	label="Flip left switch off",
	blurb=(
	"Flip the left switch down into the off position. The arm flips it "
	"back up between rollouts, so the scene resets itself."
	),
	dataset="armnet/busybox_flip_left_switch_off",
	reference_policy="armnet/busybox_flip_left_switch_off_diffusion",
	reference_label="Diffusion · flip left switch off — reference",
	),
	)
	}

	TASK_LABELS = list(TASKS)
	DEFAULT_TASK_LABEL = TASK_LABELS[0]


	def task_for(label: str) -> Task:
	"""The selected task, falling back to the default for an unknown label."""

	return TASKS.get(label) or TASKS[DEFAULT_TASK_LABEL]

	# Fallback cell id written on leaderboard rows that predate per-job attribution
	# (and on a run whose assigned cell is not yet known). Live availability no
	# longer names a particular cell — the UI reports how many cells are free.
	CELL_ID = os.environ.get("ARMNET_DEMO_CELL_ID", "cell-01")

	# How often the background poller refreshes fleet availability. This is a
	# single process-wide poll shared by every viewer (not per tab), so it can
	# stay brisk without adding to the per-viewer request budget.
	CELL_STATUS_POLL_S = float(os.environ.get("ARMNET_DEMO_CELL_STATUS_POLL_S", "5.0"))

	# Number of rollouts per eval run. 20 to keep per-policy variance down.
	N_EPISODES = int(os.environ.get("ARMNET_DEMO_N_EPISODES", "20"))
	# Protect the shared public queue from unbounded submissions. The orchestrator
	# Client-side abuse guard: only queued jobs submitted by this Space count.
	MAX_QUEUED_JOBS = int(os.environ.get("ARMNET_DEMO_MAX_QUEUED_JOBS", "10"))
	SUBMISSION_SOURCE = "armnet_eval_space"

	# SDK-side wait budget. 10 rollouts x ~30s + resets + policy load; generous.
	TIMEOUT_SECONDS = int(os.environ.get("ARMNET_DEMO_TIMEOUT_S", "2400"))

	# Pre-built + pre-pushed runtime image ref. A Space can't build Docker, so the
	# lerobot-eval runtime (which bakes in the molmoact2 plugin) is built and pushed
	# elsewhere and named here. The policy itself is NOT baked in — it is loaded at
	# runtime from the repo id passed as a job arg, so one image serves every policy.
	#
	# The ARMNET_DEMO_EVAL_IMAGE variable is what actually decides, and the
	# armnet-demo-eval-image workflow keeps it pointing at the current runtime. This
	# constant is only the fallback for a Space (or a local run) without it set, so
	# treat it as a floor rather than the truth: it is bumped by hand and will lag.
	_DEFAULT_EVAL_IMAGE = (
	"europe-north1-docker.pkg.dev/alpha-robotics-476817/customer-images/"
	"ville/armnet-lerobot-eval:armnet-01c6ad4641e6"
	)
	EVAL_IMAGE = os.environ.get("ARMNET_DEMO_EVAL_IMAGE") or _DEFAULT_EVAL_IMAGE

	# Per-rollout wall-clock and control rate (mirror armnet-lerobot-eval defaults).
	EPISODE_TIME_S = float(os.environ.get("ARMNET_DEMO_EPISODE_TIME_S", "30.0"))
	POLICY_FPS = int(os.environ.get("ARMNET_DEMO_POLICY_FPS", "20"))
	START_SEED = 1000

	# Every leaderboard policy is scored across the same spread of rail positions,
	# camera angles and lighting rather than one staged pose, so the ranking
	# rewards generalisation. Not exposed to the submitter and deliberately not
	# read from the environment: a seed anyone could change is a leaderboard whose
	# entries are not comparable.
	VARIATION_SEED = 42

	# Optional HF token. Two uses: (1) the Space stores it once as a named
	# orchestrator secret (EVAL_HF_SECRET_NAME) under the demo user and the eval
	# references it as $HF_TOKEN, so the container can pull a gated policy; (2) it
	# authenticates the leaderboard dataset writes. Set it as a Space secret; without
	# it only fully-public policies load and the leaderboard is read-only.
	HF_TOKEN = os.environ.get("HF_TOKEN")

	# Name of the orchestrator secret (under the demo user) that holds the HF token.
	# The job references it as {"HF_TOKEN": EVAL_HF_SECRET_NAME}; the orchestrator
	# resolves the name to its value — it never accepts an inline token value.
	EVAL_HF_SECRET_NAME = os.environ.get("ARMNET_DEMO_HF_SECRET_NAME", "hf_token")

	# Push the recorded eval dataset to the Hub. Off for the demo: it avoids needing
	# Hub write access and spamming datasets; the live results still show in the UI.
	PUSH_DATASET = os.environ.get("ARMNET_DEMO_PUSH_DATASET", "").lower() in ("1", "true", "yes")

	# Namespace for that dataset. The eval otherwise calls whoami() and uses the
	# token owner's personal account, which is not the org the token can write to.
	DATASET_NAMESPACE = os.environ.get("ARMNET_DEMO_DATASET_NAMESPACE", "armnet")

	# Live-video stream size, passed to the eval as job args. Small values keep the
	# NATS frames tiny so the stream stays live and doesn't drop frames (e.g. the
	# success/fail moment) under backpressure. scale=0.3 -> ~9% of the pixels;
	# jpeg_quality=30 -> strong compression. Together ~10% of the original size.
	RERUN_IMAGE_SCALE = float(os.environ.get("ARMNET_DEMO_RERUN_IMAGE_SCALE", "0.3"))
	RERUN_JPEG_QUALITY = int(os.environ.get("ARMNET_DEMO_RERUN_JPEG_QUALITY", "30"))


	# Leaderboard persistence. Gradio has no built-in DB and Space disk is
	# ephemeral, so completed runs are appended to this HF dataset repo (reloaded
	# on restart). Reads are public; writes need HF_TOKEN (same one used for gated
	# policy pulls). Pooled by (policy repo, revision).
	LEADERBOARD_REPO = os.environ.get(
	"ARMNET_DEMO_LEADERBOARD_REPO", "armnet/armnet-demo-leaderboard"
	)


	def eval_image_configured() -> bool:
	return bool(EVAL_IMAGE)