feat: switch pusht transformer logging to swanlab

tests/test_pusht_image_runner_metrics.py

@@ -0,0 +1,110 @@
+import pathlib
+import sys
+
+import gym
+from gym import spaces
+import numpy as np
+import pytest
+import torch
+
+ROOT_DIR = pathlib.Path(__file__).resolve().parents[1]
+if str(ROOT_DIR) not in sys.path:
+    sys.path.append(str(ROOT_DIR))
+
+import diffusion_policy.env_runner.pusht_image_runner as runner_module
+from diffusion_policy.env_runner.pusht_image_runner import summarize_rollout_metrics
+
+
+class FakePushTImageEnv(gym.Env):
+    metadata = {'render.modes': ['rgb_array']}
+
+    def __init__(self, legacy=False, render_size=96):
+        del legacy, render_size
+        self.observation_space = spaces.Dict({
+            'image': spaces.Box(low=0, high=255, shape=(3, 4, 4), dtype=np.uint8),
+        })
+        self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(2,), dtype=np.float32)
+        self.seed_value = 0
+        self.step_count = 0
+
+    def seed(self, seed=None):
+        self.seed_value = 0 if seed is None else seed
+
+    def reset(self):
+        self.step_count = 0
+        return {'image': np.zeros((3, 4, 4), dtype=np.uint8)}
+
+    def step(self, action):
+        del action
+        self.step_count += 1
+        reward = 0.1 if self.seed_value < 10000 else 0.9
+        done = self.step_count >= 1
+        obs = {'image': np.full((3, 4, 4), self.step_count, dtype=np.uint8)}
+        return obs, reward, done, {}
+
+    def render(self, *args, **kwargs):
+        raise AssertionError('render should not be called for scalar-only PushT image rollouts')
+
+
+class FakePolicy:
+    device = torch.device('cpu')
+    dtype = torch.float32
+
+    def reset(self):
+        return None
+
+    def predict_action(self, obs_dict):
+        n_envs = next(iter(obs_dict.values())).shape[0]
+        return {
+            'action': torch.zeros((n_envs, 2, 2), dtype=torch.float32),
+        }
+
+
+def test_summarize_rollout_metrics_keeps_scalar_rewards_renames_means_and_omits_videos():
+    log_data = summarize_rollout_metrics(
+        env_seeds=[11, 12, 101],
+        env_prefixs=['train/', 'train/', 'test/'],
+        all_rewards=[
+            [0.2, 0.8],
+            [0.1, 0.4],
+            [0.5, 0.9],
+        ],
+        all_video_paths=[
+            '/tmp/train-11.mp4',
+            '/tmp/train-12.mp4',
+            '/tmp/test-101.mp4',
+        ],
+    )
+
+    assert log_data['train/sim_max_reward_11'] == 0.8
+    assert log_data['train/sim_max_reward_12'] == 0.4
+    assert log_data['test/sim_max_reward_101'] == 0.9
+    assert log_data['train_mean_score'] == pytest.approx(0.6)
+    assert log_data['test_mean_score'] == pytest.approx(0.9)
+    assert not any(key.startswith('train/sim_video_') for key in log_data)
+    assert not any(key.startswith('test/sim_video_') for key in log_data)
+
+
+def test_runner_ignores_vis_flags_and_never_emits_sim_videos(tmp_path, monkeypatch):
+    monkeypatch.setattr(runner_module, 'PushTImageEnv', FakePushTImageEnv)
+
+    runner = runner_module.PushTImageRunner(
+        output_dir=tmp_path,
+        n_train=1,
+        n_train_vis=1,
+        n_test=1,
+        n_test_vis=1,
+        n_envs=2,
+        max_steps=2,
+        n_obs_steps=2,
+        n_action_steps=2,
+        tqdm_interval_sec=0.0,
+    )
+
+    log_data = runner.run(FakePolicy())
+
+    assert log_data['train/sim_max_reward_0'] == pytest.approx(0.1)
+    assert log_data['test/sim_max_reward_10000'] == pytest.approx(0.9)
+    assert log_data['train_mean_score'] == pytest.approx(0.1)
+    assert log_data['test_mean_score'] == pytest.approx(0.9)
+    assert not any('sim_video' in key for key in log_data)

tests/test_train_diffusion_transformer_workspace_logging.py

