qing
45e019952e
- config.py: action_dim默认值 8→18（对齐bot4cj 18列） - mt_parser.py: 轻量级.mt @Action解析器（支持-占位符前向填充） - pretrain_mt.py: 预训练入口（零图像占位，checkpoint兼容现有推理） - _run.sh: 新增pretrain子命令 - README.md: 新增预训练文档 + action_dim说明 - test_actpred.py: 测试断言更新为18
220 lines
8.3 KiB
Python
220 lines
8.3 KiB
Python
"""actpred 模块单元测试
|
|
|
|
运行方式：
|
|
cd actpred && python -m pytest test_actpred.py -v
|
|
# 或
|
|
cd actpred && python test_actpred.py
|
|
"""
|
|
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
|
|
from config import Config, DataConfig, ModelConfig, build_argparser, apply_overrides
|
|
from dataset import _natural_key, ActionDataset, DataRecorder
|
|
from model import ActionPredictor
|
|
from predict import parse_action_history
|
|
|
|
|
|
class TestNaturalKey(unittest.TestCase):
|
|
"""测试自然排序键函数"""
|
|
|
|
def test_natural_order_no_padding(self):
|
|
"""无前导零的文件名应按数字序排列"""
|
|
files = [Path(f"{i}.jpg") for i in [3, 1, 10, 2, 20]]
|
|
sorted_files = sorted(files, key=_natural_key)
|
|
names = [f.name for f in sorted_files]
|
|
self.assertEqual(names, ["1.jpg", "2.jpg", "3.jpg", "10.jpg", "20.jpg"])
|
|
|
|
def test_natural_order_with_padding(self):
|
|
"""有前导零的文件名排序也应正确"""
|
|
files = [Path(f"{i:04d}.jpg") for i in [3, 1, 10, 2, 20]]
|
|
sorted_files = sorted(files, key=_natural_key)
|
|
names = [f.name for f in sorted_files]
|
|
self.assertEqual(names, ["0001.jpg", "0002.jpg", "0003.jpg", "0010.jpg", "0020.jpg"])
|
|
|
|
|
|
class TestParseActionHistory(unittest.TestCase):
|
|
"""测试动作历史解析函数"""
|
|
|
|
def test_single_step(self):
|
|
history = parse_action_history("1,2,3,4", action_dim=4)
|
|
self.assertEqual(history.shape, (1, 4))
|
|
np.testing.assert_array_equal(history[0], [1, 2, 3, 4])
|
|
|
|
def test_multi_step(self):
|
|
history = parse_action_history("1,0,0,0:0,1,0,0:0,0,1,0", action_dim=4)
|
|
self.assertEqual(history.shape, (3, 4))
|
|
|
|
def test_wrong_dim_raises(self):
|
|
with self.assertRaises(AssertionError):
|
|
parse_action_history("1,2,3", action_dim=4)
|
|
|
|
|
|
class TestActionPredictor(unittest.TestCase):
|
|
"""测试模型前向传播"""
|
|
|
|
def test_forward_simple_mlp(self):
|
|
data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5)
|
|
model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp")
|
|
model = ActionPredictor(model_cfg, data_cfg)
|
|
model.eval()
|
|
|
|
image = torch.randn(2, 3, 64, 64)
|
|
history = torch.randn(2, 5, 18)
|
|
with torch.no_grad():
|
|
out = model(image, history)
|
|
self.assertEqual(out.shape, (2, 18))
|
|
|
|
def test_forward_simple_gru(self):
|
|
data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5)
|
|
model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="gru")
|
|
model = ActionPredictor(model_cfg, data_cfg)
|
|
model.eval()
|
|
|
|
image = torch.randn(2, 3, 64, 64)
|
|
history = torch.randn(2, 5, 18)
|
|
with torch.no_grad():
|
|
out = model(image, history)
|
|
self.assertEqual(out.shape, (2, 18))
|
|
|
|
def test_forward_dynamic_batch(self):
|
|
data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3)
|
|
model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp")
|
|
model = ActionPredictor(model_cfg, data_cfg)
|
|
model.eval()
|
|
|
|
for batch_size in [1, 4, 8]:
|
|
image = torch.randn(batch_size, 3, 32, 32)
|
|
history = torch.randn(batch_size, 3, 4)
|
|
with torch.no_grad():
|
|
out = model(image, history)
|
|
self.assertEqual(out.shape, (batch_size, 4))
|
|
|
|
|
|
class TestDataRecorder(unittest.TestCase):
|
|
"""测试数据录制器"""
|
|
|
|
def test_record_and_flush(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
recorder = DataRecorder(tmpdir, action_dim=4)
|
|
for _ in range(10):
|
|
img = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)
|
|
action = np.random.randn(4).astype(np.float32)
|
|
recorder.record(img, action)
|
|
episode_path = recorder.flush()
|
|
|
|
self.assertTrue(os.path.exists(os.path.join(episode_path, "actions.npy")))
|
|
images_dir = os.path.join(episode_path, "images")
|
|
self.assertTrue(os.path.isdir(images_dir))
|
|
self.assertEqual(len(os.listdir(images_dir)), 10)
|
|
|
|
actions = np.load(os.path.join(episode_path, "actions.npy"))
|
|
self.assertEqual(actions.shape, (10, 4))
|
|
|
|
def test_multiple_episodes_increment_index(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
r1 = DataRecorder(tmpdir, action_dim=4)
|
|
r1.flush()
|
|
r2 = DataRecorder(tmpdir, action_dim=4)
|
|
r2.flush()
|
|
|
|
self.assertNotEqual(r1.episode_dir, r2.episode_dir)
|
|
self.assertIn("episode_0001", str(r2.episode_dir))
|
|
|
|
|
|
class TestActionDataset(unittest.TestCase):
|
|
"""测试数据集加载"""
|
|
|
|
@staticmethod
|
|
def _create_dummy_episode(base_dir, n_frames=10, action_dim=4, image_size=32):
|
|
ep_dir = base_dir / "episode_0001"
|
|
images_dir = ep_dir / "images"
|
|
images_dir.mkdir(parents=True, exist_ok=True)
|
|
actions = []
|
|
for i in range(n_frames):
|
|
img = Image.fromarray(
|
|
np.random.randint(0, 255, (image_size, image_size, 3), dtype=np.uint8))
|
|
img.save(images_dir / f"{i:04d}.jpg")
|
|
actions.append(np.random.randn(action_dim).astype(np.float32))
|
|
np.save(ep_dir / "actions.npy", np.array(actions))
|
|
|
|
def test_dataset_loading(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
self._create_dummy_episode(Path(tmpdir), n_frames=10, action_dim=4)
|
|
ds = ActionDataset(tmpdir, history_len=3, action_dim=4, image_size=32)
|
|
self.assertEqual(len(ds), 7) # 10 - 3 = 7 samples
|
|
|
|
sample = ds[0]
|
|
self.assertIn("image", sample)
|
|
self.assertIn("action_history", sample)
|
|
self.assertIn("target", sample)
|
|
self.assertEqual(sample["image"].shape, (3, 32, 32))
|
|
self.assertEqual(sample["action_history"].shape, (3, 4))
|
|
self.assertEqual(sample["target"].shape, (4,))
|
|
|
|
def test_episode_count(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
self._create_dummy_episode(Path(tmpdir))
|
|
count = ActionDataset.get_episode_count(tmpdir)
|
|
self.assertEqual(count, 1)
|
|
|
|
def test_empty_stats_raises(self):
|
|
"""空数据集计算归一化统计量应报错"""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with self.assertRaises(ValueError):
|
|
ActionDataset(tmpdir, history_len=5, action_dim=4,
|
|
normalize_actions=True)
|
|
|
|
def test_invalid_stats_path_raises(self):
|
|
"""指定的 stats 路径不存在应报错"""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with self.assertRaises(FileNotFoundError):
|
|
ActionDataset(tmpdir, history_len=3, action_dim=4,
|
|
normalize_actions=True,
|
|
action_stats_path="/nonexistent/path.npy")
|
|
|
|
|
|
class TestConfig(unittest.TestCase):
|
|
"""测试配置和 CLI 参数"""
|
|
|
|
def test_default_config(self):
|
|
cfg = Config()
|
|
self.assertEqual(cfg.data.image_size, 64)
|
|
self.assertEqual(cfg.data.action_dim, 18)
|
|
self.assertEqual(cfg.data.history_len, 5)
|
|
self.assertEqual(cfg.model.cnn_backbone, "simple")
|
|
self.assertEqual(cfg.model.history_encoder, "mlp")
|
|
self.assertEqual(cfg.train.epochs, 100)
|
|
|
|
def test_apply_overrides(self):
|
|
parser = build_argparser("test")
|
|
args = parser.parse_args(["--epochs", "50", "--lr", "0.01", "--backbone", "mobilenet"])
|
|
cfg = apply_overrides(Config(), args)
|
|
self.assertEqual(cfg.train.epochs, 50)
|
|
self.assertEqual(cfg.train.lr, 0.01)
|
|
self.assertEqual(cfg.model.cnn_backbone, "mobilenet")
|
|
|
|
def test_normalize_actions_flag(self):
|
|
"""--normalize-actions 和 --no-normalize-actions 都应可用"""
|
|
parser = build_argparser("test")
|
|
args = parser.parse_args(["--normalize-actions"])
|
|
self.assertTrue(args.normalize_actions)
|
|
|
|
parser = build_argparser("test")
|
|
args = parser.parse_args(["--no-normalize-actions"])
|
|
self.assertFalse(args.normalize_actions)
|
|
|
|
parser = build_argparser("test")
|
|
args = parser.parse_args([])
|
|
self.assertIsNone(args.normalize_actions)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2)
|