leagents 0.0.4__py3-none-any.whl

leagents/__init__.py

@@ -0,0 +1,3 @@
+"""leagents — agentic orchestration for the LeRobot pipeline."""
+
+__version__ = "0.0.4"

leagents/agents/__init__.py

@@ -0,0 +1,21 @@
+from leagents.agents.base import RunResult, Runner, dry_runner, subprocess_runner
+from leagents.agents.data_agent import DataAgent
+from leagents.agents.eval_agent import EvalAgent, EvalError
+from leagents.agents.improve_agent import ImproveAgent, ImproveError
+from leagents.agents.knowledge_agent import KnowledgeAgent
+from leagents.agents.train_agent import TrainAgent, TrainError
+
+__all__ = [
+    "RunResult",
+    "Runner",
+    "dry_runner",
+    "subprocess_runner",
+    "DataAgent",
+    "EvalAgent",
+    "EvalError",
+    "ImproveAgent",
+    "ImproveError",
+    "KnowledgeAgent",
+    "TrainAgent",
+    "TrainError",
+]

leagents/agents/base.py

@@ -0,0 +1,42 @@
+"""Shared subprocess runner for agents that wrap LeRobot CLIs.
+
+Runners are injectable so tests (and --dry-run) exercise the full loop
+without lerobot installed or a GPU present.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Sequence
+
+
+@dataclass
+class RunResult:
+    cmd: list[str]
+    exit_code: int
+    duration_s: float
+    log_path: Path | None = None
+
+
+Runner = Callable[[Sequence[str], Path], RunResult]
+
+
+def subprocess_runner(cmd: Sequence[str], log_path: Path) -> RunResult:
+    """Run a command, teeing stdout+stderr to a log file.
+
+    Training jobs stream output for hours — never buffer it in memory.
+    """
+    log_path.parent.mkdir(parents=True, exist_ok=True)
+    start = time.monotonic()
+    with log_path.open("w") as log:
+        proc = subprocess.run(list(cmd), stdout=log, stderr=subprocess.STDOUT)
+    return RunResult(list(cmd), proc.returncode, time.monotonic() - start, log_path)
+
+
+def dry_runner(cmd: Sequence[str], log_path: Path) -> RunResult:
+    log_path.parent.mkdir(parents=True, exist_ok=True)
+    log_path.write_text("DRY RUN: " + " ".join(cmd) + "\n")
+    return RunResult(list(cmd), 0, 0.0, log_path)

leagents/agents/data_agent.py

@@ -0,0 +1,38 @@
+"""Data Agent — M0 scope (DESIGN.md §3.2, §8).
+
+M0 collects no new data: the agent resolves the seed dataset from the
+proposal and records provenance for the cycle. RoboGene-style curation,
+MimicGen-style amplification, and GenAug-style augmentation land in M1.
+"""
+
+from __future__ import annotations
+
+from leagents.contracts import DatasetRef, Proposal
+from leagents.events import Event, EventBus
+
+
+class DataAgent:
+    def __init__(self, bus: EventBus):
+        self.bus = bus
+
+    def run(self, *, run_id: str, cycle: int, proposal: Proposal) -> DatasetRef:
+        ref = DatasetRef(
+            repo_id=proposal.dataset,
+            num_episodes=proposal.num_episodes,
+            notes=proposal.notes,
+        )
+        self.bus.emit(
+            Event(
+                run_id=run_id,
+                stage="collect",
+                kind="dataset_resolved",
+                cycle=cycle,
+                payload={
+                    "action": proposal.action,
+                    "repo_id": ref.repo_id,
+                    "num_episodes": ref.num_episodes,
+                    "notes": ref.notes,
+                },
+            )
+        )
+        return ref

leagents/agents/eval_agent.py

@@ -0,0 +1,105 @@
+"""Eval Agent — `lerobot-eval` LIBERO gate (DESIGN.md §3.4)."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from leagents.agents.base import Runner, subprocess_runner
+from leagents.config import EvalConfig
+from leagents.contracts import CheckpointRecord, EvalReport
+from leagents.events import Event, EventBus
+from leagents.orchestrator.constitution import Constitution, ConstitutionError
+
+
+class EvalError(Exception):
+    pass
+
+
+class EvalAgent:
+    def __init__(
+        self,
+        cfg: EvalConfig,
+        constitution: Constitution,
+        bus: EventBus,
+        runner: Runner = subprocess_runner,
+    ):
+        self.cfg = cfg
+        self.constitution = constitution
+        self.bus = bus
+        self.runner = runner
+
+    def build_command(self, checkpoint: CheckpointRecord, output_dir: Path) -> list[str]:
+        return [
+            "lerobot-eval",
+            f"--policy.path={checkpoint.path}",
+            f"--env.type={self.cfg.env_type}",
+            f"--env.task={self.cfg.task_suite}",
+            f"--eval.n_episodes={self.cfg.n_episodes}",
+            # lerobot-eval rejects batch_size > n_episodes (default batch is 50)
+            f"--eval.batch_size={min(self.cfg.batch_size, self.cfg.n_episodes)}",
+            f"--output_dir={output_dir}",
+            *self.cfg.extra_args,
+        ]
+
+    def run(
+        self, *, run_id: str, cycle: int, checkpoint: CheckpointRecord, workdir: Path
+    ) -> EvalReport:
+        output_dir = workdir / f"cycle_{cycle}" / "eval"
+        cmd = self.build_command(checkpoint, output_dir)
+
+        for verdict in (self.constitution.check_eval(self.cfg.env_type, self.cfg.n_episodes),
+                        self.constitution.check_command(cmd)):
+            if not verdict.allowed:
+                self.bus.emit(Event(run_id, "eval", "constitution_denied", cycle,
+                                    {"rule": verdict.rule, "reason": verdict.reason}))
+                raise ConstitutionError(verdict)
+
+        self.bus.emit(Event(run_id, "eval", "job_started", cycle, {"cmd": cmd}))
+        # log lives outside output_dir: lerobot CLIs refuse a pre-existing output_dir
+        result = self.runner(cmd, output_dir.parent / "eval.log")
+        self.bus.emit(Event(run_id, "eval", "job_finished", cycle,
+                            {"exit_code": result.exit_code, "duration_s": result.duration_s,
+                             "log": str(result.log_path)}))
+        if result.exit_code != 0:
+            raise EvalError(f"lerobot-eval exited {result.exit_code}, see {result.log_path}")
+
+        success_rate, per_task, raw = self._parse_eval_info(output_dir / "eval_info.json")
+        report = EvalReport(
+            checkpoint=checkpoint.path,
+            env_type=self.cfg.env_type,
+            task_suite=self.cfg.task_suite,
+            n_episodes=self.cfg.n_episodes,
+            success_rate=success_rate,
+            per_task=per_task,
+            raw=raw,
+        )
+        self.bus.emit(Event(run_id, "eval", "report", cycle,
+                            {"success_rate": success_rate, "task_suite": self.cfg.task_suite}))
+        return report
+
+    @staticmethod
+    def _parse_eval_info(path: Path) -> tuple[float, dict[str, float], dict[str, Any]]:
+        """Parse lerobot-eval's eval_info.json.
+
+        lerobot 0.5 schema: {"overall": {"pc_success": <0-100>, ...},
+        "per_group": {<group>: {"pc_success": ...}}, "per_task": [...]}.
+        Returns (success fraction, per-group success fractions, raw data).
+        """
+        if not path.exists():
+            raise EvalError(f"eval output missing: {path}")
+        data: dict[str, Any] = json.loads(path.read_text())
+        aggregate = data.get("overall") or data.get("aggregated") or data
+        if "pc_success" in aggregate:  # percentage 0-100 by definition
+            success = float(aggregate["pc_success"]) / 100.0
+        elif "success_rate" in aggregate:  # fraction 0-1
+            success = float(aggregate["success_rate"])
+        else:
+            raise EvalError(f"no success metric found in {path}")
+        per_task = {
+            group: float(metrics["pc_success"]) / 100.0
+            for group, metrics in (data.get("per_group") or {}).items()
+            if isinstance(metrics, dict) and "pc_success" in metrics
+        }
+        return success, per_task, data

