"""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()