Mujoco_WASM/web_platform/e2e/lekiwi.training.spec.ts

import { expect, test } from '@playwright/test';
import { spawn } from 'node:child_process';
import { createServer } from 'node:net';
import { existsSync } from 'node:fs';
import { resolve } from 'node:path';
import { mkdir, readFile, readdir, writeFile } from 'node:fs/promises';

const mobilePython = process.env.MOBILE_TRAINING_PYTHON ?? resolve('build/venvs/mobile/bin/python');
const bundle =
  process.env.MOBILE_BUNDLE_ZIP ??
  resolve('../Reference_Projects/LeKiwi/New_urdf/robot_urdf_bundle (1).zip');
for (const [taskId, file] of [
  ['MobileManipulator-LeKiwi-v1', resolve('build/lekiwi/lekiwi-v1.zip')],
  ['MobileManipulator-LeKiwi-Bundle', bundle],
]) {
  test(`主控制台真实一键训练及策略导入：${taskId}`, async ({ page }, info) => {
    test.skip(
      !existsSync(mobilePython) || !existsSync(file),
      '需要 MuJoCo 3.11/SB3 独立环境及机器人资产',
    );
    const port = await new Promise<number>((resolvePort) => {
      const listener = createServer();
      listener.listen(0, '127.0.0.1', () => {
        const port = (listener.address() as { port: number }).port;
        listener.close(() => resolvePort(port));
      });
    });
    const root = info.outputPath('training');
    await mkdir(root, { recursive: true });
    const token = 'oneclick-test-token';
    const endpoint = `http://127.0.0.1:${port}`;
    const process = spawn(
      resolve('.venv/bin/python'),
      [
        '-u',
        'training_server/server.py',
        '--port',
        String(port),
        '--trainer-root',
        root,
        '--mobile-python',
        mobilePython,
        '--token',
        token,
      ],
      { stdio: ['ignore', 'pipe', 'pipe'] },
    );
    let output = '';
    process.stdout.on('data', (data) => {
      output += String(data);
    });
    process.stderr.on('data', (data) => {
      output += String(data);
    });
    const errors: string[] = [];
    page.on('pageerror', (error) => errors.push(error.message));
    try {
      await expect
        .poll(
          async () => {
            try {
              return (
                await fetch(`${endpoint}/api/training/health`, {
                  headers: { Authorization: `Bearer ${token}` },
                })
              ).ok;
            } catch {
              return false;
            }
          },
          { timeout: 30000 },
        )
        .toBe(true);
      await page.goto('/');
      await page.locator('input[type="file"]').first().setInputFiles(file);
      const dialog = page.getByRole('dialog', { name: '配置 URDF 仿真组件' });
      await expect(dialog).toBeVisible({ timeout: 30000 });
      // v1 also covers migration from the existing external-control profile at start.
      await dialog
        .getByLabel('URDF 机器人 profile')
        .selectOption(taskId === 'MobileManipulator-LeKiwi-v1' ? 'lekiwi-v1' : taskId);
      await dialog.getByRole('button', { name: '转换并加载' }).click();
      await expect(page.getByRole('button', { name: '▶ 播放' })).toBeEnabled({ timeout: 60000 });
      await page.getByRole('tab', { name: '控制台', exact: true }).click();
      if (taskId === 'MobileManipulator-LeKiwi-v1') {
        await page.getByRole('button', { name: 'LeKiwi 智能抓放', exact: true }).click();
        const agent = page.getByLabel('LeKiwi 智能抓放任务', { exact: true });
        await agent.getByRole('button', { name: '单技能步进', exact: true }).click();
        await expect(agent.getByTestId('agent-task-state')).toContainText('任务已暂停', {
          timeout: 60000,
        });
      }
      await page.getByRole('button', { name: '强化学习任务', exact: true }).click();
      await page.getByLabel('本地训练服务地址').fill(endpoint);
      await page.getByLabel('训练服务访问令牌').fill(token);
      await page.getByRole('button', { name: '连接', exact: true }).click();
      await expect(page.getByLabel('训练任务')).toHaveValue(taskId);
      await page.getByLabel('训练迭代', { exact: true }).fill('2');
      await page.getByLabel('每环境采样步数').fill('8');
      await expect(page.getByLabel('移动操作训练阶段')).toHaveValue('navigate');
      await page.getByLabel('独立评估回合').fill('2');
      await page.getByLabel('导航启动示教步数（仅初训）').fill('0');
      const jobResponse = page.waitForResponse(
        (response) =>
          response.url() === `${endpoint}/api/training/jobs` &&
          response.request().method() === 'POST',
      );
      await page.getByRole('button', { name: '发起本地训练' }).click();
      const created = (await (await jobResponse).json()) as { id: string };
      await expect(page.getByRole('button', { name: '导入策略', exact: true })).toBeEnabled({
        timeout: 60000,
      });
      await page.getByRole('button', { name: '导入策略', exact: true }).click();
      await expect
        .poll(
          async () =>
            page.evaluate(async () => {
              const path = '/src/stores/useAppStore.ts';
              const { useAppStore } = (await import(
                /* @vite-ignore */ path
              )) as typeof import('../src/stores/useAppStore');
              return useAppStore.getState().snapshot?.rlPolicy?.inferenceCount ?? 0;
            }),
          { timeout: 30000 },
        )
        .toBeGreaterThan(0);
      const job = await (
        await fetch(`${endpoint}/api/training/jobs/${created.id}`, {
          headers: { Authorization: `Bearer ${token}` },
        })
      ).json();
      expect(job.state).toBe('succeeded');
      expect(job.logs.join('\n')).toContain('Mean reward:');
      expect(job.deployment.input.shape).toEqual([1, 92]);
      expect(job.deployment.trainingStage).toBe('navigate');
      expect(job.deployment.evaluation.episodes).toBe(2);
      expect(job.deployment.evaluation.safetyStops).toBe(0);
      expect(job.deployment.evaluation.maxJointVelocity).toBeLessThan(2);
      expect(job.deployment.output.shape).toEqual([1, 12]);
      if (taskId === 'MobileManipulator-LeKiwi-v1') {
        const scenes = (await readdir(root, { recursive: true })).filter((name) =>
          name.endsWith('.training.xml'),
        );
        expect(scenes.length).toBeGreaterThan(0);
        for (const scene of scenes)
          expect(await readFile(resolve(root, scene), 'utf8')).not.toContain('__agent_');
        await expect(page.getByTestId('agent-task-state')).toContainText('已取消');
      }
      await writeFile(info.outputPath('job.json'), JSON.stringify(job, null, 2));
      await page.screenshot({ path: info.outputPath('unified-training.png') });
      expect(errors).toEqual([]);
      if (taskId === 'MobileManipulator-LeKiwi-v1') {
        await page.getByRole('button', { name: '接续此作业（保留权重）' }).click();
        await expect(page.getByLabel('接续作业 ID')).toHaveValue(created.id);
        const resumedResponse = page.waitForResponse(
          (response) =>
            response.url() === `${endpoint}/api/training/jobs` &&
            response.request().method() === 'POST',
        );
        await page.getByRole('button', { name: '发起本地训练' }).click();
        const response = await resumedResponse;
        expect(response.status()).toBe(202);
        const resumed = await response.json();
        await expect(page.getByRole('button', { name: '导入策略', exact: true })).toBeEnabled({
          timeout: 60000,
        });
        const second = await (
          await fetch(`${endpoint}/api/training/jobs/${resumed.id}`, {
            headers: { Authorization: `Bearer ${token}` },
          })
        ).json();
        expect(second.state).toBe('succeeded');
        expect(second.deployment.initialization).toBe('resumed-checkpoint');
        expect(second.deployment.trainedTimesteps).toBeGreaterThan(job.deployment.trainedTimesteps);
        await writeFile(info.outputPath('resumed-job.json'), JSON.stringify(second, null, 2));
      }
    } finally {
      await page.close();
      process.kill('SIGTERM');
      await new Promise<void>((resolveExit) => {
        if (process.exitCode !== null) resolveExit();
        else process.once('exit', () => resolveExit());
      });
      await writeFile(info.outputPath('server.log'), output.replaceAll(token, '<redacted>'));
    }
  });
}