leagents/agents/improve_agent.py

@@ -0,0 +1,140 @@
+"""Improvement Agent — DexFlyWheel step 3: success-filtered rollout collection
+(M1, DESIGN.md §3.5).
+
+Runs the blessed policy in sim via ``leagents.scripts.collect_rollouts`` (a
+subprocess, like every other agent) and keeps only successful episodes as a
+new LeRobotDataset. Residual-RL (flywheel step 2) and merging the rollout
+dataset into the training mix are the remaining M1 work.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+from leagents.agents.base import Runner, subprocess_runner
+from leagents.config import EvalConfig, ImproveConfig
+from leagents.contracts import CheckpointRecord, DatasetRef
+from leagents.events import Event, EventBus
+from leagents.orchestrator.constitution import Constitution, ConstitutionError
+
+
+class ImproveError(Exception):
+    pass
+
+
+class ImproveAgent:
+    def __init__(
+        self,
+        cfg: ImproveConfig,
+        eval_cfg: EvalConfig,
+        constitution: Constitution,
+        bus: EventBus,
+        runner: Runner = subprocess_runner,
+    ):
+        self.cfg = cfg
+        self.eval_cfg = eval_cfg
+        self.constitution = constitution
+        self.bus = bus
+        self.runner = runner
+
+    def build_command(self, checkpoint: CheckpointRecord, out_dir: Path, repo_id: str) -> list[str]:
+        cmd = [
+            sys.executable, "-m", "leagents.scripts.collect_rollouts",
+            f"--policy-path={checkpoint.path}",
+            f"--env-type={self.eval_cfg.env_type}",
+            f"--task={self.eval_cfg.task_suite}",
+            f"--episodes={self.cfg.episodes}",
+            f"--out={out_dir}",
+            f"--repo-id={repo_id}",
+            f"--device={self.cfg.device}",
+            *self.cfg.extra_args,
+        ]
+        if self.cfg.task_text:
+            cmd.append(f"--task-text={self.cfg.task_text}")
+        return cmd
+
+    def run(
+        self, *, run_id: str, cycle: int, checkpoint: CheckpointRecord, workdir: Path,
+        prev_mix: DatasetRef | None = None,
+    ) -> tuple[DatasetRef, dict]:
+        out_dir = workdir / f"cycle_{cycle}" / "rollouts"
+        repo_id = f"local/{run_id}-cycle{cycle}-rollouts"
+        cmd = self.build_command(checkpoint, out_dir, repo_id)
+
+        for verdict in (self.constitution.check_eval(self.eval_cfg.env_type, self.cfg.episodes),
+                        self.constitution.check_command(cmd)):
+            if not verdict.allowed:
+                self.bus.emit(Event(run_id, "improve", "constitution_denied", cycle,
+                                    {"rule": verdict.rule, "reason": verdict.reason}))
+                raise ConstitutionError(verdict)
+
+        self.bus.emit(Event(run_id, "improve", "job_started", cycle, {"cmd": cmd}))
+        result = self.runner(cmd, out_dir.parent / "rollouts.log")
+        self.bus.emit(Event(run_id, "improve", "job_finished", cycle,
+                            {"exit_code": result.exit_code, "duration_s": result.duration_s,
+                             "log": str(result.log_path)}))
+        if result.exit_code != 0:
+            raise ImproveError(f"rollout collection exited {result.exit_code}, "
+                               f"see {result.log_path}")
+
+        summary = self._parse_summary(result.log_path)
+        self.bus.emit(Event(run_id, "improve", "rollouts_collected", cycle, summary))
+        ref = DatasetRef(
+            repo_id=repo_id,
+            root=str(out_dir),
+            num_episodes=summary.get("kept"),
+            notes=f"success-filtered rollouts from cycle {cycle} blessed checkpoint",
+        )
+        if summary.get("kept", 0) == 0:
+            return (prev_mix if prev_mix else ref), summary
+        if prev_mix is None:
+            return ref, summary
+        return self._merge(run_id, cycle, workdir, prev_mix, ref), summary
+
+    def _merge(
+        self, run_id: str, cycle: int, workdir: Path, a: DatasetRef, b: DatasetRef
+    ) -> DatasetRef:
+        """Accumulate rollouts across cycles: mix_{n} = mix_{n-1} ∪ rollouts_n
+        (DexFlyWheel's growing dataset), via lerobot-edit-dataset merge."""
+        mix_repo_id = f"local/{run_id}-rollouts-mix-c{cycle}"
+        mix_root = workdir / f"cycle_{cycle}" / "rollouts_mix"
+        cmd = [
+            "lerobot-edit-dataset",
+            "--operation.type=merge",
+            f"--operation.repo_ids=[{a.repo_id}, {b.repo_id}]",
+            f"--operation.roots=[{a.root}, {b.root}]",
+            f"--new_repo_id={mix_repo_id}",
+            f"--new_root={mix_root}",
+            "--push_to_hub=false",
+        ]
+        self.bus.emit(Event(run_id, "improve", "job_started", cycle,
+                            {"cmd": cmd, "stage": "merge"}))
+        result = self.runner(cmd, mix_root.parent / "rollouts_merge.log")
+        self.bus.emit(Event(run_id, "improve", "job_finished", cycle,
+                            {"exit_code": result.exit_code, "duration_s": result.duration_s,
+                             "log": str(result.log_path)}))
+        if result.exit_code != 0:
+            raise ImproveError(f"rollout merge exited {result.exit_code}, see {result.log_path}")
+        total = (a.num_episodes or 0) + (b.num_episodes or 0)
+        merged = DatasetRef(repo_id=mix_repo_id, root=str(mix_root),
+                            num_episodes=total or None,
+                            notes=f"rollout mix through cycle {cycle}")
+        self.bus.emit(Event(run_id, "improve", "rollouts_merged", cycle,
+                            {"repo_id": mix_repo_id, "episodes": total}))
+        return merged
+
+    @staticmethod
+    def _parse_summary(log_path: Path | None) -> dict:
+        """The collector prints a JSON summary as its LAST stdout line."""
+        if log_path is None or not Path(log_path).exists():
+            raise ImproveError(f"rollout log missing: {log_path}")
+        for line in reversed(Path(log_path).read_text().splitlines()):
+            line = line.strip()
+            if line.startswith("{"):
+                try:
+                    return json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+        raise ImproveError(f"no JSON summary found in {log_path}")

leagents/agents/knowledge_agent.py

@@ -0,0 +1,156 @@
+"""Knowledge Agent — OKF knowledge layer (DESIGN.md §3.6, M1).
+
+Two operations mirroring the Karpathy-LLM-Wiki workflow:
+- ingest: after each cycle, distill the eval outcome into task/policy
+  pages (markdown + YAML frontmatter, per Google's OKF spec shape).
+- lint: health-check pass over the wiki (schema, provenance).
+
+Pages are advisory context only — loop control stays deterministic.
+Layer-1 artifacts (events.jsonl, SQLite, eval reports) are never edited;
+this agent writes only under the knowledge root.
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from leagents.contracts import EvalReport
+from leagents.events import Event, EventBus
+from leagents.llm import LLMClient, NullLLM
+
+_PAGE_TYPES = {"task", "policy", "experiment"}
+_STATUSES = {"observed-once", "replicated", "human-confirmed"}
+
+_LESSONS_HEADING = "## Lessons"
+
+_LESSONS_SYSTEM = (
+    "You maintain one page of a robotics run-knowledge wiki. Given the page and a new "
+    "observation, rewrite ONLY the Lessons section as a short markdown bullet list of "
+    "durable, actionable lessons. No preamble, bullets only."
+)
+
+
+def _slug(name: str) -> str:
+    return name.lower().replace(" ", "-").replace("/", "-")
+
+
+def _read_page(path: Path) -> tuple[dict[str, Any], str]:
+    text = path.read_text()
+    if text.startswith("---\n"):
+        _, frontmatter, body = text.split("---\n", 2)
+        return yaml.safe_load(frontmatter) or {}, body.lstrip("\n")
+    return {}, text
+
+
+def _write_page(path: Path, meta: dict[str, Any], body: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    frontmatter = yaml.safe_dump(meta, sort_keys=False, allow_unicode=True).strip()
+    path.write_text(f"---\n{frontmatter}\n---\n\n{body.strip()}\n")
+
+
+class KnowledgeAgent:
+    def __init__(self, root: Path, bus: EventBus, llm: LLMClient | None = None):
+        self.root = Path(root)
+        self.bus = bus
+        self.llm = llm or NullLLM()
+
+    # -- ingest ----------------------------------------------------------
+    def ingest(
+        self,
+        *,
+        run_id: str,
+        cycle: int,
+        policy: str,
+        report: EvalReport,
+        decision: str,
+    ) -> list[Path]:
+        observation = (
+            f"- {time.strftime('%Y-%m-%d')} · run {run_id} · cycle {cycle} · "
+            f"{report.task_suite} · policy {policy} · "
+            f"success {report.success_rate:.1%} · decision {decision}"
+        )
+        touched = [
+            self._upsert_page(
+                self.root / "policies" / f"{_slug(policy)}.md",
+                page_type="policy",
+                title=f"Policy: {policy}",
+                name=policy,
+                observation=observation,
+                provenance={"run": run_id, "cycle": cycle},
+            ),
+            self._upsert_page(
+                self.root / "tasks" / f"{_slug(report.task_suite)}.md",
+                page_type="task",
+                title=f"Task: {report.task_suite}",
+                name=report.task_suite,
+                observation=observation,
+                provenance={"run": run_id, "cycle": cycle},
+            ),
+        ]
+        self.bus.emit(
+            Event(run_id, "knowledge", "knowledge_updated", cycle,
+                  {"pages": [str(p) for p in touched]})
+        )
+        return touched
+
+    def _upsert_page(
+        self,
+        path: Path,
+        *,
+        page_type: str,
+        title: str,
+        name: str,
+        observation: str,
+        provenance: dict[str, Any],
+    ) -> Path:
+        if path.exists():
+            meta, body = _read_page(path)
+        else:
+            meta = {"name": name, "type": page_type, "status": "observed-once",
+                    "provenance": []}
+            body = f"# {title}\n\n## Observations\n\n{_LESSONS_HEADING}\n"
+
+        meta.setdefault("provenance", []).append(provenance)
+        if meta.get("status") != "human-confirmed":  # human verdicts are never downgraded
+            meta["status"] = "replicated" if len(meta["provenance"]) >= 2 else "observed-once"
+        meta["updated"] = time.strftime("%Y-%m-%d")
+
+        head, sep, lessons = body.partition(_LESSONS_HEADING)
+        body = head.rstrip("\n") + f"\n{observation}\n\n" + sep + lessons
+
+        lessons_update = self.llm.complete(
+            f"PAGE:\n{body}\n\nNEW OBSERVATION:\n{observation}", system=_LESSONS_SYSTEM
+        ).strip()
+        if lessons_update:
+            head, _, _ = body.partition(_LESSONS_HEADING)
+            body = head + f"{_LESSONS_HEADING}\n{lessons_update}\n"
+
+        _write_page(path, meta, body)
+        return path
+
+    # -- lint --------------------------------------------------------------
+    def lint(self) -> list[dict[str, str]]:
+        """Health check over all pages; findings feed the next experiments."""
+        findings: list[dict[str, str]] = []
+        for path in sorted(self.root.rglob("*.md")):
+            if path.name == "KNOWLEDGE.md":
+                continue
+            try:
+                meta, _ = _read_page(path)
+            except Exception as exc:
+                findings.append({"page": str(path), "problem": f"unparseable: {exc!r}"})
+                continue
+            if meta.get("type") not in _PAGE_TYPES:
+                findings.append({"page": str(path), "problem": f"bad type {meta.get('type')!r}"})
+            if meta.get("status") not in _STATUSES:
+                findings.append({"page": str(path),
+                                 "problem": f"bad status {meta.get('status')!r}"})
+            if not meta.get("provenance"):
+                findings.append({"page": str(path), "problem": "missing provenance"})
+        self.bus.emit(Event("knowledge-lint", "knowledge", "lint_report",
+                            payload={"findings": findings}))
+        return findings

leagents/agents/train_agent.py

@@ -0,0 +1,98 @@
+"""Training Agent — subprocess wrapper over `lerobot-train` (DESIGN.md §3.3)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from leagents.agents.base import Runner, subprocess_runner
+from leagents.config import PolicyRung, TrainConfig
+from leagents.contracts import CheckpointRecord, DatasetRef
+from leagents.events import Event, EventBus
+from leagents.orchestrator.constitution import Constitution, ConstitutionError
+
+
+class TrainError(Exception):
+    pass
+
+
+class TrainAgent:
+    def __init__(
+        self,
+        cfg: TrainConfig,
+        constitution: Constitution,
+        bus: EventBus,
+        runner: Runner = subprocess_runner,
+    ):
+        self.cfg = cfg
+        self.constitution = constitution
+        self.bus = bus
+        self.runner = runner
+
+    def build_command(
+        self,
+        rung: PolicyRung,
+        dataset: DatasetRef,
+        output_dir: Path,
+        init_override: str | None = None,
+        steps: int | None = None,
+    ) -> list[str]:
+        init = init_override or rung.init
+        policy_arg = f"--policy.path={init}" if init else f"--policy.type={rung.name}"
+        episodes_args = (
+            [f"--dataset.episodes={list(range(dataset.num_episodes))}"]
+            if dataset.num_episodes and dataset.root is None
+            else []
+        )
+        root_args = [f"--dataset.root={dataset.root}"] if dataset.root else []
+        return [
+            "lerobot-train",
+            policy_arg,
+            # lerobot >= 0.5 refuses to train without policy.repo_id unless hub
+            # push is explicitly off; leagents manages checkpoints locally.
+            "--policy.push_to_hub=false",
+            f"--dataset.repo_id={dataset.repo_id}",
+            *root_args,
+            *episodes_args,
+            f"--output_dir={output_dir}",
+            f"--steps={steps or self.cfg.steps}",
+            f"--batch_size={self.cfg.batch_size}",
+            *self.cfg.extra_args,
+        ]
+
+    def run(
+        self,
+        *,
+        run_id: str,
+        cycle: int,
+        rung: PolicyRung,
+        dataset: DatasetRef,
+        workdir: Path,
+        init_override: str | None = None,
+        steps: int | None = None,
+        stage: str = "train",
+    ) -> CheckpointRecord:
+        output_dir = workdir / f"cycle_{cycle}" / stage
+        cmd = self.build_command(rung, dataset, output_dir, init_override, steps)
+
+        for verdict in (self.constitution.check_train(steps or self.cfg.steps),
+                        self.constitution.check_command(cmd)):
+            if not verdict.allowed:
+                self.bus.emit(Event(run_id, "train", "constitution_denied", cycle,
+                                    {"rule": verdict.rule, "reason": verdict.reason}))
+                raise ConstitutionError(verdict)
+
+        self.bus.emit(Event(run_id, "train", "job_started", cycle, {"cmd": cmd, "stage": stage}))
+        # log lives outside output_dir: lerobot-train refuses a pre-existing output_dir
+        result = self.runner(cmd, output_dir.parent / f"{stage}.log")
+        self.bus.emit(Event(run_id, "train", "job_finished", cycle,
+                            {"exit_code": result.exit_code, "duration_s": result.duration_s,
+                             "log": str(result.log_path)}))
+        if result.exit_code != 0:
+            raise TrainError(f"lerobot-train exited {result.exit_code}, see {result.log_path}")
+
+        # lerobot-train checkpoint layout: output_dir/checkpoints/last/pretrained_model
+        checkpoint = output_dir / "checkpoints" / "last" / "pretrained_model"
+        if not checkpoint.exists():
+            self.bus.emit(Event(run_id, "train", "warning", cycle,
+                                {"msg": f"expected checkpoint missing: {checkpoint}"}))
+        return CheckpointRecord(policy_type=rung.name, path=str(checkpoint), cycle=cycle)