creature-lab 0.3.0__py3-none-any.whl

creature_lab/__init__.py

@@ -0,0 +1,18 @@
+"""Creature Lab: a minimal, visual, backend-agnostic creature simulation lab."""
+
+from creature_lab.schema import (
+    CreatureSpec,
+    EpisodeTrace,
+    FrameState,
+    TaskSpec,
+)
+
+VERSION = "0.3.0"
+
+__all__ = [
+    "VERSION",
+    "CreatureSpec",
+    "EpisodeTrace",
+    "FrameState",
+    "TaskSpec",
+]

creature_lab/agents/__init__.py

@@ -0,0 +1,5 @@
+"""Agent design loop: validated tools, policies, and the hill-climbing loop.
+
+Agents act only through validated tools that return a fresh, validated
+``CreatureSpec`` — they never touch a backend or unvalidated state.
+"""

creature_lab/agents/baseline.py

@@ -0,0 +1,53 @@
+"""A no-LLM policy that proposes random, valid tool calls.
+
+Lets the agent loop (and ``creature-lab ask --offline``) run without any model
+provider, so the lab is useful and testable out of the box. Deterministic for a
+given seed.
+"""
+
+from __future__ import annotations
+
+import random
+
+from creature_lab.agents.loop import Observation, Proposal
+from creature_lab.schema.creature import ShapeType
+
+_LENGTH_SHAPES = (ShapeType.CAPSULE, ShapeType.CYLINDER)
+
+
+class RandomToolPolicy:
+    """Propose a random parameter tweak via the validated tool layer."""
+
+    def __init__(self, seed: int = 0) -> None:
+        self._rng = random.Random(seed)
+
+    def __call__(self, observation: Observation) -> Proposal:
+        creature = observation.creature
+        rng = self._rng
+        options: list[str] = []
+        if creature.motors:
+            options += ["amplitude", "frequency", "phase"]
+        if any(part.shape in _LENGTH_SHAPES for part in creature.parts):
+            options.append("length")
+        if not options:
+            # Nothing tweakable; a no-op proposal the loop will reject as invalid.
+            return Proposal("set_motor", {"joint": "_none_"}, "no tunable parameters")
+
+        choice = rng.choice(options)
+        if choice == "length":
+            limbs = [part for part in creature.parts if part.shape in _LENGTH_SHAPES]
+            limb = rng.choice(limbs)
+            new_length = round(max(0.05, limb.length * rng.uniform(0.85, 1.15)), 4)
+            return Proposal("resize_limb", {"part": limb.id, "length": new_length}, "resize limb")
+
+        motor = rng.choice(creature.motors)
+        if choice == "amplitude":
+            value = round(max(0.0, motor.amplitude * rng.uniform(0.7, 1.3)), 4)
+            args = {"joint": motor.joint, "amplitude": value}
+            return Proposal("set_motor", args, "tune amplitude")
+        if choice == "frequency":
+            value = round(max(0.1, motor.frequency * rng.uniform(0.7, 1.3)), 4)
+            args = {"joint": motor.joint, "frequency": value}
+            return Proposal("set_motor", args, "tune frequency")
+        value = round(motor.phase + rng.uniform(-0.5, 0.5), 4)
+        return Proposal("set_motor", {"joint": motor.joint, "phase": value}, "tune phase")

creature_lab/agents/llm.py

@@ -0,0 +1,33 @@
+"""LiteLLM-backed policy (optional ``llm`` extra).
+
+Imports litellm lazily so the rest of the agent stack works without it. The
+prompt building and response parsing live in ``prompts.py`` and are tested
+without any network call.
+"""
+
+from __future__ import annotations
+
+from creature_lab.agents.loop import Observation, Proposal
+from creature_lab.agents.prompts import SYSTEM_PROMPT, build_prompt, parse_proposal
+
+
+class LLMPolicy:
+    """Ask a model (via LiteLLM) for the next tool call."""
+
+    def __init__(self, model: str = "gpt-4o-mini", temperature: float = 0.3) -> None:
+        self.model = model
+        self.temperature = temperature
+
+    def __call__(self, observation: Observation) -> Proposal:
+        import litellm
+
+        response = litellm.completion(
+            model=self.model,
+            temperature=self.temperature,
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": build_prompt(observation)},
+            ],
+        )
+        text = response.choices[0].message.content or ""
+        return parse_proposal(text)

creature_lab/agents/loop.py

