import contextlib import io import json import tempfile import unittest from pathlib import Path from unittest.mock import patch import torch from ml.eval.eval_cl import json_finite, run from ml.eval.diag_collision import replay_certificates, replay_dense, replay_first_hits from ml.sim.closed_loop import TeacherOracle, rollout from ml.common.kin_torch import Kin, seg_seg_dist class EvalMetricsTest(unittest.TestCase): def setUp(self): self.tasks = [ {'q0': [0] * 6, 'T': [0] * 3, 'exec_s': 1, 'travel': 5}, {'q0': [0] * 6, 'T': [0] * 3, 'exec_s': None, 'travel': None}, {'q0': [0] * 6, 'T': [0] * 3, 'exec_s': 2, 'travel': 6}, ] self.per = {k: torch.tensor(v) for k, v in { 'ok': [True, False, True], 'done': [True, False, True], 't_done': [0.5, float('nan'), 2.0], 'col': [1.0, 0.0, 0.0], 'travel': [6.0, 0.0, 8.0], 'brakes': [2.0, 4.0, 6.0], 'start_fail': [False, True, False], 'goal_ratio_2s': [0.3, 0.9, 0.1], 'err': [0.2, 200.0, 0.8], 'dls_it': [4, -1, 20], }.items()} def evaluate(self, tasks, bs=2, **kw): offset = 0 def fake_rollout(model, kin, batch, device, **kwargs): nonlocal offset p = {k: v[offset:offset + len(batch)] for k, v in self.per.items()} offset += len(batch) return {'time_budget': 'mixed'}, p with patch('ml.eval.eval_cl.rollout', side_effect=fake_rollout), contextlib.redirect_stdout(io.StringIO()): return run(None, None, tasks, 'cpu', bs=bs, **kw)[0] def test_global_denominators_and_collision_free_arrival(self): result = self.evaluate(self.tasks) self.assertEqual(result['n_ref'], 2) self.assertAlmostEqual(result['success'], 2 / 3, places=6) self.assertAlmostEqual(result['collision_free_success'], 1 / 3, places=6) self.assertAlmostEqual(result['collision_rate'], 1 / 3, places=6) self.assertEqual(result['time_ok_1p3'], 1.0) self.assertEqual(result['time_mean_s'], 1.25) self.assertAlmostEqual(result['travel_ratio_median'], 1.2, places=6) self.assertAlmostEqual(result['goal_ratio_2s'], 0.3, places=6) self.assertEqual(result['err_fail_med'], 200.0) self.assertEqual(result['brakes_mean'], 4.0) def test_empty_conditional_metrics_and_explicit_limit(self): task = {**self.tasks[0], 'exec_s': None, 'travel': None} result = self.evaluate([task], max_sec=0.04) self.assertIsNone(result['err_fail_med']) self.assertIsNone(result['time_ok_1p3']) self.assertIsNone(result['travel_ratio_median']) self.assertEqual(result['time_budget'], 'fixed 0.040 s') self.assertEqual(result['collision_free_success'], 0.0) def test_json_replaces_nonfinite_without_changing_zero(self): value = {'items': [float('nan'), float('inf'), -float('inf'), 0.0]} text = json.dumps(json_finite(value), allow_nan=False) self.assertEqual(json.loads(text), {'items': [None, None, None, 0.0]}) def test_first_hit_probe_detects_collision(self): class PlaneKin: def __init__(self, **kwargs): self.cont = torch.zeros(6, dtype=torch.bool) self.lo, self.hi = torch.full((6,), -10.0), torch.full((6,), 10.0) def gaps(self, q, cap): return {'gm': q[:, :1].expand(-1, 49).clamp(max=cap)} traj = torch.zeros(1, 3, 6) traj[0, :, 0] = torch.tensor([1.0, 1.0, -1.0]) per = {'col': torch.tensor([1.0]), 'first_hit': torch.tensor([0.04]), 'dls_it': torch.tensor([-1]), 'traj': traj} with tempfile.TemporaryDirectory() as folder: trace = Path(folder) / 'trace.pt' torch.save({'tasks': [self.tasks[0]], 'batches': [({}, per, 1)]}, trace) with patch('ml.eval.diag_collision.Kin', PlaneKin), contextlib.redirect_stdout(io.StringIO()): result = replay_first_hits(trace) self.assertEqual(result['reported_collision_tasks'], 1) self.assertEqual(result['replayed_first_hits32'], 1) self.assertEqual(result['replayed_first_hits64'], 1) self.assertEqual(result['rows'][0]['gap64_mm'], [0.5, 0.5, 0.0, -1.0]) self.assertEqual(result['rows'][0]['dense_gap64_min_mm'], -1.0) def test_dense_probe_detects_between_checkpoints_and_respects_deadline(self): class ThinObstacle: fp = 'probe' def __init__(self, **kwargs): pass def gaps(self, q, cap): return {'gm': ((q[:, :1] - 0.25).abs() - 0.05).clamp(max=cap)} traj = torch.zeros(4, 3, 6) traj[:, :, 0] = torch.tensor([[0, 1, 1], [0, -0.5, -1], [0, 0, 0.25], [0, 0, 0.25]]) per = {'traj': traj, 'col': torch.zeros(4), 'limit': torch.tensor([0.04, 0.04, 0.02, 0.04]), 't_done': torch.tensor([float('nan')] * 3 + [0.02])} with tempfile.TemporaryDirectory() as folder: trace = Path(folder) / 'trace.pt' output = Path(folder) / 'audit.json' torch.save({'kin_fp': 'probe', 'tasks': [{}] * 4, 'batches': [({}, per, 4)]}, trace) with patch('ml.eval.diag_collision.Kin', ThinObstacle), contextlib.redirect_stdout(io.StringIO()): result = replay_dense(trace, [0, 1, 2, 3], 0.01, out=output) with self.assertRaisesRegex(ValueError, 'valid task indices'): replay_dense(trace, [4]) with self.assertRaisesRegex(ValueError, 'positive max_step_deg'): replay_dense(trace, [0], 0.0) self.assertEqual(json.loads(output.read_text())['collision_tasks'], 1) self.assertEqual(result['collision_tasks'], 1) self.assertEqual(result['rows'][0]['collision_segments'], 1) self.assertAlmostEqual(result['rows'][0]['min_gap64_mm'], -0.05) self.assertGreater(result['rows'][1]['min_gap64_mm'], 0) for row in result['rows'][2:]: self.assertGreater(row['min_gap64_mm'], 0) self.assertEqual(row['steps'], 1) def test_certificate_replay_samples_actual_movement_before_deadlines(self): kin = SegmentCertificateTest.ThinObstacle() kin.fp = 'probe' traj = torch.zeros(3, 4, 6) traj[:, :, 0] = torch.tensor([[0, 1, 1, 1], [0, -0.5, -1, -1], [0, 0, 1, 1]]) per = {'traj': traj, 'limit': torch.tensor([0.06, 0.02, 0.02]), 't_done': torch.tensor([float('nan'), 0.02, float('nan')]), 'plans': {'it': torch.arange(3), 'commit': torch.zeros(3, 3, 1, 6), 'certified': torch.ones(3, 3, 1, dtype=torch.bool), 'gaps': torch.full((3, 3, 2), 0.5)}} with tempfile.TemporaryDirectory() as folder: trace, output = Path(folder) / 'trace.pt', Path(folder) / 'audit.json' torch.save({'kin_fp': 'probe', 'tasks': [{}] * 3, 'batches': [({}, per, 3)]}, trace) with patch('ml.eval.diag_collision.Kin', return_value=kin), contextlib.redirect_stdout(io.StringIO()): result = replay_certificates(trace, [0, 1, 2], out=output, repeats=1, extra_segments=[(0, 0)]) with self.assertRaisesRegex(ValueError, 'extra segment'): replay_certificates(trace, [1], repeats=1, extra_segments=[(1, 1)]) self.assertEqual(json.loads(output.read_text())['dense_collision_segments'], 1) self.assertEqual([(r['i'], r['start_it']) for r in result['rows']], [(0, 0), (1, 0)]) self.assertEqual(result['budgets']['32']['safe'], 1) self.assertEqual(result['budgets']['32']['collision'], 1) self.assertEqual(result['accepted_dense_collision_segments'], 0) self.assertTrue(result['batch_single_match']) self.assertEqual(result['timing']['calls'], 2) class DlsSafetyTest(unittest.TestCase): def test_rejected_refinement_retains_safe_candidate(self): class ObstacleKin: def clamp_box(self, q): return q def gaps(self, q, cap): return {'gm': torch.where(q[:, :1] > 0.3, -torch.ones_like(q[:, :1]), torch.ones_like(q[:, :1]))} task = {'q0': [0] * 6, 'T': [2, 0, 0]} traj = torch.zeros(1, 2, 6) traj[0, 1, 0] = 0.1 ref = torch.zeros(1, 1, 6) ref[0, 0, 0] = 1.0 patches = { 'features': lambda *args: torch.zeros(1, 83), 'goal_err': lambda kin, q, T, R: (torch.full((q.shape[0],), 2.0), torch.zeros(q.shape[0])), 'dls_steps': lambda *args, **kwargs: ref, } results = [] with patch.multiple('ml.sim.closed_loop', **patches): for gate in (0, 1): _, per = rollout(None, ObstacleKin(), [task], 'cpu', K=1, safety=True, max_sec=0.02, return_traj=True, use_sub=False, select='first', oracle=TeacherOracle(traj), replan='empty', commit_n=1, safety_partial=1, dls_safety=gate) results.append(per) self.assertEqual(results[0]['col'].item(), 1.0) self.assertEqual(results[0]['dls_it'].item(), 0) self.assertEqual(results[1]['col'].item(), 0.0) self.assertEqual(results[1]['dls_it'].item(), -1) self.assertAlmostEqual(results[1]['traj'][0, 1, 0].item(), 0.1) def test_feedback_is_reduced_only_after_full_rejection(self): class ObstacleKin: def clamp_box(self, q): return q def gaps(self, q, cap): return {'gm': torch.where(q[:, :1] > 0.3, -torch.ones_like(q[:, :1]), torch.ones_like(q[:, :1]))} task = {'q0': [0] * 6, 'T': [10, 0, 0]} traj = torch.zeros(1, 17, 6) traj[0, :, 0] = -0.2 * torch.arange(17) ref = torch.zeros(1, 1, 6) ref[0, 0, 0] = 1.0 patches = { 'features': lambda *args: torch.zeros(1, 83), 'goal_err': lambda kin, q, T, R: (torch.full((q.shape[0],), 10.0), torch.zeros(q.shape[0])), 'dls_steps': lambda *args, **kwargs: ref, } results = [] with patch.multiple('ml.sim.closed_loop', **patches): for gain in (1.0, 0.5): _, per = rollout(None, ObstacleKin(), [task], 'cpu', K=1, safety=True, max_sec=0.04, return_traj=True, use_sub=False, select='first', oracle=TeacherOracle(traj), replan='empty', commit_n=1, safety_partial=1, fb_gain=0.6, fb_block_gain=gain) results.append(per) self.assertEqual(results[0]['brakes'].item(), 2.0) self.assertEqual(results[1]['brakes'].item(), 1.0) self.assertEqual(results[1]['traj'][0, 1, 0].item(), 0.0) self.assertAlmostEqual(results[1]['traj'][0, 2, 0].item(), 0.1) self.assertEqual(results[1]['col'].item(), 0.0) class ClampedPathSafetyTest(unittest.TestCase): def rollout_candidate(self, increments, safety_partial, goal_error=10.0, **kwargs): class ObstacleKin: def clamp_box(self, q): clamped = q.clone() clamped[..., 4] = clamped[..., 4].clamp(max=0.0) return clamped def gaps(self, q, cap): collision = (q[:, :1] > 0.1) & (q[:, 4:5] < 0.2) return {'gm': torch.where(collision, -torch.ones_like(q[:, :1]), torch.ones_like(q[:, :1]))} task = {'q0': [0] * 6, 'T': [goal_error, 0, 0]} commit_n = len(increments) traj = torch.zeros(1, commit_n + 1, 6) traj[0, 1:] = torch.tensor(increments).cumsum(0) patches = { 'features': lambda *args: torch.zeros(1, 83), 'goal_err': lambda kin, q, T, R: (torch.full((q.shape[0],), goal_error), torch.zeros(q.shape[0])), } with patch.multiple('ml.sim.closed_loop', **patches): _, per = rollout(None, ObstacleKin(), [task], 'cpu', K=1, safety=True, max_sec=0.02 * commit_n, return_traj=True, use_sub=False, select='first', oracle=TeacherOracle(traj), replan='empty', commit_n=commit_n, safety_partial=safety_partial, **kwargs) return per def test_full_safety_rejects_collision_after_joint_clamp(self): # Raw endpoint/midpoint are safe, but the J5 upper limit puts the endpoint at (0.2, 0). per = self.rollout_candidate([[0.2, 0, 0, 0, 0.6, 0]], safety_partial=0) self.assertEqual(per['col'].item(), 0.0) torch.testing.assert_close(per['traj'], torch.zeros(1, 2, 6), rtol=0, atol=0) def test_partial_safety_rejects_collision_after_joint_clamp(self): per = self.rollout_candidate([[0.2, 0, 0, 0, 0.6, 0]], safety_partial=1) self.assertEqual(per['col'].item(), 0.0) torch.testing.assert_close(per['traj'], torch.zeros(1, 2, 6), rtol=0, atol=0) def test_partial_safety_keeps_only_clamped_safe_prefix(self): per = self.rollout_candidate([ [0.05, 0, 0, 0, 0.6, 0], [0.2, 0, 0, 0, 0.6, 0], ], safety_partial=1) self.assertAlmostEqual(per['traj'][0, -1, 0].item(), 0.05) self.assertEqual(per['traj'][0, -1, 4].item(), 0.0) self.assertEqual(per['col'].item(), 0.0) expected = torch.zeros(1, 3, 6) expected[0, 1:, 0] = 0.05 torch.testing.assert_close(per['traj'], expected) torch.testing.assert_close(per['plans']['commit'][0, 0, 1], torch.zeros(6), rtol=0, atol=0) def test_near_dls_rejects_clamped_collision_and_keeps_candidate(self): candidate = [[0.05, 0, 0, 0, 0.6, 0]] ref = torch.tensor([[[0.2, 0, 0, 0, 0.6, 0]]]) for safety_partial in (0, 1): with self.subTest(safety_partial=safety_partial): with patch('ml.sim.closed_loop.dls_steps', return_value=ref) as dls: per = self.rollout_candidate(candidate, safety_partial, goal_error=2.0, dls_safety=1) dls.assert_called_once() self.assertEqual(per['col'].item(), 0.0) self.assertEqual(per['dls_it'].item(), -1) self.assertEqual(per['brakes'].item(), 0.0) expected = torch.zeros(1, 2, 6) expected[0, 1, 0] = 0.05 torch.testing.assert_close(per['traj'], expected) torch.testing.assert_close(per['plans']['commit'], torch.tensor(candidate)[None, None]) self.assertTrue(per['plans']['certified'].all().item()) def test_r3_clamped_queue_origin_prefix_and_certificates_match_execution(self): class SafeKin: def clamp_box(self, q): clamped = q.clone() clamped[..., 4] = clamped[..., 4].clamp(max=0.0) return clamped def gaps(self, q, cap): return {'gm': torch.ones_like(q[:, :1])} prefixes = [] def capture_features(kin, q, qp, T, R, pre, r, sub, sf): prefixes.append((pre.clone(), int(r.item()))) return torch.zeros(1, 83) teacher = torch.zeros(1, 25, 6) teacher[0, :, 0] = 0.02 * torch.arange(25) teacher[0, :, 4] = 0.6 * torch.arange(25) task = {'q0': [0] * 6, 'T': [10, 0, 0]} patches = { 'features': capture_features, 'goal_err': lambda kin, q, T, R: (torch.full((q.shape[0],), 10.0), torch.zeros(q.shape[0])), } with patch.multiple('ml.sim.closed_loop', **patches): _, per = rollout(None, SafeKin(), [task], 'cpu', K=1, safety=True, max_sec=0.16, return_traj=True, use_sub=False, select='first', oracle=TeacherOracle(teacher), replan='r3', commit_n=3, safety_partial=1) self.assertEqual([r for _, r in prefixes], [0, 2, 3, 3]) for pre, r in prefixes[1:]: torch.testing.assert_close(pre[..., 4], torch.zeros(1, 3), rtol=0, atol=0) torch.testing.assert_close(pre[:, :r, 0], torch.full((1, r), 0.02)) self.assertEqual(per['col'].item(), 0.0) expected = torch.zeros(1, 9, 6) expected[0, :, 0] = 0.02 * torch.arange(9) torch.testing.assert_close(per['traj'], expected) plans = per['plans'] self.assertEqual(plans['it'].tolist(), [0, 3, 6, 9]) torch.testing.assert_close(plans['commit'][..., 4], torch.full((1, 4, 3), 0.6)) self.assertTrue(plans['certified'].all().item()) self.assertTrue((plans['gaps'] > 0).all().item()) checked = 0 for i, start in enumerate(plans['it'].tolist()): n = min(3, 8 - start) if n <= 0: continue torch.testing.assert_close(plans['qe'][:, i], per['traj'][:, start], rtol=0, atol=0) before = per['traj'][:, start:start + n] after = per['traj'][:, start + 1:start + n + 1] torch.testing.assert_close(plans['points'][:, i, :n], after, rtol=0, atol=0) torch.testing.assert_close(plans['midpoints'][:, i, :n], before + 0.5 * (after - before), rtol=0, atol=0) checked += n self.assertEqual(checked, 8) def rollout_fallback(self, safety_partial, midpoint_collision=False): class ObstacleKin: def clamp_box(self, q): return q def gaps(self, q, cap): collision = q[:, :1] > 0.3 if midpoint_collision: collision |= (q[:, :1] > 0.09) & (q[:, :1] < 0.11) return {'gm': torch.where(collision, -torch.ones_like(q[:, :1]), torch.ones_like(q[:, :1]))} teacher = torch.zeros(1, 4, 6) teacher[0, :, 0] = torch.arange(4) ref = torch.zeros(1, 1, 6) ref[0, 0, 0] = 0.2 task = {'q0': [0] * 6, 'T': [10, 0, 0]} patches = { 'features': lambda *args: torch.zeros(1, 83), 'goal_err': lambda kin, q, T, R: (torch.full((q.shape[0],), 10.0), torch.zeros(q.shape[0])), 'dls_steps': lambda *args, **kwargs: ref, } with patch.multiple('ml.sim.closed_loop', **patches): _, per = rollout(None, ObstacleKin(), [task], 'cpu', K=1, safety=True, max_sec=0.06, return_traj=True, use_sub=False, select='first', oracle=TeacherOracle(teacher), replan='empty', commit_n=3, safety_partial=safety_partial, fb_gain=1.0, fb_fallback=1.0) return per def test_fallback_executes_one_safe_dls_step_then_certified_zeros(self): for safety_partial in (0, 1): with self.subTest(safety_partial=safety_partial): per = self.rollout_fallback(safety_partial) self.assertEqual(per['col'].item(), 0.0) self.assertEqual(per['brakes'].item(), 1.0) expected = torch.zeros(1, 4, 6) expected[0, 1:, 0] = 0.2 torch.testing.assert_close(per['traj'], expected) commit = torch.zeros(1, 1, 3, 6) commit[0, 0, 0, 0] = 0.2 plans = per['plans'] self.assertEqual(plans['it'].tolist(), [0]) torch.testing.assert_close(plans['commit'], commit, rtol=0, atol=0) self.assertTrue(plans['certified'].all().item()) self.assertTrue((plans['gaps'] > 0).all().item()) torch.testing.assert_close(plans['points'][:, 0], expected[:, 1:]) torch.testing.assert_close(plans['midpoints'][:, 0], expected[:, :-1] + 0.5 * (expected[:, 1:] - expected[:, :-1])) def test_fallback_rejects_safe_endpoint_with_colliding_midpoint(self): for safety_partial in (0, 1): with self.subTest(safety_partial=safety_partial): per = self.rollout_fallback(safety_partial, midpoint_collision=True) self.assertEqual(per['col'].item(), 0.0) self.assertEqual(per['brakes'].item(), 1.0) torch.testing.assert_close(per['traj'], torch.zeros(1, 4, 6), rtol=0, atol=0) torch.testing.assert_close(per['plans']['commit'], torch.zeros(1, 1, 3, 6), rtol=0, atol=0) self.assertFalse(per['plans']['certified'].any().item()) class SegmentCertificateTest(unittest.TestCase): class ThinObstacle(Kin): def __init__(self): self.dtype = torch.float64 self.lo, self.hi = torch.full((6,), -5.0), torch.full((6,), 5.0) self.cont = torch.zeros(6, dtype=torch.bool) self.cert_groups = torch.tensor([0]) self.cert_lever = torch.tensor([[1.0, 0, 0, 0, 0, 0]], dtype=self.dtype) def gaps(self, q, cap): return {'gm': ((q[:, :1] - 0.25).abs() - 0.05).clamp(max=cap)} def test_bound_detects_hidden_collision_and_accepts_clear_segment(self): a, b = torch.zeros(3, 6, dtype=torch.float64), torch.zeros(3, 6, dtype=torch.float64) a[:, 0] = torch.tensor([0.0, -1.0, 0.25]) b[:, 0] = torch.tensor([1.0, 0.0, 0.25]) result = self.ThinObstacle().certify_segments(a, b) self.assertEqual(result['safe'].tolist(), [False, True, False]) self.assertEqual(result['collision'].tolist(), [True, False, True]) self.assertEqual(result['queries'].tolist(), [4, 2, 2]) def test_budget_and_depth_exhaustion_do_not_accept_unknown_segment(self): a, b = torch.zeros(1, 6, dtype=torch.float64), torch.zeros(1, 6, dtype=torch.float64) b[:, 0] = 1.0 kin = self.ThinObstacle() for kwargs in ({'max_queries': 2}, {'max_depth': 0}): with self.subTest(**kwargs): result = kin.certify_segments(a, b, **kwargs) self.assertFalse(result['safe'].item()) self.assertTrue(result['unresolved'].item()) self.assertEqual(result['queries'].item(), 2) def test_budget_exhaustion_scales_with_gap_and_bound(self): class ConstantGap(self.ThinObstacle): def gaps(self, q, cap): return {'gm': torch.full((q.shape[0], 1), 0.1, dtype=q.dtype, device=q.device)} # 界取 1 mm/度：1° 段要证明无碰撞，必须细分到位移界 < 两端间隙之和(0.2 mm)。 kin = ConstantGap() kin.cert_lever = torch.ones(1, 6, dtype=torch.float64) a, b = torch.zeros(1, 6, dtype=torch.float64), torch.zeros(1, 6, dtype=torch.float64) b[:, 0] = 1.0 for budget, safe, queries in ((4, False, 3), (8, False, 5), (9, True, 9)): with self.subTest(budget=budget): result = kin.certify_segments(a, b, max_queries=budget) self.assertEqual(bool(result['safe'].item()), safe) self.assertEqual(int(result['queries'].item()), queries) def test_tensor_budget_cannot_wrap_depth(self): a, b = torch.zeros(1, 6, dtype=torch.float64), torch.zeros(1, 6, dtype=torch.float64) a[:, 0] = b[:, 0] = 0.25 for tensor in (torch.tensor(255, dtype=torch.uint8), torch.tensor(255)): with self.subTest(dtype=tensor.dtype): with self.assertRaisesRegex(ValueError, 'plain ints'): self.ThinObstacle().certify_segments(a, b, max_depth=tensor) def test_batch_matches_per_root_and_ignores_order(self): kin = self.ThinObstacle() a = torch.tensor([[0.0, 0, 0, 0, 0, 0], [0.25, 0, 0, 0, 0, 0], [-1.0, 0, 0, 0, 0, 0]], dtype=torch.float64) b = torch.tensor([[1.0, 0, 0, 0, 0, 0], [0.25, 0, 0, 0, 0, 0], [0.0, 0, 0, 0, 0, 0]], dtype=torch.float64) batch = kin.certify_segments(a, b, max_queries=16) self.assertEqual(batch['collision'].tolist(), [True, True, False]) self.assertEqual(batch['safe'].tolist(), [False, False, True]) for perm in ([0, 1, 2], [2, 1, 0], [1, 0, 2]): index = torch.tensor(perm) moved = kin.certify_segments(a[index], b[index], max_queries=16) for key in ('safe', 'collision', 'unresolved', 'queries'): torch.testing.assert_close(moved[key], batch[key][index], rtol=0, atol=0) def test_empty_batch_and_all_rejected_frontier(self): kin = self.ThinObstacle() empty = kin.certify_segments(torch.zeros(0, 6, dtype=torch.float64), torch.zeros(0, 6, dtype=torch.float64)) self.assertEqual(empty['safe'].shape, (0,)) a = torch.tensor([[0.25, 0, 0, 0, 0, 0], [6.0, 0, 0, 0, 0, 0]], dtype=torch.float64) result = kin.certify_segments(a, a.clone(), max_queries=2) self.assertEqual(result['safe'].tolist(), [False, False]) self.assertEqual(result['collision'].tolist(), [True, False]) self.assertEqual(result['outside_box'].tolist(), [False, True]) def test_nonfinite_and_outside_box_are_rejected(self): a, b = torch.zeros(1, 6, dtype=torch.float64), torch.zeros(1, 6, dtype=torch.float64) kin = self.ThinObstacle() with patch.object(kin, 'gaps', side_effect=lambda q, cap: {'gm': torch.full((q.shape[0], 1), float('nan'))}): result = kin.certify_segments(a, b) self.assertFalse(result['safe'].item()) self.assertTrue(result['unresolved'].item()) for value in (float('inf'), -float('inf')): with self.subTest(value=value): with patch.object(kin, 'gaps', side_effect=lambda q, cap, v=value: {'gm': torch.full((q.shape[0], 1), v)}): result = kin.certify_segments(a, b) self.assertFalse(result['safe'].item()) a[:, 0] = b[:, 0] = 6.0 result = kin.certify_segments(a, b) self.assertFalse(result['safe'].item()) self.assertTrue(result['outside_box'].item()) def test_point_segment_distance_is_symmetric_and_differentiable(self): p = torch.tensor([[0.0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]], dtype=torch.float64, requires_grad=True) q = torch.tensor([[2.0, 0, 0]] * 4, dtype=torch.float64) point = torch.tensor([[1.0, 0, 0], [1, 3, 0], [4, 0, 0], [-1, 0, 0]], dtype=torch.float64) forward = seg_seg_dist(p, q, point, point) reverse = seg_seg_dist(point, point, p, q) torch.testing.assert_close(forward, torch.tensor([0.0, 3, 2, 1], dtype=torch.float64), atol=1e-8, rtol=0) torch.testing.assert_close(forward, reverse, atol=1e-8, rtol=0) forward.sum().backward() self.assertTrue(torch.isfinite(p.grad).all()) torch.testing.assert_close(seg_seg_dist(p.detach(), p.detach(), point, point), (p.detach() - point).norm(dim=-1), atol=1e-8, rtol=0) if __name__ == '__main__': unittest.main()