Mujoco_WASM/decision_server/providers/jev.py

"""TypeSafe System One / OpenRouter Decisions protocol, not a chat endpoint.

Protocol shape informed by MIT-licensed jev-libero / embodied-jev; see THIRD_PARTY_NOTICES.
"""

import json
import math

from ..protocol import (
    LANGUAGE_SKILLS,
    LANGUAGE_VERSION,
    SKILLS,
    DecisionError,
    schema_for,
    validate_jev,
)
from .http import post, usage

# Criteria are the decision model's option semantics, not just display labels.
# In particular, `close` means close the gripper, NOT close/terminate the task.
SKILL_CRITERIA = {
    "stow": "Raise the empty arm to its safe travel pose before approaching the source.",
    "approach": (
        "After completedSkill=stow with failure=none, the arm has reached its safe travel pose. "
        "Proceed to dock the base at the source; the gripper should still be empty."
    ),
    "open": "Open the gripper before grasping a supported block (also safe recovery).",
    "pregrasp": "Move the TCP above the block while the base is stopped.",
    "descend": "Lower the open gripper from above to align its TCP with the block.",
    "close": (
        "Close the gripper fingers around the aligned, supported block; TCP-object distance "
        "must be <=0.012 m. Zero finger forces and secure=false BEFORE closing are expected."
    ),
    "verify": (
        "Lift to verify grasp AFTER closing, with both finger forces >=0.2 N. "
        "secure=false is expected BEFORE this lift; this skill establishes verified grasp."
    ),
    "carry": "Transport the block with verified grasp: evidence.secure=true is required.",
    "stop-base": "Stop at the target dock while maintaining evidence.secure=true.",
    "place": "Lower the transported block onto the target support with the base stopped.",
    "release": (
        "Open the fingers after evidence.onGoalSupport=true. secure may already be false "
        "because the block is now supported; this is expected, not slipping."
    ),
    "retreat": "Withdraw the empty gripper after releasing the supported block.",
    "settle": "Hold still to measure released block stability on the target support.",
    "stop": (
        "Choose stop when failure indicates a hard fault, safety flags are present, or the "
        "next skill's physical precondition is contradicted or cannot be established; "
        "NOT merely because the whole task is unfinished or the gripper is empty before grasp."
    ),
}
STAGE_CONTEXT = (
    "Judge the offered NEXT skill, not completion of the whole task. observation.phase names "
    "the current/just-completed skill; candidates names the next permitted skill or stop. "
    "transition.completedSkill means the local executor completed that skill and its exit "
    "checks without failure; null means initial entry or recovery, not completion. "
    "The local simulator freezes physics while waiting for this response; time is simulation "
    "seconds, not a wall-clock timestamp. The stamped observation is the current snapshot. "
    "An empty gripper before close is normal. verify establishes secure grasp by lifting; "
    "only carry/stop-base require an already verified grasp. During place/release the object "
    "returns to support, so secure=false there is not by itself a lost grasp. "
    "Never ignore an explicit failure, safety flag or missing required evidence. "
    "The local controller independently checks physical preconditions before any actuation. "
)


def question(options, instruction, criteria=None):
    return {
        "type": "choice",
        "instructions": instruction,
        "criteria": {option: (criteria or {}).get(option, option) for option in options},
    }


def answer(result, name, options):
    answers = result.get("answers")
    item = answers.get(name) if isinstance(answers, dict) else None
    if not isinstance(item, dict) or item.get("choice") not in options:
        raise DecisionError("jev_invalid_choice", 502)
    probabilities = item.get("probabilities", {})
    if (
        not isinstance(probabilities, dict)
        or probabilities.keys() - set(options)
        or any(
            type(p) not in (int, float) or not math.isfinite(p) or not 0 <= p <= 1
            for p in probabilities.values()
        )
    ):
        raise DecisionError("jev_invalid_probabilities", 502)
    # Keep official choice even when probabilities do not rank it highest.
    return item["choice"]


