"""Obstacle-specific schema, orchestration, objective math and opt-in real rollout.""" import json import os import sys import tempfile import threading import unittest from copy import deepcopy from pathlib import Path from types import SimpleNamespace from unittest.mock import patch SERVICE = Path(__file__).resolve().parents[1] for root in (SERVICE, SERVICE / "rl"): sys.path.insert(0, str(root)) from task_config import build_terrain_layout, validate_task_config # noqa: E402 from tuning import obstacle_scoring as scoring # noqa: E402 from tuning.advisor import DeepSeekAdvisor # noqa: E402 from tuning.manager import TuningError, TuningManager # noqa: E402 from tuning.process import GpuLease, ResourceBusyError # noqa: E402 from tuning.schema import ( # noqa: E402 OBSTACLE_TASK, RewardConfigError, apply_reward_configuration, base_configuration, merge_proposal, validate_configuration, validate_constraints, validate_proposal, ) from tuning.scoring import EvaluationError # noqa: E402 def sample(**kw): return ( dict( distance=1.0, clearance=0.5, action_delta=0.0, ray_hit=0.0, collision=0.0, fall=0.0, terminal=0.0, **{}, ) | kw ) def evaluation(custom, n=2): metrics = scoring.score_trajectory([sample()] * scoring.STEPS) return { "protocol": scoring.protocol(custom, n), "metrics": metrics, "seedMetrics": [ {"seed": seed, "episodes": n, "rolloutSteps": scoring.STEPS, "metrics": metrics} for seed in scoring.SEEDS ], } class ObstacleSchemaTest(unittest.TestCase): def test_ranges_task_isolation_nonfinite_and_patch_no_mutation(self): base = base_configuration(OBSTACLE_TASK) original = deepcopy(base) for section, key, low, high in ( ("weights", "avoidance_weight", 0.5, 5), ("weights", "collision_penalty", -10, -0.5), ("weights", "action_smoothness", -0.05, -0.001), ("params", "target_velocity", 0.3, 1.2), ): for value in (low, high): changed = deepcopy(base) changed[section][key] = value self.assertEqual(validate_configuration(changed, OBSTACLE_TASK), changed) for value in (low - 0.00001, high + 0.00001, float("nan"), float("inf"), True): with self.subTest(key=key, value=value), self.assertRaises(RewardConfigError): validate_proposal({section: {key: value}}, base, task_id=OBSTACLE_TASK) for proposal in ( {"weights": {"pose": 1.1}}, {"weights": {"target_velocity": 0.8}}, {"params": {"seed": 101}}, {"unknown": {}}, {"weights": {"avoidance_weight": 4.01}}, ): with self.assertRaises(RewardConfigError): validate_proposal(proposal, base, task_id=OBSTACLE_TASK) with self.assertRaises(RewardConfigError): validate_configuration(base) with self.assertRaises(RewardConfigError): validate_configuration(base_configuration(), OBSTACLE_TASK) with self.assertRaises(RewardConfigError): validate_constraints({"weights.pose": {"kind": "fixed", "value": 1}}, OBSTACLE_TASK) candidate = merge_proposal( base, {"params": {"target_velocity": 1.2}}, task_id=OBSTACLE_TASK ) self.assertEqual(candidate["params"]["target_velocity"], 1.2) self.assertEqual(base, original) def test_real_training_config_mapping_and_command(self): import torch from scripts.train import ( _configure_task_and_rewards, _load_reward_config, _load_task_config, ) from src.tasks.obstacle_avoidance.env_cfg import ( apply_obstacle_configuration, unitree_go2_obstacle_env_cfg, ) base = base_configuration(OBSTACLE_TASK) base["weights"].update(avoidance_weight=3, collision_penalty=-7, action_smoothness=-0.02) base["params"]["target_velocity"] = 0.9 custom = validate_task_config(OBSTACLE_TASK, {"sensorCfg": {"fov": 60}}, 42) with tempfile.TemporaryDirectory() as directory: reward, task = Path(directory) / "reward.json", Path(directory) / "task.json" reward.write_text(json.dumps(base)) task.write_text(json.dumps(custom)) cfg = unitree_go2_obstacle_env_cfg() apply_obstacle_configuration(cfg, _load_task_config(OBSTACLE_TASK, str(task), 42)) apply_reward_configuration( cfg, _load_reward_config(str(reward), None, OBSTACLE_TASK), OBSTACLE_TASK ) deployment, checked = _configure_task_and_rewards( OBSTACLE_TASK, SimpleNamespace( env=unitree_go2_obstacle_env_cfg(), task_config=str(task), reward_config=str(reward), reward_config_json=None, agent=SimpleNamespace(seed=42), ), ) self.assertEqual(checked, base) self.assertEqual(deployment["navigation"]["speed"], 0.9) self.assertEqual(deployment["sensorCfg"]["avoidanceWeight"], 3) self.assertEqual(cfg.rewards["obstacle_proximity"].weight, -3) self.assertEqual(cfg.rewards["obstacle_collision"].weight, -7) self.assertIs( cfg.rewards["obstacle_collision"].func, cfg.terminations["illegal_contact"].func ) self.assertEqual(cfg.rewards["action_rate_l2"].weight, -0.02) self.assertNotIn("target_velocity", cfg.rewards) self.assertEqual(cfg.commands["twist"].speed, 0.9) class Scene(dict): env_origins = torch.tensor([[-5.0, 0, 0]]) scene = Scene( robot=SimpleNamespace( data=SimpleNamespace( root_link_pos_w=torch.tensor([[-5.0, 0, 0.32]]), root_link_quat_w=torch.tensor([[1.0, 0, 0, 0]]), ) ) ) command = cfg.commands["twist"].build( SimpleNamespace(num_envs=1, device="cpu", scene=scene) ) self.assertAlmostEqual(command.command[0, 0].item(), 0.9, places=6) def test_mock_deepseek_task_context_no_network(self): advisor = DeepSeekAdvisor() output = SimpleNamespace( weights={"avoidance_weight": 2.2}, params={"target_velocity": 0.7}, rationale="绕行与目标导航", expected_impact={}, confidence=0.8, ) fake = SimpleNamespace(run_sync=lambda prompt: SimpleNamespace(output=output)) with patch.object(advisor, "_agent", return_value=fake): result = advisor.propose({"task": OBSTACLE_TASK}, base_configuration(OBSTACLE_TASK)) self.assertEqual(result["patch"]["params"], {"target_velocity": 0.7}) class ObstacleScoringTest(unittest.TestCase): def test_hand_calculated_scores_and_first_terminal(self): samples = [ sample(clearance=0.25, action_delta=0.5), sample(distance=0.4, clearance=0.25, action_delta=0.5), ] metrics = scoring.score_trajectory(samples, 2) self.assertEqual(metrics["success"], 1) self.assertEqual(metrics["time"], 0) self.assertEqual(metrics["clearance"], 0.5) self.assertEqual(metrics["smooth"], 0.5) self.assertAlmostEqual(scoring.score_evaluation(metrics)["score"], 0.65) failed = scoring.score_trajectory([sample(distance=0.1, terminal=1, fall=1)], 1000) self.assertEqual((failed["success"], failed["time"], failed["clearance"]), (0, 0, 0)) contact = scoring.score_trajectory([sample(distance=0.1, terminal=1, collision=1)], 1000) self.assertEqual((contact["arrival_rate"], contact["success"], contact["time"]), (1, 0, 0)) with self.assertRaises(EvaluationError): scoring.score_trajectory([sample()], 1000) with self.assertRaises(EvaluationError): scoring.score_trajectory([sample(terminal=1), sample()], 2) with self.assertRaises(EvaluationError): scoring.score_trajectory([sample(clearance=float("nan"))], 1) def test_recorder_excludes_floor_and_keeps_terminal_before_reset(self): import torch from scripts.evaluate_obstacle import FirstEpisodeRecorder class Scene(dict): env_origins = torch.zeros((2, 3)) robot = SimpleNamespace( root_link_pos_w=torch.tensor([[0.0, 0, 0.32], [0.0, 0, 0.1]]), projected_gravity_b=torch.tensor([[0.0, 0, -1.0], [0.0, 1.0, 0.0]]), ) scene = Scene( robot=SimpleNamespace(data=robot), nonfoot_ground_touch=SimpleNamespace( data=SimpleNamespace(force_history=torch.zeros((2, 1, 4, 3))) ), forward_scan=SimpleNamespace( data=SimpleNamespace(distances=torch.ones((2, 32))), cfg=SimpleNamespace(max_distance=4), ), ) command = SimpleNamespace(errors=lambda: (None, torch.tensor([0.2, 0.2]), None)) env = SimpleNamespace( scene=scene, device="cpu", num_envs=2, action_manager=SimpleNamespace(action=torch.ones((2, 12)) * 0.5), command_manager=SimpleNamespace(get_term=lambda name: command), termination_manager=SimpleNamespace(compute=lambda: torch.ones(2, dtype=torch.bool)), ) floor = {"pos": [0, 0, -0.1], "size": [6, 6, 0.1]} recorder = FirstEpisodeRecorder(env, {"spawn": [0, 0, 0.32], "boxes": [floor]}) env.termination_manager.compute() robot.root_link_pos_w[:] = torch.tensor([0.0, 0, 0.32]) # Simulated auto-reset overwrite. robot.projected_gravity_b[:] = torch.tensor([0.0, 0, -1.0]) env.termination_manager.compute() self.assertEqual([len(s) for s in recorder.samples], [1, 1]) upright, lying = [scoring.score_trajectory(s, 2) for s in recorder.samples] self.assertEqual(upright["clearance"], 0.5) # 1 valid safe step / 2, not floor distance 0. self.assertEqual(upright["success"], 1) self.assertEqual(lying["clearance"], 0) self.assertEqual(lying["success"], 0) self.assertEqual(lying["time"], 0) self.assertEqual(lying["fall_rate"], 1) def test_threshold_boundary_and_fail_closed(self): base = scoring.score_trajectory([sample(distance=0.1)], 1) base.update(success=0.5, fall_rate=0.1, no_fall=0.9) current = dict(base, success=0.48, fall_rate=0.12, no_fall=0.88) self.assertTrue(scoring.score_evaluation(current, base)["eligible"]) for changed in ( dict(current, success=0.479999), dict(current, fall_rate=0.120001, no_fall=0.879999), ): self.assertFalse(scoring.score_evaluation(changed, base)["eligible"]) self.assertEqual(scoring.score_evaluation(changed, base)["score"], -1) custom = validate_task_config(OBSTACLE_TASK, {}, 42) expected = scoring.protocol(custom, 2) valid = evaluation(custom) scoring.validate_evaluation(valid, expected) for mutate in ( lambda v: v["seedMetrics"].pop(), lambda v: v["metrics"].pop("success"), lambda v: v["protocol"].update(stepsPerSeed=999), lambda v: v["seedMetrics"][0].update(episodes=1), lambda v: v["metrics"].update(smooth=float("inf")), ): bad = deepcopy(valid) mutate(bad) with self.assertRaises(EvaluationError): scoring.validate_evaluation(bad, expected) def test_fixed_three_seed_and_custom_authoritative_map(self): custom = validate_task_config(OBSTACLE_TASK, {}, 42) scenes = scoring.evaluation_scenarios(custom) self.assertEqual([s["seed"] for s in scenes], [101, 202, 303]) self.assertNotEqual(scenes[0]["terrain"], scenes[1]["terrain"]) layout = build_terrain_layout(custom) layout["approximation"] = True custom = validate_task_config( OBSTACLE_TASK, {"terrainPreset": "custom_boxes", "customTerrainBoxes": layout}, 42 ) fixed = scoring.protocol(custom, 2) self.assertEqual(fixed["sceneMode"], "fixed-custom-map") self.assertEqual([s["terrain"] for s in fixed["scenarios"]], [layout] * 3) self.assertEqual(fixed, scoring.protocol(deepcopy(custom), 2)) class ObstacleManagerTest(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() root = Path(self.temp.name) self.manager = TuningManager(SERVICE / "rl", sys.executable, root, GpuLease()) payload = { "taskId": OBSTACLE_TASK, "mode": "approval", "evalNumEnvs": 2, "numEnvs": 2, "initialIterations": 1, "middleIterations": 1, "finalIterations": 1, "sensorCfg": {"fov": 60, "avoidanceWeight": 3}, } mode, config, objective, fallback = self.manager.parse_create(payload) self.session = self.manager.storage.create_session(mode, config, objective, fallback) self.trial = self.manager.storage.create_trial( self.session["id"], 0, 0, 1, self.manager._base_configuration(self.session), None, "trial-000-rung-0", ) def tearDown(self): self.manager.shutdown() self.temp.cleanup() def test_argv_json_real_train_and_eval_ingress(self): commands = [] custom = self.session["config"]["taskConfig"] def run(session_id, command, cwd, environment, log_path): commands.append(command) task_path = Path(command[command.index("--task-config") + 1]) reward_path = Path(command[command.index("--reward-config") + 1]) from scripts.train import _load_reward_config, _load_task_config self.assertEqual(_load_task_config(OBSTACLE_TASK, str(task_path), 42), custom) self.assertEqual( _load_reward_config(str(reward_path), None, OBSTACLE_TASK), self.trial["rewardConfig"], ) if "scripts/train.py" in command: (log_path.parent / "model_0.pt").write_bytes(b"mock-checkpoint") (log_path.parent / "policy.onnx").write_bytes(b"mock-onnx") else: Path(command[command.index("--output") + 1]).write_text( json.dumps(evaluation(custom)) ) return 0 self.manager.cancel_events[self.session["id"]] = threading.Event() with ( patch.object(self.manager, "_run_command", side_effect=run), patch.object(self.manager.studies, "record", return_value=0), ): result = self.manager._execute_trial(self.session, self.trial) self.assertEqual(result["state"], "completed") self.assertIn("--steps-per-seed=1000", commands[1]) self.assertIn(OBSTACLE_TASK, commands[0]) context = self.manager._proposal_context(self.session) self.assertIn("32", context["taskContext"]) self.assertEqual(set(context["allowlist"]["params"]), {"target_velocity"}) def test_cas_stale_guardrails_and_mode_roundtrip(self): sid = self.session["id"] constraint = {"params.target_velocity": {"kind": "range", "min": 0.4, "max": 0.8}} self.manager.set_constraints(sid, {"revision": 0, "constraints": constraint}) before = self.manager.detail(sid) with self.assertRaises(ResourceBusyError): self.manager.set_constraints(sid, {"revision": 0, "constraints": {}}) self.assertEqual(self.manager.detail(sid), before) for bad in ( {"weights.pose": {"kind": "fixed", "value": 1}}, {"params.target_velocity": {"kind": "fixed", "value": float("nan")}}, ): with self.assertRaises(RewardConfigError): self.manager.set_constraints(sid, {"revision": 1, "constraints": bad}) self.assertEqual(self.manager.detail(sid), before) self.manager.storage.update_session(sid, state="awaiting_approval") proposal = self.manager.storage.create_proposal( sid, self.trial["id"], {"params": {"target_velocity": 0.7}}, "test", {}, 0.8, "agent" ) self.assertEqual(self.manager.set_mode(sid, {"mode": "automatic"})["mode"], "automatic") self.assertEqual(self.manager.storage.get_proposal(proposal["id"])["state"], "approved") self.assertEqual(self.manager.set_mode(sid, {"mode": "approval"})["mode"], "approval") with self.assertRaises(RewardConfigError): validate_proposal( {"params": {"target_velocity": 0.9}}, self.trial["rewardConfig"], constraint, OBSTACLE_TASK, ) def test_protocol_fields_cannot_be_changed(self): for values in ( {"evalSteps": 999}, {"seeds": [1, 2, 3]}, {"objectiveWeights": {}}, {"taskConfig": {"unknown": 1}}, {"taskConfig": {"seed": True}}, {"seed": 1, "taskConfig": {"seed": True}}, ): with self.assertRaises(TuningError): self.manager.parse_create({"taskId": OBSTACLE_TASK, **values}) class IsolatedEvaluationTest(unittest.TestCase): def test_spawn_results_and_failure_cancel_missing_seed_fail_closed(self): import subprocess from scripts.evaluate import EvaluateConfig from scripts.evaluate_obstacle import evaluate_isolated_seeds custom = validate_task_config(OBSTACLE_TASK, {}, 42) fixed = scoring.protocol(custom, 2) reward = base_configuration(OBSTACLE_TASK) with tempfile.TemporaryDirectory() as directory: checkpoint = Path(directory) / "model.pt" checkpoint.write_bytes(b"test checkpoint snapshot") cfg = EvaluateConfig(checkpoint=str(checkpoint), output="unused", num_envs=2) def run(command, **kwargs): self.assertFalse(kwargs["shell"]) self.assertEqual(command[0], sys.executable) request = json.loads(Path(command[2]).read_text()) result = evaluation(custom)["seedMetrics"][0] result["seed"] = request["scenario"]["seed"] result["checkpointSha256"] = request["checkpointSha256"] Path(command[3]).write_text(json.dumps(result)) with patch("scripts.evaluate_obstacle.subprocess.run", side_effect=run) as mocked: results = evaluate_isolated_seeds(OBSTACLE_TASK, cfg, fixed, reward) self.assertEqual([r["seed"] for r in results], [101, 202, 303]) self.assertEqual(mocked.call_count, 3) for failure in ( subprocess.CalledProcessError(1, "worker"), subprocess.TimeoutExpired("worker", 1200), KeyboardInterrupt(), ): with patch( "scripts.evaluate_obstacle.subprocess.run", side_effect=failure ) as mocked: with self.assertRaises(type(failure)): evaluate_isolated_seeds(OBSTACLE_TASK, cfg, fixed, reward) self.assertEqual(mocked.call_count, 1) with ( patch("scripts.evaluate_obstacle.subprocess.run"), self.assertRaises(FileNotFoundError), ): evaluate_isolated_seeds(OBSTACLE_TASK, cfg, fixed, reward) def wrong_seed(command, **kwargs): run(command, **kwargs) path = Path(command[3]) result = json.loads(path.read_text()) result["seed"] = 0 path.write_text(json.dumps(result)) with ( patch("scripts.evaluate_obstacle.subprocess.run", side_effect=wrong_seed), self.assertRaises(EvaluationError), ): evaluate_isolated_seeds(OBSTACLE_TASK, cfg, fixed, reward) @unittest.skipUnless( os.environ.get("GO2_RUN_TUNING_SMOKE") == "1", "opt-in real GPU checkpoint rollout" ) class ObstacleRolloutSmoke(unittest.TestCase): def test_actual_checkpoint_policy_statistics_and_first_terminal(self): from dataclasses import asdict import src.tasks # noqa: F401 import torch from mjlab.envs import ManagerBasedRlEnv from mjlab.rl import RslRlVecEnvWrapper from mjlab.tasks.registry import load_rl_cfg, load_runner_cls from scripts.evaluate import EvaluateConfig from scripts.evaluate_obstacle import FirstEpisodeRecorder, configure_seed cfg = EvaluateConfig( checkpoint=os.environ["GO2_TUNING_CHECKPOINT"], output="/tmp/unused.json", num_envs=2, device="cuda:0", ) custom = validate_task_config(OBSTACLE_TASK, {}, 42) scenario = scoring.evaluation_scenarios(custom)[0] env_cfg = configure_seed(OBSTACLE_TASK, cfg, scenario, base_configuration(OBSTACLE_TASK)) env_cfg.episode_length_s = 0.02 # Test-only one-step timeout exercises auto-reset capture. env = ManagerBasedRlEnv(env_cfg, device="cuda:0") agent_cfg = load_rl_cfg(OBSTACLE_TASK) wrapped = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) try: runner = load_runner_cls(OBSTACLE_TASK)( wrapped, asdict(agent_cfg), log_dir=None, device="cuda:0" ) runner.load( cfg.checkpoint, load_cfg={"actor": True}, strict=True, map_location="cuda:0" ) saved = torch.load(cfg.checkpoint, weights_only=False)["actor_state_dict"] actual = runner.alg.actor.state_dict() normalizers = [k for k in saved if "normaliz" in k] self.assertTrue(normalizers) for key in normalizers: torch.testing.assert_close(actual[key], saved[key].to(actual[key].device)) obs, _ = env.reset(seed=101) recorder = FirstEpisodeRecorder(env, scenario["terrain"]) with torch.inference_mode(): policy = runner.get_inference_policy(device="cuda:0") for _ in range(2): obs, _, _, _ = wrapped.step(policy(obs)) self.assertEqual([len(s) for s in recorder.samples], [1, 1]) self.assertTrue(all(s[0]["terminal"] for s in recorder.samples)) scoring.validate_metrics(recorder.metrics(horizon=2)) finally: wrapped.close()