File size: 13,264 Bytes
7399b6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 | """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)
|