async def decide(session, conn, request):
    version = request["observation"]["version"]
    props = schema_for(version)["$defs"]["JevDecision"]["properties"]
    options = {name: props[name]["enum"] for name in ("grasp", "diagnosis", "recovery")}
    options["choice"] = request["candidates"]
    prompts = {
        "choice": (
            STAGE_CONTEXT + "Choose the safe offered next skill according to its criteria. "
            "Stop when its required evidence is absent or a hard safety fault is present."
        ),
        "grasp": (
            "secure requires two finger forces >=0.2 N, verified lift and stable grasp evidence. "
            "empty means no held object. slipping means a previously secure grasp is being lost. "
            "Use uncertain if evidence is insufficient. An open gripper is not a secure grasp."
        ),
        "diagnosis": (
            "Report none when failure=none and there are no safety flags. "
            "An empty gripper before closing or after release is expected, not a fault. "
            "Otherwise diagnose empty, slipping, misaligned, unreachable, stalled or uncertain."
        ),
        "recovery": (
            "Continue with a safe next skill if failure=none; retry only recoverable alignment "
            "or empty grasp failures before transport; replan when retries are insufficient; "
            "stop for hard safety faults or unsafe uncertainty."
        ),
    }
    if version == LANGUAGE_VERSION:
        options.update({name: props[name]["enum"] for name in ("noul", "score", "reason")})
        prompts.update(
            {
                "noul": (
                    STAGE_CONTEXT + "Veto unsafe NEXT skills, not incomplete tasks: allow with "
                    "sufficient evidence for the offered stage and no hard safety fault; "
                    "deny on danger, uncertain on missing evidence required for that stage. "
                    "This is not physical authorization."
                ),
                "score": (
                    "Grade progress at this skill boundary: poor, partial, good, excellent, "
                    "or unavailable. This is a discrete quality grade, NOT probability or "
                    "final success. Empty fingers before grasp/after release are expected."
                ),
                "reason": (
                    "Select the main reason for safety/quality: none, unsafe, "
                    "insufficient_evidence, tracking_error, verified_progress."
                ),
            }
        )
    observation = request["observation"]
    candidates = [skill for skill in request["candidates"] if skill != "stop"]
    next_skill = candidates[0] if len(candidates) == 1 else None
    failure = request.get("failure", "none")
    sequence = LANGUAGE_SKILLS if version == LANGUAGE_VERSION else SKILLS
    adjacent = (
        next_skill in sequence
        and sequence.index(next_skill) == sequence.index(observation["phase"]) + 1
    )
    # This is the caller's skill-boundary protocol, not a new safety authorization.
    # Preserve raw observations and vetoes; never infer grasp from task progress.
    transition = {
        "nextSkill": next_skill,
        "completedSkill": (
            observation["phase"]
            if adjacent and failure == "none" and not observation["safety"]
            else None
        ),
    }
    criteria = {
        "choice": SKILL_CRITERIA,
        "grasp": {
            "secure": "evidence.secure=true and both fingerForces >=0.2 N; verified lifted grasp.",
            "empty": "No held block; normal before closing and after release.",
            "slipping": "Previously held block is being lost, not intentionally placed on support.",
            "uncertain": "Grasp not yet verified: fingers closed but lift verification pending.",
        },
        "noul": {
            "allow": "The offered next skill is safe with sufficient evidence for THAT stage.",
            "deny": "Explicit danger, safety fault or violated precondition for the next skill.",
            "uncertain": "Evidence required for the next skill is missing or ambiguous.",
        },
        "reason": {
            "none": "No safety concern and no specific progress finding.",
            "unsafe": "Explicit danger or safety fault motivates a veto/stop.",
            "insufficient_evidence": "Missing stage evidence motivates uncertainty or a stop.",
            "tracking_error": "Measured alignment or tracking failure.",
            "verified_progress": "Physical evidence supports progress at this skill boundary.",
        },
    }
    result = await post(
        session,
        conn,
        "",
        {
            "model": conn.model,
            **(
                {"provider": {"allow_fallbacks": False}}
                if conn.protocol == "openrouter-decisions"
                else {}
            ),
            "state": json.dumps(
                {
                    "observation": observation,
                    "failure": failure,
                    "transition": transition,
                    "candidates": request["candidates"],
                    "nextSkillCriteria": {
                        skill: SKILL_CRITERIA[skill] for skill in request["candidates"]
                    },
                },
                ensure_ascii=False,
            ),
            "questions": {
                name: question(values, prompts[name], criteria.get(name))
                for name, values in options.items()
            },
        },
    )
    value = {
        "version": version,
        **{name: answer(result, name, values) for name, values in options.items()},
    }
    return validate_jev(value, request["candidates"], version), usage(result)