Mujoco_WASM/web_platform/src/tuning/domain.ts

import type {
  ObjectiveWeights,
  RewardConfiguration,
  TuningSessionState,
} from '../training/types';

export type RewardSection = keyof RewardConfiguration;

export interface ParameterDefinition {
  path: string;
  section: RewardSection;
  key: string;
  label: string;
  minimum: number;
  maximum: number;
  defaultValue: number;
  allowZero: boolean;
  precision: number;
}

const weight = (
  key: string,
  label: string,
  minimum: number,
  maximum: number,
  defaultValue: number,
  allowZero = true,
  precision = 4,
): ParameterDefinition => ({
  path: `weights.${key}`,
  section: 'weights',
  key,
  label,
  minimum,
  maximum,
  defaultValue,
  allowZero,
  precision,
});

const parameter = (
  key: string,
  label: string,
  minimum: number,
  maximum: number,
  defaultValue: number,
  precision = 4,
): ParameterDefinition => ({
  path: `params.${key}`,
  section: 'params',
  key,
  label,
  minimum,
  maximum,
  defaultValue,
  allowZero: false,
  precision,
});

/** 与 training_server/tuning/schema.py 同步的只读前端目录；服务端仍是安全边界。 */
export const PARAMETER_DEFINITIONS: readonly ParameterDefinition[] = [
  weight('track_linear_velocity', '线速度跟踪', 0.5, 3, 1, false),
  weight('track_angular_velocity', '角速度跟踪', 0.25, 2, 1, false),
  weight('body_orientation_l2', '躯干姿态', -3, -0.1, -1, false),
  weight('pose', '目标姿态', 0, 2.5, 1),
  weight('body_ang_vel', '机身角速度', -0.2, 0, -0.05),
  weight('angular_momentum', '角动量', -0.1, 0, -0.025),
  weight('is_terminated', '跌倒终止', -400, -50, -200, false, 2),
  weight('joint_acc_l2', '关节加速度', -2e-6, 0, -2.5e-7, true, 8),
  weight('joint_pos_limits', '关节限位', -30, -2, -10, false, 2),
  weight('action_rate_l2', '动作平滑', -0.2, -0.005, -0.05),
  weight('foot_gait', '步态相位', 0, 1.5, 0.5),
  weight('foot_clearance', '抬脚高度', -3, 0, -1),
  weight('foot_slip', '足端滑移', -1, 0, -0.25),
  weight('soft_landing', '柔和落足', -0.005, 0, -0.001, true, 6),
  weight('stand_still', '静止姿态', -3, 0, -1),
  weight('electrical_power', '电功率', -0.005, 0, 0, true, 6),
  parameter('track_linear_velocity.std', '线速度核宽', 0.25, 1, 0.5),
  parameter('track_angular_velocity.std', '角速度核宽', 0.35, 1.2, Math.sqrt(0.5)),
  parameter('pose.std_standing_scale', '站立姿态尺度', 0.5, 2, 1),
  parameter('pose.std_walking_scale', '行走姿态尺度', 0.5, 2, 1),
  parameter('pose.std_running_scale', '奔跑姿态尺度', 0.5, 2, 1),
  parameter('pose.walking_threshold', '行走阈值', 0.05, 0.5, 0.1),
  parameter('pose.running_threshold', '奔跑阈值', 1, 2.5, 1.5),
  parameter('foot_gait.period', '步态周期', 0.4, 0.8, 0.6),
  parameter('foot_gait.threshold', '步态阈值', 0.45, 0.65, 0.56),
  parameter('foot_gait.command_threshold', '步态命令阈值', 0.02, 0.3, 0.1),
  parameter('foot_clearance.target_height', '目标抬脚高度', 0.06, 0.16, 0.1),
  parameter('foot_clearance.command_threshold', '抬脚命令阈值', 0.02, 0.3, 0.1),
  parameter('foot_slip.command_threshold', '滑移命令阈值', 0.02, 0.3, 0.1),
  parameter('soft_landing.command_threshold', '落足命令阈值', 0.02, 0.3, 0.1),
  parameter('stand_still.command_threshold', '静止命令阈值', 0.02, 0.3, 0.1),
] as const;

export const OBSTACLE_PARAMETERS: readonly ParameterDefinition[] = [
  weight('avoidance_weight', '近障避让', 0.5, 5, 2, false),
  weight('collision_penalty', '非法接触惩罚', -10, -0.5, -5, false),
  weight('action_smoothness', '动作平滑', -0.05, -0.001, -0.05, false),
  parameter('target_velocity', '目标导航速度（m/s，非奖励）', 0.3, 1.2, 0.6),
];
export const parameterDefinitions = (taskId?: string) =>
  taskId === 'Unitree-Go2-ObstacleAvoidance' ? OBSTACLE_PARAMETERS : PARAMETER_DEFINITIONS;
