91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
"""Regenerate CPU mj_ray references: run from repository root in .venv."""
|
|
|
|
import json
|
|
import math
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import mujoco
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
from task_config import OBSTACLE_TASK, deployment_metadata, validate_task_config # noqa: E402
|
|
|
|
|
|
def generate():
|
|
config = validate_task_config(
|
|
OBSTACLE_TASK, {"sensorCfg": {"sensorMode": "multi_ring_raycast"}}, 42
|
|
)
|
|
deployment = deployment_metadata(OBSTACLE_TASK, config, 42)
|
|
boxes = deployment["terrain"]["boxes"]
|
|
# Low obstacle directly ahead; platform edge is a bounded-floor drop (no invented height map).
|
|
layouts = {
|
|
"default": boxes,
|
|
"low": [boxes[0], {"pos": [0.9, 0, 0.025], "size": [0.2, 0.5, 0.025], "yaw": 0}],
|
|
"edge": [boxes[0]],
|
|
}
|
|
cases = []
|
|
for layout, position, euler in [
|
|
("default", [-5, 0, 0.32], [0, 0, 0]),
|
|
("default", [-3, 1, 0.42], [0.23, -0.31, 0.51]),
|
|
("default", [-2, -2, 0.5], [-0.35, 0.27, -0.63]),
|
|
("low", [0, 0, 0.32], [0, 0, 0]),
|
|
("low", [0, 0, 0.32], [0.2, 0.1, -0.1]),
|
|
("edge", [5.45, 0, 0.32], [0, 0, 0]),
|
|
]:
|
|
xml = (
|
|
"<mujoco><worldbody>"
|
|
+ "".join(
|
|
'<geom type="box" pos="{}" size="{}"/>'.format(
|
|
" ".join(map(str, b["pos"])), " ".join(map(str, b["size"]))
|
|
)
|
|
for b in layouts[layout]
|
|
)
|
|
+ "</worldbody></mujoco>"
|
|
)
|
|
model = mujoco.MjModel.from_xml_string(xml)
|
|
data = mujoco.MjData(model)
|
|
mujoco.mj_forward(model, data)
|
|
q = np.zeros(4)
|
|
mujoco.mju_euler2Quat(q, np.array(euler), "xyz")
|
|
matrix = np.zeros(9)
|
|
mujoco.mju_quat2Mat(matrix, q)
|
|
matrix = matrix.reshape(3, 3)
|
|
origin = np.array(position) + matrix @ np.array([0.3, 0, 0.05])
|
|
distances, ids = [], []
|
|
for pitch in [0, -20, -45]:
|
|
for i in range(16):
|
|
yaw = math.radians(-45 + i * 90 / 15)
|
|
p = math.radians(pitch)
|
|
direction = matrix @ np.array(
|
|
[math.cos(p) * math.cos(yaw), math.cos(p) * math.sin(yaw), math.sin(p)]
|
|
)
|
|
geom = np.array([-1], dtype=np.int32)
|
|
distance = mujoco.mj_ray(model, data, origin, direction, None, 1, -1, geom)
|
|
distances.append(distance)
|
|
ids.append(int(geom[0]))
|
|
cases.append(
|
|
dict(
|
|
layout=layout,
|
|
position=position,
|
|
quaternion=q.tolist(),
|
|
distances=distances,
|
|
hitIds=ids,
|
|
depth=[1 if d < 0 else min(1, d / 4) for d in distances],
|
|
)
|
|
)
|
|
root = Path("web_platform/src/rl/fixtures")
|
|
(root / "multiRingGolden.json").write_text(
|
|
json.dumps(
|
|
dict(source=f"CPU MuJoCo {mujoco.__version__} mj_ray", layouts=layouts, cases=cases),
|
|
indent=2,
|
|
)
|
|
+ "\n"
|
|
)
|
|
(root / "multiRingDeployment.json").write_text(json.dumps(deployment, indent=2) + "\n")
|
|
Path("/tmp/go2-multi-ring-stage5/task.json").write_text(json.dumps(config))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
generate()
|