255 lines
11 KiB
Python
255 lines
11 KiB
Python
"""Authoritative custom boxes validation and CPU compilation, no training/RNG."""
|
|
|
|
import copy
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
for path in (ROOT, ROOT / "rl"):
|
|
sys.path.insert(0, str(path))
|
|
from task_config import ( # noqa: E402
|
|
OBSTACLE_TASK,
|
|
TaskConfigError,
|
|
build_terrain_layout,
|
|
deployment_metadata,
|
|
validate_custom_terrain,
|
|
validate_task_config,
|
|
)
|
|
|
|
|
|
def layout():
|
|
return json.loads((Path(__file__).parent / "fixtures/custom-boxes.json").read_text())
|
|
|
|
|
|
class CustomBoxesTest(unittest.TestCase):
|
|
def test_payload_roundtrip_without_rng_and_aliasing(self):
|
|
source = layout()
|
|
with patch("task_config.random.Random", side_effect=AssertionError("must not run RNG")):
|
|
config = validate_task_config(
|
|
OBSTACLE_TASK,
|
|
{
|
|
"terrainPreset": "custom_boxes",
|
|
"customTerrainBoxes": source,
|
|
},
|
|
42,
|
|
)
|
|
self.assertEqual(build_terrain_layout(config), source)
|
|
self.assertEqual(deployment_metadata(OBSTACLE_TASK, config, 42)["terrain"], source)
|
|
source["boxes"][1]["pos"][0] = 9
|
|
self.assertEqual(config["customTerrainBoxes"], layout())
|
|
|
|
def test_strict_fields_numbers_floor_safety_and_count(self):
|
|
mutations = [
|
|
lambda t: t.update(path="../../etc/passwd"),
|
|
lambda t: t.pop("target"),
|
|
lambda t: t.update(approximation=False),
|
|
lambda t: t.update(actualObstacleCount=True),
|
|
lambda t: t.update(actualObstacleCount=9),
|
|
lambda t: t.update(size=25),
|
|
lambda t: t.update(friction=True),
|
|
lambda t: t.update(spawnQuaternion=[2, 0, 0, 0]),
|
|
lambda t: t.update(spawn=[-2, -1, 0.4]),
|
|
lambda t: t.update(target=[6, 0]),
|
|
lambda t: t["boxes"][0]["pos"].__setitem__(2, 0),
|
|
lambda t: t["boxes"][1].update(mesh="../../model.stl"),
|
|
lambda t: t["boxes"][1].update(yaw=True),
|
|
lambda t: t["boxes"][1].update(pos=[6, 2, 0.5]),
|
|
lambda t: t["boxes"][1].update(pos=[1, 2, -0.3]),
|
|
lambda t: t["boxes"][1].update(pos=[1, 2, 12]),
|
|
lambda t: t.update(
|
|
boxes=t["boxes"] + [copy.deepcopy(t["boxes"][1])] * 256, actualObstacleCount=257
|
|
),
|
|
]
|
|
for bad in (float("nan"), float("inf"), -float("inf"), 10**400, 0, -0.0, -1, True):
|
|
mutations.append(lambda t, bad=bad: t["boxes"][1]["size"].__setitem__(0, bad))
|
|
for mutate in mutations:
|
|
value = layout()
|
|
mutate(value)
|
|
with self.subTest(value=str(value)[:200]), self.assertRaises(TaskConfigError):
|
|
validate_custom_terrain(value)
|
|
for key in ("spawn", "target"):
|
|
value = layout()
|
|
# Circle tangent to the right face: reject; floor alone is exempt.
|
|
value[key][:2] = [1.9, 2]
|
|
with self.assertRaisesRegex(TaskConfigError, "安全区"):
|
|
validate_custom_terrain(value)
|
|
value[key][:2] = [1.91, 2]
|
|
validate_custom_terrain(value)
|
|
# Precise circular corner test, not an expanded square approximation.
|
|
value[key][:2] = [1.8, 2.7]
|
|
validate_custom_terrain(value)
|
|
|
|
def test_conflict_missing_path_and_training_entry(self):
|
|
from scripts.train import _load_task_config
|
|
|
|
for payload in (
|
|
{"terrainPreset": "custom_boxes"},
|
|
{"terrainPreset": "custom_boxes", "customTerrainBoxes": "../../file.json"},
|
|
{"terrainPreset": "plane", "customTerrainBoxes": layout()},
|
|
{
|
|
"terrainPreset": "custom_boxes",
|
|
"customTerrainBoxes": layout(),
|
|
"terrainParams": {"size": 8, "friction": 0.8},
|
|
},
|
|
):
|
|
with self.assertRaises(TaskConfigError):
|
|
validate_task_config(OBSTACLE_TASK, payload, 42)
|
|
config = validate_task_config(
|
|
OBSTACLE_TASK, {"terrainPreset": "custom_boxes", "customTerrainBoxes": layout()}, 42
|
|
)
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "task.json"
|
|
path.write_text(json.dumps(config))
|
|
self.assertEqual(_load_task_config(OBSTACLE_TASK, str(path), 42), config)
|
|
config["customTerrainBoxes"]["boxes"][1]["size"][0] = -0.0
|
|
path.write_text(json.dumps(config))
|
|
with self.assertRaises(TaskConfigError):
|
|
_load_task_config(OBSTACLE_TASK, str(path), 42)
|
|
config["customTerrainBoxesPath"] = "/etc/passwd"
|
|
path.write_text(json.dumps(config))
|
|
with self.assertRaises(ValueError):
|
|
_load_task_config(OBSTACLE_TASK, str(path), 42)
|
|
|
|
def test_server_payload_and_limit(self):
|
|
from server import DEFAULT_TASKS, MAX_REQUEST_BYTES, ApiError, TrainingManager
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
manager = TrainingManager(
|
|
Path(directory), sys.executable, DEFAULT_TASKS, check_environment=False
|
|
)
|
|
payload = dict(
|
|
taskId=OBSTACLE_TASK,
|
|
terrainPreset="custom_boxes",
|
|
customTerrainBoxes=layout(),
|
|
numEnvs=2,
|
|
maxIterations=1,
|
|
seed=42,
|
|
device="cpu",
|
|
gpuIds=[],
|
|
)
|
|
config = manager.parse_config(payload)
|
|
self.assertEqual(config.deployment["terrain"], layout())
|
|
self.assertEqual(config.task_config["customTerrainBoxes"], layout())
|
|
self.assertIn("custom_boxes", manager.health()["taskMetadata"][2]["terrainPresets"])
|
|
self.assertIn(
|
|
"--task-config", manager.command_for(config, Path(directory) / "server-owned.json")
|
|
)
|
|
with self.assertRaises(ApiError):
|
|
manager.parse_config({**payload, "customTerrainBoxesPath": "/etc/passwd"})
|
|
full = layout()
|
|
full["boxes"] += [copy.deepcopy(full["boxes"][1])] * 255
|
|
full["actualObstacleCount"] = 256
|
|
# Worst typical double precision expansion still fits the bounded request envelope.
|
|
full["boxes"][1:] = [
|
|
dict(
|
|
pos=[1.123456789012345, 2.123456789012345, 0.5123456789012345],
|
|
size=[0.4123456789012345, 0.3123456789012345, 0.5123456789012345],
|
|
yaw=0,
|
|
)
|
|
for _ in range(256)
|
|
]
|
|
self.assertLess(
|
|
len(json.dumps({**payload, "customTerrainBoxes": full}).encode()), MAX_REQUEST_BYTES
|
|
)
|
|
manager.parse_config({**payload, "customTerrainBoxes": full})
|
|
|
|
def test_http_body_length_is_bounded_and_allows_full_layout(self):
|
|
import io
|
|
|
|
from server import MAX_REQUEST_BYTES, ApiError, TrainingRequestHandler
|
|
|
|
value = layout()
|
|
value["boxes"] += [copy.deepcopy(value["boxes"][1])] * 255
|
|
value["actualObstacleCount"] = 256
|
|
body = json.dumps({"terrainPreset": "custom_boxes", "customTerrainBoxes": value}).encode()
|
|
handler = object.__new__(TrainingRequestHandler)
|
|
handler.headers = {"Content-Length": str(len(body))}
|
|
handler.rfile = io.BytesIO(body)
|
|
self.assertEqual(handler._payload()["customTerrainBoxes"], value)
|
|
for size in (0, MAX_REQUEST_BYTES + 1):
|
|
handler.headers = {"Content-Length": str(size)}
|
|
with self.assertRaises(ApiError):
|
|
handler._payload()
|
|
|
|
def test_real_cpu_compilation_and_custom_origin_quaternion_goal(self):
|
|
import mujoco
|
|
import numpy as np
|
|
from mjlab.terrains import TerrainGenerator
|
|
from src.tasks.obstacle_avoidance.env_cfg import (
|
|
apply_obstacle_configuration,
|
|
unitree_go2_obstacle_env_cfg,
|
|
)
|
|
|
|
value = layout()
|
|
config = validate_task_config(
|
|
OBSTACLE_TASK, {"terrainPreset": "custom_boxes", "customTerrainBoxes": value}, 42
|
|
)
|
|
cfg = unitree_go2_obstacle_env_cfg()
|
|
apply_obstacle_configuration(cfg, config)
|
|
generator = TerrainGenerator(cfg.scene.terrain.terrain_generator)
|
|
spec = mujoco.MjSpec()
|
|
generator.compile(spec)
|
|
model = spec.compile()
|
|
np.testing.assert_allclose(model.geom_pos, [b["pos"] for b in value["boxes"]], atol=1e-12)
|
|
np.testing.assert_allclose(model.geom_size, [b["size"] for b in value["boxes"]], atol=1e-12)
|
|
np.testing.assert_allclose(generator.terrain_origins[0, 0], [-2, -1, 0])
|
|
self.assertEqual(
|
|
cfg.scene.entities["robot"].init_state.rot, tuple(value["spawnQuaternion"])
|
|
)
|
|
self.assertEqual(cfg.commands["twist"].goal_offset, (4, 0))
|
|
self.assertEqual(cfg.terminations["outside_map"].params, {"size": 12})
|
|
self.assertTrue(math.isfinite(model.geom_size.sum()))
|
|
|
|
@unittest.skipUnless(os.environ.get("GO2_CUSTOM_BOXES_SMOKE") == "1", "opt-in 2env/1step smoke")
|
|
def test_two_env_one_step_custom_layout(self):
|
|
import numpy as np
|
|
import torch
|
|
import warp as wp
|
|
from mjlab.envs import ManagerBasedRlEnv
|
|
from src.tasks.obstacle_avoidance.env_cfg import (
|
|
apply_obstacle_configuration,
|
|
unitree_go2_obstacle_env_cfg,
|
|
)
|
|
from warp._src import context
|
|
|
|
if not torch.cuda.is_available():
|
|
self.skipTest("CUDA unavailable")
|
|
if not hasattr(wp, "context"):
|
|
wp.context = context
|
|
cfg = unitree_go2_obstacle_env_cfg()
|
|
custom = validate_task_config(
|
|
OBSTACLE_TASK, {"terrainPreset": "custom_boxes", "customTerrainBoxes": layout()}, 42
|
|
)
|
|
apply_obstacle_configuration(cfg, custom, randomize_navigation=False)
|
|
cfg.scene.num_envs = 2
|
|
env = ManagerBasedRlEnv(cfg, device="cuda:0")
|
|
try:
|
|
obs, _ = env.reset()
|
|
self.assertEqual(tuple(obs["actor"].shape), (2, 81))
|
|
np.testing.assert_allclose(env.scene.env_origins.cpu(), [[-2, -1, 0]] * 2, atol=1e-6)
|
|
np.testing.assert_allclose(
|
|
env.scene["robot"].data.root_link_pos_w.cpu(), [layout()["spawn"]] * 2, atol=1e-6
|
|
)
|
|
np.testing.assert_allclose(
|
|
env.scene["robot"].data.root_link_quat_w.cpu(),
|
|
[layout()["spawnQuaternion"]] * 2,
|
|
atol=1e-6,
|
|
)
|
|
np.testing.assert_allclose(
|
|
env.command_manager.get_term("twist").errors()[1].cpu(), [4, 4], atol=1e-6
|
|
)
|
|
obs, reward, terminated, truncated, _ = env.step(
|
|
torch.zeros((2, 12), device=env.device)
|
|
)
|
|
self.assertTrue(torch.isfinite(obs["actor"]).all())
|
|
self.assertTrue(torch.isfinite(reward).all())
|
|
self.assertFalse(terminated.any() or truncated.any())
|
|
finally:
|
|
env.close()
|