From dd9fd11b789558ffd16f92252a480ca334e8c3a5 Mon Sep 17 00:00:00 2001 From: ravindu somawansa Date: Thu, 8 Oct 2026 16:53:21 +0200 Subject: [PATCH] feat(inference): show the task field for any policy not known to ignore it --- lelab/jobs.py | 11 +++++++++-- tests/test_jobs.py | 14 ++++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/lelab/jobs.py b/lelab/jobs.py index 10810a16c..b4e41cc3d 100644 --- a/lelab/jobs.py +++ b/lelab/jobs.py @@ -579,7 +579,14 @@ def _list_imported_hub(api, repo_id: str) -> list[JobCheckpoint]: return [] -_LANGUAGE_CONDITIONED_POLICY_TYPES = {"smolvla", "pi0", "pi0_fast", "pi05"} +# Policies known to ignore the task string. Everything else (including plugin +# policies we can't enumerate) gets the task field, since an unused task is +# harmless but a missing one makes a language-conditioned policy unusable. +_TASK_FREE_POLICY_TYPES = {"act", "diffusion", "vqbet", "tdmpc", "sac"} + + +def _policy_requires_task(policy_type: object) -> bool: + return policy_type not in _TASK_FREE_POLICY_TYPES _HUB_CKPT_REF_RE = re.compile(r"^(?P[^@]+)@checkpoints/(?P\d+)$") @@ -1086,7 +1093,7 @@ def get_policy_config_summary(self, job_id: str, step: int) -> dict[str, object] return { "policy_type": policy_type, "image_features": image_features, - "requires_task": policy_type in _LANGUAGE_CONDITIONED_POLICY_TYPES, + "requires_task": _policy_requires_task(policy_type), } def delete(self, job_id: str) -> None: diff --git a/tests/test_jobs.py b/tests/test_jobs.py index 10384b1f5..40fecd810 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -271,3 +271,17 @@ def list_repo_files(self, repo_id, repo_type): assert rec.output_dir == "" cks = reg.list_checkpoints(rec.id) assert [c.ref for c in cks] == ["user/some-model@root"] + + +@pytest.mark.parametrize("policy_type", ["act", "diffusion", "vqbet", "tdmpc", "sac"]) +def test_policy_requires_task_false_for_task_free_policies(policy_type: str) -> None: + from lelab.jobs import _policy_requires_task + + assert _policy_requires_task(policy_type) is False + + +@pytest.mark.parametrize("policy_type", ["smolvla", "pi0", "pi05", "groot", "xvla", "some_plugin", None]) +def test_policy_requires_task_true_for_language_or_unknown_policies(policy_type: object) -> None: + from lelab.jobs import _policy_requires_task + + assert _policy_requires_task(policy_type) is True