export const OBSTACLE_OBJECTIVES: ObjectiveWeights = {
  success: 0.4,
  time: 0.2,
  clearance: 0.2,
  smooth: 0.1,
  no_fall: 0.1,
};
export const OBSTACLE_OBJECTIVE_META = [
  { key: 'success', label: '无碰撞到达率', shortLabel: '成功', metric: 'success' },
  { key: 'time', label: '成功到达时间', shortLabel: '时间', metric: 'time' },
  { key: 'clearance', label: '障碍净距', shortLabel: '净距', metric: 'clearance' },
  { key: 'smooth', label: '动作平滑', shortLabel: '平滑', metric: 'smooth' },
  { key: 'no_fall', label: '不跌倒', shortLabel: '防跌', metric: 'no_fall' },
];
export const objectiveMeta = (taskId?: string) =>
  taskId === 'Unitree-Go2-ObstacleAvoidance' ? OBSTACLE_OBJECTIVE_META : OBJECTIVE_META;

export const PARAMETER_BY_PATH = new Map(
  [...PARAMETER_DEFINITIONS, ...OBSTACLE_PARAMETERS].map((definition) => [
    definition.path,
    definition,
  ]),
);

export const OBJECTIVE_META: ReadonlyArray<{
  key: keyof ObjectiveWeights;
  label: string;
  shortLabel: string;
  metric: string;
}> = [
  {
    key: 'velocity_tracking',
    label: '速度跟踪',
    shortLabel: '速度',
    metric: 'linear_velocity_rmse',
  },
  {
    key: 'action_smoothness',
    label: '动作平滑度',
    shortLabel: '平滑',
    metric: 'mean_action_acc',
  },
  {
    key: 'posture_stability',
    label: '躯干姿态稳定',
    shortLabel: '姿态',
    metric: 'orientation_error',
  },
  {
    key: 'fall_avoidance',
    label: '跌倒规避',
    shortLabel: '防跌',
    metric: 'fall_rate',
  },
  { key: 'foot_slip', label: '接触滑移抑制', shortLabel: '滑移', metric: 'slip_velocity' },
  { key: 'energy', label: '机械能耗', shortLabel: '能耗', metric: 'mechanical_power' },
];

export const STATE_META = {
  queued: { label: '分析排队', phase: 0, tone: 'neutral' },
  running: { label: '策略训练中', phase: 1, tone: 'accent' },
  evaluating: { label: '固定协议评估中', phase: 2, tone: 'warning' },
  awaiting_approval: { label: '等待人工审批', phase: 3, tone: 'warning' },
  paused: { label: '安全暂停', phase: 3, tone: 'neutral' },
  interrupted: { label: '服务已中断', phase: 3, tone: 'warning' },
  succeeded: { label: '调优完成', phase: 4, tone: 'success' },
  failed: { label: '调优失败', phase: 4, tone: 'warning' },
  cancelled: { label: '已取消', phase: 4, tone: 'neutral' },
} as const satisfies Record<
  TuningSessionState,
  { label: string; phase: number; tone: 'neutral' | 'accent' | 'success' | 'warning' }
>;

export const ACTIVE_SESSION_STATES = new Set<TuningSessionState>([
  'queued',
  'running',
  'evaluating',
  'awaiting_approval',
  'paused',
  'interrupted',
]);

export const TRIAL_STATE_LABELS: Record<string, string> = {
  queued: '等待调度',
  training: '训练中',
  evaluating: '评估中',
  completed: '完成',
  interrupted: '中断',
  failed: '失败',
  cancelled: '取消',
};

export function parameterValue(config: RewardConfiguration | undefined, path: string): number {
  const definition = PARAMETER_BY_PATH.get(path);
  if (!definition) return Number.NaN;
  return config?.[definition.section][definition.key] ?? definition.defaultValue;
}

export interface RewardDiffEntry {
  path: string;
  before: number;
  after: number;
}

export function rewardConfigurationDiff(
  before: RewardConfiguration,
  after: RewardConfiguration,
): RewardDiffEntry[] {
  const changes: RewardDiffEntry[] = [];
  for (const section of ['weights', 'params'] as const) {
    const keys = new Set([...Object.keys(before[section]), ...Object.keys(after[section])]);
    for (const key of keys) {
      const previous = before[section][key];
      const next = after[section][key];
      if (previous !== next)
        changes.push({ path: `${section}.${key}`, before: previous, after: next });
    }
  }
  return changes;
}

export function mergeRewardPatch(
  base: RewardConfiguration,
  patch: Partial<{ weights: Record<string, number>; params: Record<string, number> }>,
): RewardConfiguration {
  return {
    weights: { ...base.weights, ...(patch.weights ?? {}) },
    params: { ...base.params, ...(patch.params ?? {}) },
  };
}

export function formatMetric(value: number | null | undefined, digits = 4): string {
  if (value === null || value === undefined || !Number.isFinite(value)) return '—';
  const magnitude = Math.abs(value);
  if (magnitude !== 0 && (magnitude < 1e-3 || magnitude >= 1e4)) return value.toExponential(2);
  return value.toFixed(digits);
}