@@ -0,0 +1,198 @@
+import importlib
+import pathlib
+import sys
+
+import pytest
+from omegaconf import OmegaConf
+
+ROOT_DIR = pathlib.Path(__file__).resolve().parents[1]
+if str(ROOT_DIR) not in sys.path:
+    sys.path.append(str(ROOT_DIR))
+
+MODULE_NAME = 'diffusion_policy.workspace.train_diffusion_transformer_hybrid_workspace'
+
+
+def load_workspace_module(monkeypatch, *, wandb_missing=False):
+    sys.modules.pop(MODULE_NAME, None)
+    if wandb_missing:
+        monkeypatch.setitem(sys.modules, 'wandb', None)
+    return importlib.import_module(MODULE_NAME)
+
+
+def test_init_logger_uses_swanlab_backend_mapping_without_loading_wandb(tmp_path, monkeypatch):
+    workspace_module = load_workspace_module(monkeypatch, wandb_missing=True)
+    events = []
+
+    class FakeRun:
+        def log(self, payload, step=None):
+            events.append(('log', payload, step))
+
+        def finish(self):
+            events.append(('finish',))
+
+    class FakeSwanLab:
+        def init(self, **kwargs):
+            events.append(('init', kwargs))
+            return FakeRun()
+
+    monkeypatch.setattr(workspace_module, '_load_swanlab', lambda: FakeSwanLab())
+    monkeypatch.setattr(
+        workspace_module,
+        '_load_wandb',
+        lambda: pytest.fail('wandb should not be loaded for the SwanLab backend'),
+    )
+
+    cfg = OmegaConf.create({
+        'logging': {
+            'backend': 'swanlab',
+            'project': 'demo-project',
+            'name': 'demo-run',
+            'group': 'demo-group',
+            'tags': ['pusht', 'dit'],
+            'id': 'run-123',
+            'resume': True,
+            'mode': 'online',
+        }
+    })
+
+    logger = workspace_module.init_logging_backend(cfg=cfg, output_dir=tmp_path)
+    logger.log({'metric': 1.0}, step=7)
+    logger.finish()
+
+    assert events[0][0] == 'init'
+    init_kwargs = events[0][1]
+    assert init_kwargs['project'] == 'demo-project'
+    assert init_kwargs['experiment_name'] == 'demo-run'
+    assert init_kwargs['group'] == 'demo-group'
+    assert init_kwargs['tags'] == ['pusht', 'dit']
+    assert init_kwargs['id'] == 'run-123'
+    assert init_kwargs['resume'] is True
+    assert init_kwargs['mode'] == 'cloud'
+    assert init_kwargs['logdir'] == str(tmp_path / 'swanlog')
+    assert ('log', {'metric': 1.0}, 7) in events
+    assert events.count(('finish',)) == 1
+
+
+def test_init_logger_defaults_to_legacy_wandb_path_when_backend_missing(tmp_path, monkeypatch):
+    workspace_module = load_workspace_module(monkeypatch)
+    events = []
+
+    class FakeRun:
+        def log(self, payload, step=None):
+            events.append(('log', payload, step))
+
+        def finish(self):
+            events.append(('finish',))
+
+    class FakeConfig:
+        def update(self, payload):
+            events.append(('config.update', payload))
+
+    class FakeWandb:
+        def __init__(self):
+            self.config = FakeConfig()
+
+        def init(self, **kwargs):
+            events.append(('init', kwargs))
+            return FakeRun()
+
+    monkeypatch.setattr(workspace_module, '_load_wandb', lambda: FakeWandb())
+
+    cfg = OmegaConf.create({
+        'logging': {
+            'project': 'demo-project',
+            'name': 'demo-run',
+            'group': None,
+            'tags': ['shared'],
+            'id': None,
+            'resume': True,
+            'mode': 'online',
+        }
+    })
+
+    logger = workspace_module.init_logging_backend(cfg=cfg, output_dir=tmp_path)
+    logger.log({'metric': 2.0}, step=3)
+    logger.finish()
+
+    assert events[0][0] == 'init'
+    init_kwargs = events[0][1]
+    assert init_kwargs['dir'] == str(tmp_path)
+    assert init_kwargs['project'] == 'demo-project'
+    assert init_kwargs['name'] == 'demo-run'
+    assert init_kwargs['mode'] == 'online'
+    assert ('config.update', {'output_dir': str(tmp_path)}) in events
+    assert ('log', {'metric': 2.0}, 3) in events
+    assert events.count(('finish',)) == 1
+
+
+def test_init_logger_rejects_unknown_backends(tmp_path, monkeypatch):
+    workspace_module = load_workspace_module(monkeypatch)
+    cfg = OmegaConf.create({
+        'logging': {
+            'backend': 'tensorboard',
+            'project': 'demo-project',
+            'name': 'demo-run',
+            'mode': 'offline',
+        }
+    })
+
+    with pytest.raises(ValueError, match='Unknown logging backend'):
+        workspace_module.init_logging_backend(cfg=cfg, output_dir=tmp_path)
+
+
+
+
+def test_logging_backend_session_preserves_primary_exception_when_finish_fails(tmp_path, monkeypatch):
+    workspace_module = load_workspace_module(monkeypatch)
+    events = []
+
+    class FakeBackend:
+        def log(self, payload, step=None):
+            events.append(('log', payload, step))
+
+        def finish(self):
+            events.append(('finish',))
+            raise RuntimeError('finish boom')
+
+    monkeypatch.setattr(
+        workspace_module,
+        'init_logging_backend',
+        lambda cfg, output_dir: FakeBackend(),
+    )
+
+    cfg = OmegaConf.create({'logging': {'mode': 'offline'}})
+
+    with pytest.raises(ValueError, match='primary boom'):
+        with workspace_module.logging_backend_session(cfg=cfg, output_dir=tmp_path) as logger:
+            logger.log({'metric': 6.0}, step=12)
+            raise ValueError('primary boom')
+
+    assert ('log', {'metric': 6.0}, 12) in events
+    assert events.count(('finish',)) == 1
+
+def test_logging_backend_session_finishes_on_exception(tmp_path, monkeypatch):
+    workspace_module = load_workspace_module(monkeypatch)
+    events = []
+
+    class FakeBackend:
+        def log(self, payload, step=None):
+            events.append(('log', payload, step))
+
+        def finish(self):
+            events.append(('finish',))
+
+    monkeypatch.setattr(
+        workspace_module,
+        'init_logging_backend',
+        lambda cfg, output_dir: FakeBackend(),
+    )
+
+    cfg = OmegaConf.create({'logging': {'mode': 'offline'}})
+
+    with pytest.raises(RuntimeError, match='boom'):
+        with workspace_module.logging_backend_session(cfg=cfg, output_dir=tmp_path) as logger:
+            logger.log({'metric': 5.0}, step=11)
+            raise RuntimeError('boom')
+
+    assert ('log', {'metric': 5.0}, 11) in events
+    assert events.count(('finish',)) == 1