Robot_Arm/tests/test_eval_cl.py

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_dense, replay_first_hits
from ml.sim.closed_loop import TeacherOracle, rollout


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)


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


if __name__ == '__main__':
    unittest.main()