13

tests/test_rl_revision.py

@@ -0,0 +1,95 @@
+"""训练/部署观测和动作的一致性，以及RL改进的目标口径。"""
+import importlib.util
+import math
+from pathlib import Path
+
+import pytest
+
+np = pytest.importorskip('numpy')
+pytest.importorskip('gymnasium')
+
+from mcm26b.rl.env import LocalizationEnv
+from mcm26b.rl.features import action_mask, build_observation, candidate_actions
+
+
+def test_training_observation_matches_runtime_channel_belief():
+    env = LocalizationEnv(problem=3, handoff_probability=1)
+    obs, _ = env.reset(seed=1)
+    belief = env._belief.channels[env._channel]
+    expected = build_observation(belief.region, env._client.state.position, env._channel,
+                                 env.config, env.problem, env._blocked, belief,
+                                 virtual_time_s=env._client.state.virtual_time_s,
+                                 receiver_channel=env._client.state.channel)
+    np.testing.assert_array_equal(obs, expected)
+    assert obs[7] == np.float32(1/12)
+    assert tuple(obs[14:17]) == (0,1,0)
+    assert env._client.state.channel != env._channel
+    env.close()
+
+
+def test_mask_rejects_measured_and_duplicate_candidates():
+    env = LocalizationEnv(problem=3)
+    env.reset(seed=17)
+    belief = env._belief.channels[env._channel]
+    region = belief.region
+    position = env._client.state.position
+    candidates = candidate_actions(region, position, env.config)
+    belief.measurements.append(dict(position=candidates[0], result='no_signal', svd_deg=None))
+    mask = env.action_masks()
+    assert not mask[0]
+    keys = [(round(p[0]*2),round(p[1]*2)) for p,v in zip(candidates,mask) if v]
+    assert len(keys) == len(set(keys))
+    assert any(mask)
+    assert all(math.dist(p,region.mec()[0]) <= 1000-region.mec()[1]+1e-6
+               for p,v in zip(candidates,mask) if v)
+    env.close()
+
+
+def test_q4_blocked_sequence_resets_on_direction():
+    env = LocalizationEnv(problem=4,directional_probability=1,handoff_probability=0)
+    env.reset(seed=7)
+    for expected in ('no_signal','direction'):
+        region = env._belief.candidate_set(env._channel)
+        candidates = candidate_actions(region,env._client.state.position,env.config,env._blocked)
+        # 只有测试端可用真值选择一个指定分支，策略从不接触该接口。
+        index = next(i for i,p in enumerate(candidates)
+                     if env._simulator._observe(p,env._channel)[0] == expected)
+        _, _, terminated, truncated, _ = env.step(index)
+        assert len(env._blocked) == (1 if expected=='no_signal' else 0)
+        if expected=='no_signal':
+            assert not terminated and not truncated
+    env.close()
+
+
+def test_success_reward_is_constant_minus_actual_virtual_cost():
+    env = LocalizationEnv(problem=3,handoff_probability=0)
+    env.reset(seed=3)
+    total = 0
+    for _ in range(16):
+        legal = np.flatnonzero(env.action_masks())
+        action = int(0 if 0 in legal else legal[0])
+        _, reward, terminated, truncated, info = env.step(action)
+        total += reward
+        if terminated or truncated:
+            break
+    assert info['cleared']
+    assert total == pytest.approx(20-info['virtual_cost_s']/100)
+    assert env._client.state.exited
+
+
+def test_bc_and_ppo_forward_mapping_agree():
+    torch = pytest.importorskip('torch')
+    from stable_baselines3 import PPO
+    spec = importlib.util.spec_from_file_location('pretrain',Path('scripts/pretrain_rl_expert.py'))
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    torch.set_num_threads(1)
+    env = LocalizationEnv()
+    ppo = PPO('MlpPolicy',env,policy_kwargs={'net_arch':[128,128]},device='cpu',n_steps=16,batch_size=8)
+    bc = module.BCMlp()
+    state = bc.net.state_dict()
+    ppo.policy.mlp_extractor.policy_net.load_state_dict({k:v for k,v in state.items() if not k.startswith('4.')})
+    ppo.policy.action_net.load_state_dict({'weight':state['4.weight'],'bias':state['4.bias']})
+    x = torch.randn(10,401)
+    torch.testing.assert_close(bc(x),ppo.policy.action_net(ppo.policy.mlp_extractor.policy_net(x)))
+    env.close()