source/bimanual/yam/envs/base_env.py

"""YamTaskEnv: the environment a task runs in.

Lifecycle mirrors ManiSkill's, so a task file reads the same way:

    _load_scene()          build the props and objects this task needs
    _initialize_episode()  per-episode randomization + settling
    solve()                run the task's solver (scripted, not a policy)
    evaluate()             return {"success": bool, ...}

The env owns the simulator, both arms and the recorder. A task subclass owns *what* is in the
scene and *what counts as done* -- it never touches action vectors or IK.
"""
from __future__ import annotations

import os
from pathlib import Path

import numpy as np

from .scene import SceneBuilder, Randomizer, TABLE_TOP

REPO = Path(__file__).resolve().parents[4]


class YamTaskEnv:
    # --- task metadata, overridden per task file (ManiSkill puts these as class attrs too) ---
    task_name = "unnamed"
    title = ""
    tags: list[str] = []
    instruction: dict = {}
    bimanual = False                 # False keeps the idle left arm frozen out of the way
    gripper_effort: float | None = None
    gripper_damping: float | None = None
    viewer_eye = (0.9, -0.9, 1.15)
    viewer_lookat = (0.05, 0.0, 0.5)
    rw_objects: dict = {}            # name -> "usd_subpath:scale:mass", registered before sim start
    rw_articulations: dict = {}      # name -> "usd_subpath:scale" for jointed props (doors, lids)

    def __init__(self, seed: int = 0, randomize: bool = True, episode: int = -1, video: str = "",
                 overrides: dict | None = None):
        from . import config as cfg
        # class attribute -> tasks/configs/<name>.yaml -> --set, weakest first
        cfg.apply(self, cfg.load(self.task_name), source=f"configs/{self.task_name}.yaml")
        cfg.apply(self, overrides or {}, source="--set")
        self.seed = seed
        self.episode = episode
        self.video_path = video or f"outputs/tasks/{self.task_name}.mp4"
        self.rand = Randomizer(seed=seed, enabled=randomize)
        self.env = None
        self.scene = None
        self.arms = {}
        self.recorder = None
        self._start_z, self._peak_z, self._peak_tilt = {}, {}, {}

    # ------------------------------------------------------------------ pre-sim
    @classmethod
    def env_vars(cls, overrides: dict | None = None) -> dict:
        """Vars that must be set BEFORE the simulator starts (asset registration, clamp force).

        `overrides` is the already-merged YAML + --set mapping. It is read here rather than on the
        instance because these decide what the simulator loads, and the instance does not exist
        until after the simulator is up.
        """
        o = overrides or {}
        rw = o.get("rw_objects", cls.rw_objects)
        artic = o.get("rw_articulations", cls.rw_articulations)
        out = {}
        if rw:
            out["YAM_RW_OBJECTS"] = ",".join(f"{k}={v}" for k, v in rw.items())
        if artic:
            out["YAM_RW_ARTIC"] = ",".join(f"{k}={v}" for k, v in artic.items())
        if o.get("gripper_effort", cls.gripper_effort):
            out["YAM_GRIP_EFFORT"] = str(o.get("gripper_effort", cls.gripper_effort))
        if o.get("gripper_damping", cls.gripper_damping):
            out["YAM_GRIP_DAMPING"] = str(o.get("gripper_damping", cls.gripper_damping))
        return out

    # ------------------------------------------------------------------ build
    def build(self, env, origin):
        from ..motion import ArmController, Recorder
        from ..motion.arm import grasp_quat

        self.env = env
        u = env.unwrapped
        self.origin = np.asarray(origin, float)
        self.scene = SceneBuilder(env, origin, self.rand)

        R, L = u.scene["right_robot"], u.scene["left_robot"]
        self._R, self._L = R, L
        self._dev = R.data.root_pos_w.device
        self._cmd = {"l": None, "r": None}
        self._grip = {"l": 1.0, "r": 1.0}
        self._quat = {"l": grasp_quat("y"), "r": grasp_quat("y")}

        def step_fn(arm, pos, quat, grip):
            side = "r" if arm.name == "right" else "l"
            self._cmd[side] = np.asarray(pos, np.float32)
            self._quat[side] = np.asarray(quat, np.float32)
            self._grip[side] = float(grip)
            self.step()

        for side, art, nm in (("right", R, "right"), ("left", L, "left")):
            self.arms[side] = ArmController(
                art, art.data.body_names,
                art.data.root_pos_w[0].cpu().numpy()-self.origin,
                art.data.root_quat_w[0].cpu().numpy(), self.origin,
                step_fn, on_step=self._on_step, name=nm)
        self._cmd["r"] = self.arms["right"].eef()
        self._cmd["l"] = self.arms["left"].eef()
        self._lhome = L.data.joint_pos[0].clone()

        self.recorder = Recorder(env, title=self.title or self.task_name, episode=self.episode)
        self.recorder.lines_fn = self._hud_lines

        self._load_scene()
        self._initialize_episode()
        self._boost_friction()
        self._start_quat = {}
        for n in self.scene.objects:
            z = float(self.scene.object_pos(n)[2])
            self._start_z[n], self._peak_z[n] = z, z
            # remember how it was RESTING, so "level" means "did not tilt from there"
            self._start_quat[n] = self.scene.objects[n].data.root_quat_w[0].cpu().numpy().copy()
        # a jointed fixture's success is "the joint MOVED", so record where it started closed
        self._start_joints = self._joint_state()
        print(f"[env] {self.task_name}: built (seed={self.seed}, randomize={self.rand.enabled}) "
              f"placements={ {k: np.round(v, 4).tolist() if hasattr(v, '__len__') else round(v, 4) for k, v in self.rand.log.items()} }",
              flush=True)

    # ---- to be provided by the task file ----
    def _load_scene(self):
        raise NotImplementedError

    def _initialize_episode(self):
        """Default: settle everything and re-seat objects at their measured height."""
        self.scene.reseat_objects(self._placed, self.step)

    def solve(self):
        raise NotImplementedError

    def evaluate(self) -> dict:
        raise NotImplementedError

    # ------------------------------------------------------------------ sim plumbing
    def step(self):
        import torch
        act = np.concatenate([self._cmd["l"], self._quat["l"], [self._grip["l"]],
                              self._cmd["r"], self._quat["r"], [self._grip["r"]]])
        self.env.step(torch.tensor(act, dtype=torch.float32, device=self._dev).view(1, -1))
        if not self.bimanual:
            self._L.write_joint_state_to_sim(
                self._lhome.view(1, -1), torch.zeros((1, self._lhome.shape[0]), device=self._dev))

    def drive_both(self, l_pos, r_pos, l_grip, r_grip):
        """Command BOTH arms and step once.

        Driving them one after the other makes the first arm park in the second one's path, and
        it also steps the sim twice per waypoint. A two-arm task must move on a single profile.
        """
        L, R = self.arms["left"], self.arms["right"]
        for arm, pos, side in ((L, l_pos, "l"), (R, r_pos, "r")):
            arm._cmd = np.asarray(pos, np.float32)
            self._cmd[side] = (np.asarray(pos, np.float32)+arm._corr).astype(np.float32)
            self._quat[side] = arm.quat
        self._grip["l"], self._grip["r"] = float(l_grip), float(r_grip)
        self.step()
        for arm in (L, R):
            e = arm._cmd-arm.eef()
            e = np.where(np.abs(e) > 0.008, e, 0.0)
            arm._corr = np.clip(arm._corr+0.08*e, -0.10, 0.10)
            arm._corr[2] = max(float(arm._corr[2]), -0.06)
        self._on_step()

    def move_both(self, l_target, r_target, l_grip, r_grip, steps=120):
        """Eased simultaneous move of both arms to their targets."""
        from ..motion.arm import ease
        L, R = self.arms["left"], self.arms["right"]
        ls = L._seg_start() if l_target is not None else None
        rs = R._seg_start() if r_target is not None else None
        lt = np.asarray(l_target, np.float32) if l_target is not None else ls
        rt = np.asarray(r_target, np.float32) if r_target is not None else rs
        for k in range(steps):
            a = ease((k+1)/float(steps))
            self.drive_both((1-a)*ls+a*lt, (1-a)*rs+a*rt, l_grip, r_grip)
        return (float(np.linalg.norm(L.eef()-lt)), float(np.linalg.norm(R.eef()-rt)))

    capture_every = 2        # sim steps per recorded frame; 2 matches the old 14 fps videos

    def _on_step(self):
        for n in self._peak_z:
            self._peak_z[n] = max(self._peak_z[n], float(self.scene.object_pos(n)[2]))
        # peak tilt away from the resting pose, for pours: the vessel is upright again by the
        # time the episode ends, so only the running maximum records that a pour happened
        for n, q0 in getattr(self, "_start_quat", {}).items():
            q = self.scene.objects[n].data.root_quat_w[0].cpu().numpy()
            d = abs(float(np.dot(q/np.linalg.norm(q), q0/np.linalg.norm(q0))))
            ang = float(np.degrees(2*np.arccos(min(1.0, d))))
            self._peak_tilt[n] = max(self._peak_tilt.get(n, 0.0), ang)
        # Record here, not only at the end: without this the video is just the final few frames
        # (12 vs 183), which reads as a broken, jumpy clip.
        self._tick = getattr(self, "_tick", 0)+1
        if self.recorder is not None and self._tick % self.capture_every == 0:
            self.recorder.capture()

    def _boost_friction(self, s=1.6, d=1.4):
        import torch
        def boost(view, tag):
            try:
                m = view.get_material_properties().clone(); m[..., 0] = s; m[..., 1] = d
                view.set_material_properties(m, torch.arange(m.shape[0], dtype=torch.int32, device=m.device))
            except Exception as e:
                print(f"[env] friction set failed on {tag}: {e}", flush=True)
        boost(self._R.root_physx_view, "right_robot"); boost(self._L.root_physx_view, "left_robot")
        for n, o in self.scene.objects.items():
            boost(o.root_physx_view, n)

    def _hud_lines(self):
        return [f"{n}=({p[0]:+.2f},{p[1]:+.2f},{p[2]:.2f})"
                for n, p in ((n, self.scene.object_pos(n)) for n in list(self.scene.objects)[:3])]

    # ------------------------------------------------------------------ scoring helpers
    def state(self):
        """The dict the condition predicates read."""
        return {"objects": {n: self.scene.object_pos(n) for n in self.scene.objects},
                "quats": {n: self.scene.objects[n].data.root_quat_w[0].cpu().numpy()
                          for n in self.scene.objects},
                "regions": self.scene.regions, "start_quats": getattr(self, "_start_quat", {}),
                "start_z": self._start_z, "peak_z": self._peak_z,
                "peak_tilt": self._peak_tilt,
                "joints": self._joint_state(), "start_joints": getattr(self, "_start_joints", {}),
                "joint_delta": getattr(self, "_joint_delta", {}),
                "links": self._link_state(),
                "joint_sane": {n: bool(
                    (art.data.joint_pos[0] >= art.data.joint_pos_limits[0, :, 0]-0.05).all()
                    and (art.data.joint_pos[0] <= art.data.joint_pos_limits[0, :, 1]+0.05).all())
                    for n, art in getattr(self.scene, "articulations", {}).items()}}

    def _link_state(self):
        """(fixture, link index) -> world xyz, for chains whose shape is the thing being judged."""
        out = {}
        for n, art in getattr(self.scene, "articulations", {}).items():
            p = art.data.body_pos_w[0].cpu().numpy()-self.origin
            for i in range(len(p)):
                out[(n, i)] = p[i]
        return out

    def _joint_state(self):
        """(fixture, joint index) -> value, for every jointed prop in the scene."""
        out = {}
        for n, art in getattr(self.scene, "articulations", {}).items():
            q = art.data.joint_pos[0].cpu().numpy()
            for i in range(len(q)):
                out[(n, i)] = float(q[i])
        return out

    def check(self, *conditions) -> dict:
        """Run condition predicates against the current state; returns {"success": ..., per-cond}."""
        from .. import conditions as C
        st = self.state()
        per = {}
        for c in conditions:
            label = getattr(c, "label", getattr(c, "__name__", "cond"))
            try:
                per[label] = bool(c(st))
            except Exception as e:
                print(f"[env] condition {label} errored: {e}", flush=True)
                per[label] = False
        ok = all(per.values()) if per else False
        self.recorder.result = "SUCCESS" if ok else "FAIL"
        for _ in range(14):
            self.recorder.capture()
        print(f"[env] EPISODE_RESULT: {self.recorder.result}", flush=True)
        for k, v in per.items():
            print(f"[env]   {'PASS' if v else 'FAIL'}  {k}", flush=True)
        return {"success": ok, **per}

    def save_video(self, path=None):
        p = path or self.video_path
        return self.recorder.save(str(REPO/p) if not os.path.isabs(p) else p)