feat(assets): a snapshot import places the robot at any heading in the room

robotsmith/assets/producers/aha3d_snapshot.py

@@ -1,7 +1,7 @@
 """AHa-3D scene export -> snapshot scenario.
 
-The robot takes one spot in the room, facing +X; the furniture it works at becomes its table,
-and the roots it can use become assets at their recorded positions:
+The robot takes one spot in the room, facing ``--yaw`` degrees from room +X; the furniture it works
+at becomes its table, and the roots it can use become assets at their recorded positions:
 
     python -m robotsmith.assets.producers.aha3d_snapshot <export_dir> --support furniture/desk \\
         --base 0.62 -1.30 --prefix office48 --scenario scenarios/office48_desk.yaml --task task.yaml
@@ -30,13 +30,25 @@
 FLOOR_TOL_M = 0.05
 
 
-def plan(scene: dict, support: str, base_xy) -> dict:
+def to_robot(xy, base_xy, yaw_deg: float = 0.0) -> np.ndarray:
+    """Room-frame xy point(s) in the robot frame: origin at the base, +X along the robot's heading."""
+    c, s = np.cos(np.radians(yaw_deg)), np.sin(np.radians(yaw_deg))
+    return (np.asarray(xy, dtype=float) - np.asarray(base_xy, dtype=float)) @ np.array([[c, -s], [s, c]])
+
+
+def box_to_robot(lo, hi, base_xy, yaw_deg: float = 0.0) -> tuple[np.ndarray, np.ndarray]:
+    """Robot-frame xy bounds of a room-frame box (bounds of its rotated corners)."""
+    corners = to_robot([[lo[0], lo[1]], [lo[0], hi[1]], [hi[0], lo[1]], [hi[0], hi[1]]], base_xy, yaw_deg)
+    return corners.min(axis=0), corners.max(axis=0)
+
+
+def plan(scene: dict, support: str, base_xy, yaw_deg: float = 0.0) -> dict:
     """Each root's fate: ``table`` (replaced by the robot's table), ``support`` (rests on it),
     ``floor`` or ``dropped`` with a reason; positions in the robot frame."""
     base = np.asarray(base_xy, dtype=float)
     roots = {r["instance_id"]: r for r in scene["roots"]}
     desk = roots[support]
-    rel = {i: (np.asarray(r["lo"][:2]) - base, np.asarray(r["hi"][:2]) - base) for i, r in roots.items() if "lo" in r}
+    rel = {i: box_to_robot(r["lo"], r["hi"], base, yaw_deg) for i, r in roots.items() if "lo" in r}
     out = {}
     for iid, r in roots.items():
         if iid == support:
@@ -59,21 +71,21 @@ def plan(scene: dict, support: str, base_xy) -> dict:
             out[iid] = ({"decision": "floor"} if np.linalg.norm(gap) <= FLOOR_REACH_M
                         else {"decision": "dropped", "reason": f"more than {FLOOR_REACH_M} m from the robot"})
         if "origin" in r:
-            out[iid]["anchor_xy"] = [round(float(v), 4) for v in np.asarray(r["origin"][:2]) - base]
+            out[iid]["anchor_xy"] = [round(float(v), 4) + 0.0 for v in to_robot(r["origin"][:2], base, yaw_deg)]
     size = table_size([rel[i] for i, d in out.items() if d["decision"] == "support"])
     cx, cy = table_center_xy(DEFAULT_WORKSPACE_XY)
     t_lo, t_hi = np.array([cx - size[0] / 2, cy - size[1] / 2]), np.array([cx + size[0] / 2, cy + size[1] / 2])
     for iid, d in out.items():
         if d["decision"] == "floor" and np.all(rel[iid][0] < t_hi) and np.all(rel[iid][1] > t_lo):
             out[iid] = {"decision": "dropped", "reason": "overlaps the robot's table", "anchor_xy": d["anchor_xy"]}