@@ -0,0 +1,112 @@
+"""Policy-driven design loop.
+
+The loop is provider-agnostic: a ``Policy`` proposes a tool call given an
+``Observation``; the loop validates and applies it through the tool layer,
+re-evaluates the creature with a caller-supplied ``evaluate`` function, keeps the
+best, and records an ``AgentTrace``. This keeps the loop testable without any LLM
+or simulator.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from creature_lab.agents.tools import ToolError, apply_tool
+from creature_lab.schema import AgentStep, AgentTrace, CreatureSpec
+
+EvaluateFn = Callable[[CreatureSpec], float]
+#: Optional: summarize why a creature is scoring the way it does, as prompt-ready text.
+DiagnoseFn = Callable[[CreatureSpec], str]
+
+
+@dataclass(frozen=True)
+class Observation:
+    """What a policy sees before proposing the next edit."""
+
+    creature: CreatureSpec
+    best_score: float
+    attempt: int
+    goal: str = ""
+    #: Plain-language failure-pattern summary for the current best (empty if none/unavailable).
+    diagnosis: str = ""
+
+
+@dataclass(frozen=True)
+class Proposal:
+    """A tool call proposed by a policy."""
+
+    tool: str
+    args: dict[str, Any] = field(default_factory=dict)
+    note: str = ""
+
+
+Policy = Callable[[Observation], Proposal]
+
+
+@dataclass
+class DesignResult:
+    best: CreatureSpec
+    best_score: float
+    trace: AgentTrace
+
+
+def design_loop(
+    creature: CreatureSpec,
+    evaluate: EvaluateFn,
+    policy: Policy,
+    *,
+    attempts: int,
+    goal: str = "",
+    task_name: str = "design",
+    diagnose: DiagnoseFn | None = None,
+) -> DesignResult:
+    """Run the policy for `attempts` steps, keeping any higher-scoring valid edit.
+
+    When ``diagnose`` is given, it is called on the current best before every proposal
+    and its text is attached to the ``Observation`` — so a policy (an LLM prompt, or a
+    future rule-based one) can react to *why* the creature is scoring the way it is,
+    not just the scalar score.
+    """
+    if attempts < 0:
+        raise ValueError("attempts must not be negative")
+
+    best = creature
+    best_score = evaluate(creature)
+    steps = [AgentStep(attempt=0, action="seed", valid=True, score=best_score, accepted=True)]
+
+    for attempt in range(1, attempts + 1):
+        diagnosis = diagnose(best) if diagnose is not None else ""
+        proposal = policy(Observation(best, best_score, attempt, goal, diagnosis))
+        action = f"{proposal.tool}({proposal.args})"
+        try:
+            candidate = apply_tool(best, proposal.tool, proposal.args)
+        except ToolError as exc:
+            steps.append(
+                AgentStep(attempt=attempt, action=action, valid=False, note=str(exc)[:160])
+            )
+            continue
+        score = evaluate(candidate)
+        accepted = score > best_score
+        if accepted:
+            best, best_score = candidate, score
+        steps.append(
+            AgentStep(
+                attempt=attempt,
+                action=action,
+                valid=True,
+                score=score,
+                accepted=accepted,
+                note=proposal.note,
+            )
+        )
+
+    trace = AgentTrace(
+        creature_name=creature.name,
+        task_name=task_name,
+        goal=goal,
+        best_score=best_score,
+        steps=steps,
+    )
+    return DesignResult(best=best, best_score=best_score, trace=trace)

creature_lab/agents/prompts.py

@@ -0,0 +1,63 @@
+"""Prompt construction and response parsing for the LLM policy.
+
+Kept free of any provider import so the parser is unit-testable without network.
+"""
+
+from __future__ import annotations
+
+import json
+
+from creature_lab.agents.loop import Observation, Proposal
+from creature_lab.agents.tools import TOOL_SIGNATURES
+
+SYSTEM_PROMPT = (
+    "You are a creature-design assistant. You improve a small simulated robot-creature "
+    "by proposing ONE tool call at a time. Respond with a single JSON object only, of the "
+    'form {"tool": <name>, "args": {<kwargs>}, "note": <short reason>}. Do not add prose.'
+)
+
+
+def _creature_summary(observation: Observation) -> str:
+    creature = observation.creature
+    motors = ", ".join(
+        f"{m.joint}(amp={m.amplitude}, freq={m.frequency}, phase={m.phase})"
+        for m in creature.motors
+    )
+    limbs = ", ".join(f"{p.id}(len={p.length})" for p in creature.parts if p.length is not None)
+    summary = (
+        f"goal: {observation.goal}\n"
+        f"current best score: {observation.best_score:.4f}\n"
+        f"motors: {motors or 'none'}\n"
+        f"resizable limbs: {limbs or 'none'}"
+    )
+    if observation.diagnosis:
+        summary += f"\ndiagnosed issues: {observation.diagnosis}"
+    return summary
+
+
+def build_prompt(observation: Observation) -> str:
+    """Build the user message describing the creature and available tools."""
+    tools = "\n".join(f"- {name}({sig})" for name, sig in TOOL_SIGNATURES.items())
+    return (
+        f"{_creature_summary(observation)}\n\navailable tools:\n{tools}\n\nPropose one tool call."
+    )
+
+
+def parse_proposal(text: str) -> Proposal:
+    """Parse a model response into a Proposal, tolerating surrounding prose/code fences."""
+    start = text.find("{")
+    end = text.rfind("}")
+    if start == -1 or end == -1 or end < start:
+        raise ValueError("no JSON object found in model response")
+    try:
+        data = json.loads(text[start : end + 1])
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"invalid JSON in model response: {exc}") from exc
+    tool = data.get("tool")
+    if not isinstance(tool, str) or not tool:
+        raise ValueError("model response missing a 'tool' name")
+    args = data.get("args", {})
+    if not isinstance(args, dict):
+        raise ValueError("'args' must be an object")
+    note = str(data.get("note", ""))
+    return Proposal(tool=tool, args=args, note=note)

creature_lab/agents/tools.py

@@ -0,0 +1,101 @@
+"""Validated creature-editing tools.
+
+Each tool takes a ``CreatureSpec`` plus keyword arguments and returns a new,
+re-validated ``CreatureSpec``. Anything invalid (unknown id, wrong shape, value
+out of range) raises ``ToolError`` rather than producing an unvalidated creature,
+so the agent loop can record the failure and move on.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import ValidationError
+
+from creature_lab.schema import CreatureSpec
+from creature_lab.schema.creature import ShapeType
+
+_LENGTH_SHAPES = (ShapeType.CAPSULE, ShapeType.CYLINDER)
+
+
+class ToolError(ValueError):
+    """Raised when a tool is called with arguments that cannot be applied."""
+
+
+def _revalidate(data: dict) -> CreatureSpec:
+    try:
+        return CreatureSpec.model_validate(data)
+    except ValidationError as exc:
+        raise ToolError(f"resulting creature is invalid: {exc.errors()[0]['msg']}") from exc
+
+
+def set_motor(
+    creature: CreatureSpec,
+    *,
+    joint: str,
+    amplitude: float | None = None,
+    frequency: float | None = None,
+    phase: float | None = None,
+) -> CreatureSpec:
+    """Update an existing motor's sinusoid parameters."""
+    data = creature.model_dump()
+    motors = [motor for motor in data["motors"] if motor["joint"] == joint]
+    if not motors:
+        raise ToolError(f"no motor drives joint {joint!r}")
+    motor = motors[0]
+    if amplitude is not None:
+        motor["amplitude"] = amplitude
+    if frequency is not None:
+        motor["frequency"] = frequency
+    if phase is not None:
+        motor["phase"] = phase
+    return _revalidate(data)
+
+
+def set_joint_limit(
+    creature: CreatureSpec, *, joint: str, lower: float, upper: float
+) -> CreatureSpec:
+    """Set the lower/upper limit of a joint."""
+    data = creature.model_dump()
+    joints = [item for item in data["joints"] if item["id"] == joint]
+    if not joints:
+        raise ToolError(f"unknown joint {joint!r}")
+    joints[0]["limit"] = [lower, upper]
+    return _revalidate(data)
+
+
+def resize_limb(creature: CreatureSpec, *, part: str, length: float) -> CreatureSpec:
+    """Change the length of a capsule/cylinder part."""
+    data = creature.model_dump()
+    parts = [item for item in data["parts"] if item["id"] == part]
+    if not parts:
+        raise ToolError(f"unknown part {part!r}")
+    if parts[0]["shape"] not in _LENGTH_SHAPES:
+        raise ToolError(f"part {part!r} has no length to resize")
+    parts[0]["length"] = length
+    return _revalidate(data)
+
+
+#: Public tool registry and a short, prompt-friendly description of each tool's args.
+TOOLS = {
+    "set_motor": set_motor,
+    "set_joint_limit": set_joint_limit,
+    "resize_limb": resize_limb,
+}
+
+TOOL_SIGNATURES = {
+    "set_motor": "joint:str, amplitude?:float, frequency?:float, phase?:float",
+    "set_joint_limit": "joint:str, lower:float, upper:float",
+    "resize_limb": "part:str, length:float",
+}
+
+
+def apply_tool(creature: CreatureSpec, name: str, args: dict[str, Any]) -> CreatureSpec:
+    """Dispatch to a tool by name, validating the call."""
+    tool = TOOLS.get(name)
+    if tool is None:
+        raise ToolError(f"unknown tool {name!r}")
+    try:
+        return tool(creature, **args)
+    except TypeError as exc:  # wrong/missing keyword arguments
+        raise ToolError(f"bad arguments for {name!r}: {exc}") from exc