-    return {"support": support, "base_xy_room": [float(v) for v in base], "yaw_deg": 0.0,
+    return {"support": support, "base_xy_room": [float(v) for v in base], "yaw_deg": float(yaw_deg),
             "table_xy": [t_lo.round(3).tolist(), t_hi.round(3).tolist()], "roots": out}
 
 
 def import_snapshot(export_dir: Path, *, support: str, base_xy, prefix: str, assets: Path, scenario: Path,
-                    task: dict, provenance: dict) -> dict:
+                    task: dict, provenance: dict, yaw_deg: float = 0.0) -> dict:
     scene = json.loads((export_dir / "scene.json").read_text(encoding="utf-8"))
-    mapping = plan(scene, support, base_xy)
+    mapping = plan(scene, support, base_xy, yaw_deg)
     roots = {r["instance_id"]: r for r in scene["roots"]}
     objects = []
     for iid, d in mapping["roots"].items():
@@ -85,7 +97,7 @@ def import_snapshot(export_dir: Path, *, support: str, base_xy, prefix: str, ass
                      description=f"{iid} from an AHa-3D scene ({r['semantic_class']})",
                      provenance=provenance | {"instance_id": iid, "semantic_class": r["semantic_class"]})
         objects.append({"id": r["item"], "asset": d["asset"], "on": d["decision"], "anchor_xy": d["anchor_xy"],
-                        "yaw_deg": 0.0})
+                        "yaw_deg": 0.0 - yaw_deg})
     doc = {"scenario": scenario.stem, "robot": "franka", "scene": {"support": "table_simple", "objects": objects},
            "task": task}
     scenario.write_text(yaml.safe_dump(doc, sort_keys=False, default_flow_style=None), encoding="utf-8", newline="\n")
@@ -98,6 +110,7 @@ def main() -> None:
     ap.add_argument("export_dir", type=Path)
     ap.add_argument("--support", required=True, help="instance_id of the furniture the robot works at")
     ap.add_argument("--base", type=float, nargs=2, required=True, help="robot base xy in the room frame")
+    ap.add_argument("--yaw", type=float, default=0.0, help="robot heading in the room frame, degrees from +X")
     ap.add_argument("--prefix", required=True)
     ap.add_argument("--assets", type=Path, default=Path("assets"))
     ap.add_argument("--scenario", type=Path, required=True)
@@ -107,7 +120,7 @@ def main() -> None:
     prov = json.loads(a.provenance.read_text(encoding="utf-8")) if a.provenance else {}
     mapping = import_snapshot(a.export_dir, support=a.support, base_xy=a.base, prefix=a.prefix, assets=a.assets,
                               scenario=a.scenario, task=yaml.safe_load(a.task.read_text(encoding="utf-8")),
-                              provenance=prov)
+                              provenance=prov, yaw_deg=a.yaw)
     kept = {i: d for i, d in mapping["roots"].items() if d["decision"] in ("support", "floor")}
     print(json.dumps(kept, indent=1))
 

tests/test_aha3d_snapshot.py

@@ -43,3 +43,15 @@ def test_plan_sorts_every_root_and_keeps_room_positions():
 def test_plan_expresses_positions_relative_to_the_robot_base():
     roots = plan(SCENE, "furniture/desk", (0.5, 0.0))["roots"]
     assert roots["decor/cup"]["anchor_xy"] == [0.74, 0.14]
+
+
+def test_plan_follows_the_robot_heading():
+    def turn(r):
+        if "lo" not in r:
+            return r
+        (x0, y0), (x1, y1) = r["lo"][:2], r["hi"][:2]
+        return _root(r["instance_id"], r["semantic_class"], [-y1, x0, r["lo"][2]], [-y0, x1, r["hi"][2]])
+
+    turned = plan({"roots": [turn(r) for r in SCENE["roots"]]}, "furniture/desk", (0.0, 0.0), yaw_deg=90.0)
+    assert turned["roots"] == plan(SCENE, "furniture/desk", (0.0, 0.0))["roots"]
+    assert turned["yaw_deg"] == 90.0