# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""DreamWaQ environment for MotrixLab — 对齐 Manaro-Alpha/DreamWaQ。

架构:
  CENet (VAE):  history(225) → [128,64] → latent(16) + vel_est(3) = code(19)
  Actor:        code(19) + obs(45) = 64 → [512,256,128] → action(12)
  Critic:       privileged_obs(247) → [512,256,128] → value(1)

观测 (45-dim, 上游顺序):
  ang_vel(3) + gravity(3) + commands(3) + joint_pos(12) + joint_vel(12) + actions(12)

特权观测 (247-dim):
  obs(45) + base_vel(3) + 足部接触力(12) + heights(187)
  注意：上游为 286（含全部 17 body × 3-axis 接触力），MotrixSim 仅用四足接触力

注册: "go1-dreamwaq-walk"
"""
import gymnasium as gym
import motrixsim as mtx
import numpy as np
from dataclasses import dataclass, field
import os

from motrix_envs import registry
from motrix_envs.locomotion.go1.cfg import Go1WalkNpEnvCfg
from motrix_envs.locomotion.go1.walk_np import Go1WalkTask
from motrix_envs.math import quaternion


_SCENE_PRINTED = False


def _scene_file():
    """选择地形场景。DREAMWAQ_TERRAIN=flat|pyramid|flat_stairs。默认 FLAT。"""
    global _SCENE_PRINTED
    terrain = os.environ.get("DREAMWAQ_TERRAIN", "flat").lower()
    if terrain in ("flat_stairs", "flatstairs"):
        fname = "scene_flat_stairs.xml"
    elif terrain in ("pyramid", "terrain", "1"):
        fname = "scene_dreamwaq_terrain.xml"
    elif terrain in ("stairs", "stairs_terrain"):
        fname = "scene_stairs_terrain.xml"
    else:
        fname = "scene_dreamwaq_flat.xml"
    if not _SCENE_PRINTED:
        print(f"[DreamWaQ] scene = {fname}  [set DREAMWAQ_TERRAIN to switch]")
        _SCENE_PRINTED = True
    return os.path.join(os.path.dirname(__file__), "xmls", fname)


# ═══════════════════════════════════════════════════════════════════════
# Config
# ═══════════════════════════════════════════════════════════════════════

@registry.envcfg("go1-dreamwaq-walk")
@dataclass
class DreamWaQCfg(Go1WalkNpEnvCfg):
    """DreamWaQ 配置 — 10×20 地形网格，5 种类型，10 个难度级别。"""

    model_file: str = field(default_factory=_scene_file)

    # CENet 参数
    num_latent: int = 16
    cenet_out_dim: int = 19         # code = vel(3) + latent(16)
    num_history: int = 5

    # 特权观测维度：obs(45) + base_vel(3) + 足部接触力(12) + heights(187) = 247
    num_privileged_obs: int = 247

    max_episode_steps: int = 1000  # 与上游一致: 20s / 0.02s ctrl_dt
    sim_dt: float = 0.005
    ctrl_dt: float = 0.02

    # 高度测量网格（与上游一致: 17×11）
    height_points_x: tuple = tuple(np.linspace(-0.8, 0.8, 17).tolist())
    height_points_y: tuple = tuple(np.linspace(-0.5, 0.5, 11).tolist())

    # 地形网格
    terrain_rows: int = 10
    terrain_cols: int = 20  # 上游 10×20 网格：5 类型 × 10 难度
    cell_size: float = 8.0
    border_size: float = 5.0

    # 命令范围（上游: [-1,1] 对称）
    @dataclass
    class Commands:
        vel_limit = [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]]

    commands: Commands = field(default_factory=Commands)

    def __init__(self):
        super().__init__()
        self.model_file = _scene_file()
        self.num_latent = 16
        self.cenet_out_dim = 19
        self.num_history = 5
        self.num_privileged_obs = 247
        self.sim_dt = 0.005
        self.ctrl_dt = 0.02
        if "flat_stairs" in self.model_file:
            self.terrain_rows = 2
        else:
            self.terrain_rows = 10
        self.terrain_cols = 20
        self.cell_size = 8.0
        self.border_size = 5.0
        self.height_points_x = tuple(np.linspace(-0.8, 0.8, 17).tolist())
        self.height_points_y = tuple(np.linspace(-0.5, 0.5, 11).tolist())
        self.commands = DreamWaQCfg.Commands()
        self.control_config.stiffness = 28.0
        self.control_config.damping = 0.7
        self.control_config.action_scale = 0.25  # 上游原值
        self.noise_config.scale_joint_angle = 0.01  # 上游原值（基类默认 0.03）
        self._apply_dreamwaq()

    def _apply_dreamwaq(self):
        """DreamWaQ 奖励尺度 — 与上游 Go1RoughCfg 一致（上游会在 _prepare 中 × dt）。"""
        r = self.reward_config.scales
        r.clear()
        r.update({
            "tracking_lin_vel": 1.0,
            "tracking_ang_vel": 0.5,
            "lin_vel_z": -2.0,
            "ang_vel_xy": -0.05,
            "orientation": -0.2,
            "dof_acc": -2.5e-7,
            "base_height": -1.0,
            "feet_air_time": 0.1,
            "action_rate": -0.01,
            "joint_power": -2e-5,
            "smoothness": -0.01,
            "power_distribution": -10e-6,
            # 注意：stand_still 在上游被注释掉
        })
        self.reward_config.only_positive_rewards = True
        self.reward_config.tracking_sigma = 0.25


# ═══════════════════════════════════════════════════════════════════════
# DreamWaQ 环境
# ═══════════════════════════════════════════════════════════════════════

@registry.env("go1-dreamwaq-walk", sim_backend="np")
class DreamWaQTask(Go1WalkTask):
    """DreamWaQ 环境：不对称观测、历史缓冲区、高度测量。"""

    _cfg: DreamWaQCfg

    def __init__(self, cfg: DreamWaQCfg = None, num_envs=1):
        if cfg is None:
            cfg = DreamWaQCfg()
        super().__init__(cfg, num_envs)

        # 覆盖观测空间（DreamWaQ 45-dim，比基类少 linvel 3-dim）
        self._observation_space = gym.spaces.Box(
            low=-np.inf, high=np.inf, shape=(45,), dtype=np.float32)
        self._num_observation = 45

        self._hx = np.array(cfg.height_points_x, dtype=np.float32)
        self._hy = np.array(cfg.height_points_y, dtype=np.float32)
        self._num_rows = cfg.terrain_rows
        self._num_cols = cfg.terrain_cols
        self._cell_size = cfg.cell_size
        self._border = cfg.border_size
        self._init_done = False
        self._all_terrain_types = np.tile(
            np.arange(self._num_cols), num_envs // self._num_cols + 1)[:num_envs]
        np.random.shuffle(self._all_terrain_types)
        self._hf_cache = None
        self._hm_cache = None

    # ── 命令：全范围，无课程（匹配上游）──

    def resample_commands(self, num_envs: int) -> np.ndarray:
        lim = np.array(self.cfg.commands.vel_limit, dtype=np.float32)
        cmds = np.random.uniform(lim[0], lim[1], size=(num_envs, 3)).astype(np.float32)
        small = np.linalg.norm(cmds[:, :2], axis=1) < 0.2
        cmds[small, :2] = 0.0
        return cmds

    # ── 动作裁剪 + 力矩计算 ──

    def apply_action(self, actions, state):
        """裁剪动作防止奖励计算溢出（上游 clip_actions=100）。"""
        actions = np.clip(actions, -100.0, 100.0)
        return super().apply_action(actions, state)

    def _compute_torques(self, actions, data):
        """PD 控制器 + 域随机化。力矩裁剪防溢出。"""
        state = getattr(self, '_state', None)
        if state is not None:
            motor_strength = state.info.get("motor_strength",
                np.ones(self._num_envs, dtype=np.float32))
            kp_factor = state.info.get("kp_factor",
                np.ones(self._num_envs, dtype=np.float32))
            kd_factor = state.info.get("kd_factor",
                np.ones(self._num_envs, dtype=np.float32))
        else:
            motor_strength = np.ones(1, dtype=np.float32)
            kp_factor = np.ones(1, dtype=np.float32)
            kd_factor = np.ones(1, dtype=np.float32)

        actions_scaled = actions * self.cfg.control_config.action_scale * motor_strength[:, np.newaxis]
        torques = (self.kps * kp_factor[:, np.newaxis]) * (
            actions_scaled + self.default_angles - self.get_dof_pos(data)
        ) - (self.kds * kd_factor[:, np.newaxis]) * self.get_dof_vel(data)
        return np.clip(torques, -80.0, 80.0)

    # ── 观测 ──

    def _get_obs(self, data: mtx.SceneData, info: dict) -> np.ndarray:
        """45-dim: ang_vel + gravity + commands + joint_pos + joint_vel + actions。"""
        gyro = self.get_gyro(data)
        pose = self._body.get_pose(data)
        base_quat = pose[:, 3:7]
        gravity = quaternion.rotate_inverse(base_quat, self.gravity_vec)
        commands = info["commands"] * self.commands_scale
        joint_pos = (self.get_dof_pos(data) - self.default_angles) * self.cfg.normalization.dof_pos
        joint_vel = self.get_dof_vel(data) * self.cfg.normalization.dof_vel
        noisy_gyro = gyro * self.cfg.normalization.ang_vel
        actions = info["current_actions"]
        obs = np.concatenate([
            noisy_gyro, gravity, commands, joint_pos, joint_vel, actions
        ], axis=-1)
        # 观测噪声
        nc = self.cfg.noise_config
        noise_scale_vec = np.zeros(45, dtype=np.float32)
        # 上游: noise_vec = noise_scale * obs_scale（noise_level 单独乘在 noise 张量上）
        n = self.cfg.normalization
        noise_scale_vec[0:3] = nc.scale_gyro * n.ang_vel
        noise_scale_vec[3:6] = nc.scale_gravity
        noise_scale_vec[6:9] = 0.0
        noise_scale_vec[9:21] = nc.scale_joint_angle * n.dof_pos
        noise_scale_vec[21:33] = nc.scale_joint_vel * n.dof_vel
        noise_scale_vec[33:45] = 0.0
        noise = (2.0 * np.random.rand(*obs.shape).astype(np.float32) - 1.0) * nc.level
        obs = obs + noise * noise_scale_vec[np.newaxis, :]
        return obs

    # ── 高度测量 ──

    def _get_heights(self, data: mtx.SceneData) -> np.ndarray:
        n = data.shape[0]
        nx, ny = len(self._hx), len(self._hy)
        pose = self._body.get_pose(data)
        base_pos = pose[:, 0:3]
        yaw = quaternion.get_yaw(pose[:, 3:7])
        cos_yaw, sin_yaw = np.cos(yaw), np.sin(yaw)
        try:
            if self._hm_cache is None:
                hf = self._model.get_hfield(0)
                self._hf_cache = hf
                hm = hf.height_matrix
                self._hm_cache = hm
                nr, nc = hm.shape
                b = hf.bound
                self._h_xmin, self._h_ymin = b[0], b[1]
                self._h_xmax, self._h_ymax = b[3], b[4]
                self._h_nr, self._h_nc = nr, nc
                cfg_base = getattr(self.cfg, "hfield_z_base", None)
                self._h_zbase = float(cfg_base) if cfg_base is not None else float(-hm[0, 0])
            hm, nr, nc = self._hm_cache, self._h_nr, self._h_nc
            xmin, ymin, xmax, ymax = self._h_xmin, self._h_ymin, self._h_xmax, self._h_ymax
            z_base = self._h_zbase
            w, h = xmax - xmin, ymax - ymin
        except Exception:
            return np.zeros((n, nx * ny), dtype=np.float32)

        heights = np.zeros((n, nx * ny), dtype=np.float32)
        idx = 0
        for iy in range(ny):
            gy = self._hy[iy]
            for ix in range(nx):
                gx = self._hx[ix]
                wx = base_pos[:, 0] + cos_yaw * gx - sin_yaw * gy
                wy = base_pos[:, 1] + sin_yaw * gx + cos_yaw * gy
                col_f = (wx - xmin) / max(w, 1e-6) * (nc - 1)
                row_f = (ymax - wy) / max(h, 1e-6) * (nr - 1)
                col = np.clip(np.nan_to_num(col_f, nan=0).astype(np.int32), 0, nc - 1)
                row = np.clip(np.nan_to_num(row_f, nan=0).astype(np.int32), 0, nr - 1)
                heights[:, idx] = hm[row, col] + z_base
                idx += 1
        return heights

    def _sample_terrain_height(self, xy: np.ndarray, radius: float = 0.0) -> np.ndarray:
        n = xy.shape[0]
        try:
            hm, nr, nc = self._hm_cache, self._h_nr, self._h_nc
            xmin, ymin, xmax, ymax = self._h_xmin, self._h_ymin, self._h_xmax, self._h_ymax
            z_base = self._h_zbase
            w, h = xmax - xmin, ymax - ymin
        except Exception:
            return np.zeros(n, dtype=np.float32)
        col = np.clip(np.nan_to_num((xy[:, 0] - xmin) / max(w, 1e-6) * (nc - 1), nan=0).astype(np.int32), 0, nc - 1)
        row = np.clip(np.nan_to_num((ymax - xy[:, 1]) / max(h, 1e-6) * (nr - 1), nan=0).astype(np.int32), 0, nr - 1)
        if radius <= 0:
            return (hm[row, col] + z_base).astype(np.float32)
        dc = max(1, int(radius / max(w, 1e-6) * (nc - 1)))
        dr = max(1, int(radius / max(h, 1e-6) * (nr - 1)))
        out = np.empty(n, dtype=np.float32)
        for i in range(n):
            r, c = int(row[i]), int(col[i])
            out[i] = hm[max(0, r - dr):r + dr + 1, max(0, c - dc):c + dc + 1].max()
        return out + z_base

    # ── 足部接触力 ──

    def _read_contact_forces(self, data: mtx.SceneData) -> np.ndarray:
        """读取四足接触力（body frame）→ (N, 12)。"""
        pose = self._body.get_pose(data)
        base_quat = pose[:, 3:7]
        forces = []
        for foot in ["FR", "FL", "RR", "RL"]:
            v = self._model.get_sensor_value(foot + "_foot_contact", data)
            v_body = quaternion.rotate_inverse(base_quat, v)
            forces.append(v_body)
        return np.concatenate(forces, axis=1)

    def _get_privileged_obs(self, data: mtx.SceneData, obs: np.ndarray) -> np.ndarray:
        base_vel = self.get_local_linvel(data) * self.cfg.normalization.lin_vel
        heights = self._get_heights(data)
        contact_forces = self._read_contact_forces(data)
        return np.concatenate([obs, base_vel, contact_forces, heights], axis=-1)

    # ── update_observation ──

    def update_observation(self, state):
        data = state.data
        # 清理物理崩溃残留的 NaN/Inf（仅在检测到时才修改，避免无谓开销）
        dv = data.dof_vel
        if np.any(~np.isfinite(dv)):
            data.set_dof_vel(np.nan_to_num(np.array(dv), nan=0.0, posinf=0.0, neginf=0.0))
        dp = data.dof_pos
        if np.any(~np.isfinite(dp)):
            data.set_dof_pos(np.nan_to_num(np.array(dp), nan=0.0, posinf=0.0, neginf=0.0), self._model)
        obs = self._get_obs(data, state.info)

        # 更新历史缓冲区
        old_history = state.info.get("obs_history",
            np.zeros((self._num_envs, self._cfg.num_history, 45), dtype=np.float32))
        new_history = np.concatenate([old_history[:, 1:, :], obs[:, np.newaxis, :]], axis=1)
        state.info["obs_history"] = new_history

        current_step = state.info.get("steps", np.zeros(self._num_envs, dtype=np.int32))

        # 中期 command 重采样（每 10 秒）
        resample_steps = int(10.0 / self.cfg.ctrl_dt)
        do_resample = (current_step % resample_steps) == 0
        if do_resample.any():
            n_resample = int(do_resample.sum())
            state.info["commands"][do_resample] = self.resample_commands(n_resample)

        # 周期性域随机化（每 4s）
        rand_interval = int(4.0 / self.cfg.ctrl_dt)
        last_rand = state.info.get("last_rand_step", np.zeros(self._num_envs, dtype=np.int32))
        do_rand = (current_step - last_rand) >= rand_interval
        if do_rand.any():
            n_rand = int(do_rand.sum())
            state.info["motor_strength"][do_rand] = np.random.uniform(0.9, 1.1, size=n_rand).astype(np.float32)
            state.info["kp_factor"][do_rand] = np.random.uniform(0.9, 1.1, size=n_rand).astype(np.float32)
            state.info["kd_factor"][do_rand] = np.random.uniform(0.9, 1.1, size=n_rand).astype(np.float32)
            state.info["friction_coeff"][do_rand] = np.random.uniform(0.2, 1.25, size=n_rand).astype(np.float32)
            state.info["added_mass"][do_rand] = np.random.uniform(-1.0, 2.0, size=n_rand).astype(np.float32)
            state.info["com_displacement"][do_rand] = np.random.uniform(-0.05, 0.05, size=(n_rand, 3)).astype(np.float32)
            state.info["last_rand_step"][do_rand] = current_step[do_rand]

        # 特权观测
        state.info["privileged_obs"] = self._get_privileged_obs(data, obs)
        state.info["base_vel"] = self.get_local_linvel(data)

        # 足部接触
        cquerys = self._model.get_contact_query(data)
        state.info["contacts"] = cquerys.is_colliding(self.foot_check).reshape(
            (self._num_envs, self.foot_check_num))
        state.info["feet_air_time"] = self.update_feet_air_time(state.info)

        # 累计命令距离和跟踪
        cmd_speed = np.linalg.norm(state.info["commands"][:, :2], axis=1)
        state.info["ep_cmd_distance"] = state.info.get("ep_cmd_distance",
            np.zeros(self._num_envs, dtype=np.float32)) + cmd_speed * self.cfg.ctrl_dt
        state.info["ep_steps"] = state.info.get("ep_steps", np.zeros(self._num_envs, dtype=np.int32)) + 1
        local_vel = self.get_local_linvel(data)[:, :2]
        cmd_vel = state.info["commands"][:, :2]
        vel_error = np.sum(np.square(cmd_vel - local_vel), axis=1)
        tracking = np.exp(-vel_error / self.cfg.reward_config.tracking_sigma)
        state.info["ep_tracking_sum"] = state.info.get("ep_tracking_sum",
            np.zeros(self._num_envs, dtype=np.float32)) + tracking

        return state.replace(obs=obs)

    # ── 出生点 ──

    def _make_origins(self, levels, indices):
        n = len(levels)
        half_x = self._border + self._num_cols * self._cell_size / 2.0
        half_y = self._border + self._num_rows * self._cell_size / 2.0
        origins = np.zeros((n, 2), dtype=np.float32)
        for i in range(n):
            row = levels[i]
            col = self._all_terrain_types[int(indices[i]) % len(self._all_terrain_types)]
            cx = -half_x + self._border + col * self._cell_size + self._cell_size / 2
            cy = half_y - self._border - row * self._cell_size - self._cell_size / 2
            origins[i, 0] = cx
            origins[i, 1] = cy
        return origins

    # ── Reset ──

    def reset(self, data) -> tuple[np.ndarray, dict]:
        num_reset = data.shape[0]
        state = getattr(self, '_state', None)
        if state is not None and hasattr(state, 'done'):
            done = state.done
            done_idx = np.where(done)[0]
        else:
            done_idx = np.arange(num_reset)

        # 游戏启发式地形课程（per-env，与上游一致）
        if not hasattr(self, '_terrain_origins'):
            all_origins = np.zeros((self._num_rows, self._num_cols, 2), dtype=np.float32)
            half_x = self._border + self._num_cols * self._cell_size / 2.0
            half_y = self._border + self._num_rows * self._cell_size / 2.0
            for row in range(self._num_rows):
                for col in range(self._num_cols):
                    cx = -half_x + self._border + col * self._cell_size + self._cell_size / 2
                    cy = half_y - self._border - row * self._cell_size - self._cell_size / 2
                    all_origins[row, col] = [cx, cy]
            self._terrain_origins = all_origins
            self._max_init_level = 3  # 先 0-3，适应后再提到 5

        if num_reset > 0 and self._init_done and state is not None and hasattr(state, 'info'):
            old_info = state.info
            old_origins = old_info.get("env_origins", np.zeros((self._num_envs, 2), dtype=np.float32))[done]
            old_levels = old_info.get("terrain_level", np.zeros(self._num_envs, dtype=np.int32))[done]
            old_commands = old_info.get("commands", np.zeros((self._num_envs, 3), dtype=np.float32))[done]
            base_pose = self._body.get_pose(state.data)
            base_pos = base_pose[done, :2]
            distance = np.linalg.norm(base_pos - old_origins, axis=1)
            move_up = distance > 3.0  # 走过去 3m 就升级（上游是 4m）
            cmd_speed = np.linalg.norm(old_commands[:, :2], axis=1)
            required_dist = cmd_speed * (self.cfg.max_episode_steps * self.cfg.ctrl_dt) * 0.5
            move_down = (distance < required_dist) & ~move_up
            new_levels = np.where(move_up, old_levels + 1, old_levels)
            new_levels = np.where(move_down, new_levels - 1, new_levels)
            at_max = new_levels >= self._num_rows
            if at_max.any():
                new_levels[at_max] = np.random.randint(0, self._num_rows, size=int(at_max.sum()))
            new_levels = np.clip(new_levels, 0, self._num_rows - 1)
        else:
            new_levels = np.random.randint(0, self._max_init_level + 1, size=num_reset, dtype=np.int32)
        self._init_done = True

        _force = getattr(self, "_force_level", None)
        if _force is not None:
            new_levels = np.full(num_reset, int(_force), dtype=np.int32)

        if "scene_stairs_terrain" in self.cfg.model_file:
            height_list = np.array([-1.0, 0.5, 1.5], dtype=np.float32)
            offset_h = [[2, 0, 2, 1, 1], [2, 2, 1, 0, 0], [1, 1, 2, 1, 2],
                        [0, 1, 0, 2, 0], [0, 1, 1, 0, 2]]
            offsets = []
            for i in range(5):
                for j in range(5):
                    offsets.append([(i-2)*8.0, (j-2)*8.0, height_list[offset_h[j][i]]])
            offset_arr = np.array(offsets, dtype=np.float32)
            idx = np.random.choice(len(offsets), size=num_reset)
            new_origins = offset_arr[idx, :2]
            self._stairs_z_offsets = offset_arr[idx, 2]
        else:
            indices = done_idx
            if hasattr(self, '_terrain_origins'):
                new_origins = np.zeros((num_reset, 2), dtype=np.float32)
                for i in range(num_reset):
                    row = int(new_levels[i])
                    col = self._all_terrain_types[int(indices[i]) % len(self._all_terrain_types)]
                    new_origins[i] = self._terrain_origins[row, col]
            else:
                new_origins = self._make_origins(new_levels, indices)
            self._stairs_z_offsets = np.zeros(num_reset, dtype=np.float32)

        data.reset(self._model)

        init_dof_pos = np.tile(self._init_dof_pos, (num_reset, 1))
        init_dof_pos[:, 0] = new_origins[:, 0]
        init_dof_pos[:, 1] = new_origins[:, 1]
        if hasattr(self, '_stairs_z_offsets') and self._stairs_z_offsets.any():
            init_dof_pos[:, 2] = self._stairs_z_offsets
        else:
            terrain_z = self._sample_terrain_height(new_origins, radius=0.35)
            spawn_abs = getattr(self, "_spawn_absolute", None)
            if spawn_abs is not None:
                init_dof_pos[:, 2] = float(spawn_abs)
            else:
                init_dof_pos[:, 2] = terrain_z + getattr(self, "_spawn_clearance", 0.45)
        yaw = np.random.uniform(-np.pi, np.pi, size=num_reset)
        init_dof_pos[:, 3] = 0.0
        init_dof_pos[:, 4] = 0.0
        init_dof_pos[:, 5] = np.sin(yaw / 2)
        init_dof_pos[:, 6] = np.cos(yaw / 2)

        init_dof_vel = np.tile(self._init_dof_vel, (num_reset, 1))
        data.set_dof_vel(init_dof_vel)
        data.set_dof_pos(init_dof_pos, self._model)
        self._model.forward_kinematic(data)

        info = {
            "current_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
            "last_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
            "last_last_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
            "commands": self.resample_commands(num_reset),
            "last_dof_vel": np.zeros((num_reset, self._num_action), dtype=np.float32),
            "feet_air_time": np.zeros((num_reset, self.foot_check_num), dtype=np.float32),
            "contacts": np.zeros((num_reset, self.foot_check_num), dtype=np.bool),
            "last_contacts": np.zeros((num_reset, self.foot_check_num), dtype=np.bool),
            "motor_strength": np.random.uniform(0.9, 1.1, size=num_reset).astype(np.float32),
            "kp_factor": np.random.uniform(0.9, 1.1, size=num_reset).astype(np.float32),
            "kd_factor": np.random.uniform(0.9, 1.1, size=num_reset).astype(np.float32),
            "friction_coeff": np.random.uniform(0.2, 1.25, size=num_reset).astype(np.float32),
            "added_mass": np.random.uniform(-1.0, 2.0, size=num_reset).astype(np.float32),
            "com_displacement": np.random.uniform(-0.05, 0.05, size=(num_reset, 3)).astype(np.float32),
            "last_rand_step": np.zeros(num_reset, dtype=np.int32),
            "ep_cmd_distance": np.zeros(num_reset, dtype=np.float32),
            "ep_steps": np.zeros(num_reset, dtype=np.int32),
            "ep_tracking_sum": np.zeros(num_reset, dtype=np.float32),
        }
        obs = self._get_obs(data, info)

        hist = np.zeros((num_reset, self._cfg.num_history, 45), dtype=np.float32)
        hist[:, -1, :] = obs
        info["obs_history"] = hist

        info["privileged_obs"] = self._get_privileged_obs(data, obs)
        info["base_vel"] = self.get_local_linvel(data)
        info["terrain_level"] = new_levels
        info["env_origins"] = new_origins

        # 清除 ep_ 奖励累计器（跳过 ep_report 它是 dict）
        if state is not None:
            for k in list(state.info.keys()):
                if k.startswith("ep_") and k != "ep_report" and k not in info:
                    info[k] = np.zeros(num_reset, dtype=np.float32)

        return obs, info

    def update_terminated(self, state):
        """基类接触终止 + hfield 边界终止（防止走出地形导致 NaN）。"""
        state = super().update_terminated(state)
        pose = self._body.get_pose(state.data)
        base_xy = pose[:, :2]
        hf = self._model.get_hfield(0)
        b = hf.bound
        out_of_bounds = (
            (base_xy[:, 0] < b[0]) | (base_xy[:, 0] > b[3]) |
            (base_xy[:, 1] < b[1]) | (base_xy[:, 1] > b[4])
        )
        state.terminated = state.terminated | out_of_bounds
        return state

    # ── 奖励（与上游对齐，update_reward 中 × dt）──

    def _get_reward(self, data: mtx.SceneData, info: dict) -> dict[str, np.ndarray]:
        """DreamWaQ 奖励项——与上游 Manaro-Alpha 对齐。"""
        commands = info["commands"]
        return {
            "tracking_lin_vel": self._reward_tracking_lin_vel(data, commands),
            "tracking_ang_vel": self._reward_tracking_ang_vel(data, commands),
            "lin_vel_z": self._reward_lin_vel_z(data),
            "ang_vel_xy": self._reward_ang_vel_xy(data),
            "orientation": self._reward_orientation(data),
            "dof_acc": self._reward_dof_acc(data, info),
            "base_height": self._reward_base_height(data),
            "feet_air_time": self._reward_feet_air_time(commands, info),
            "action_rate": self._reward_action_rate(info),
            "joint_power": self._reward_joint_power(data),
            "smoothness": self._reward_smoothness(info),
            "power_distribution": self._reward_power_distribution(data),
            # stand_still 在上游被注释掉
        }

    def update_feet_air_time(self, info: dict):
        """覆盖基类 — 在清零前保存 first_contact 供奖励函数使用。"""
        feet_air_time = info["feet_air_time"]
        contacts = info["contacts"]
        # 保存着地瞬间的状态（脚刚触地，且之前腾空时间 > 0）
        info["first_contact"] = (feet_air_time > 0.0) & contacts
        info["air_time_at_contact"] = feet_air_time.copy()
        # 基类逻辑：累加 dt，着地清零
        feet_air_time = feet_air_time + self.cfg.ctrl_dt
        feet_air_time = feet_air_time * (~contacts)
        info["feet_air_time"] = feet_air_time
        return feet_air_time

    def _reward_feet_air_time(self, commands, info):
        """足部腾空时间奖励——与上游 legged_robot.py 完全对齐。

        使用 update_feet_air_time 保存的 first_contact / air_time_at_contact。
        """
        first_contact = info.get("first_contact")
        air_time = info.get("air_time_at_contact")
        if first_contact is None or air_time is None:
            return np.zeros(self._num_envs, dtype=np.float32)
        rew = np.sum((air_time - 0.5) * first_contact, axis=1)
        rew *= np.linalg.norm(commands[:, :2], axis=1) > 0.1
        return rew

    def update_reward(self, state):
        """存储各项奖励到 TensorBoard + 更新 state.reward。

        与上游 _prepare_reward_function 对齐：所有 scales × dt。
        """
        reward_dict = self._get_reward(state.data, state.info)
        scales = self._cfg.reward_config.scales
        dt = self._cfg.ctrl_dt
        scaled_terms = {
            k: v * scales.get(k, 0.0) * dt
            for k, v in reward_dict.items()
        }
        # state.reward 直接用带 dt 的值（不调 super，super 不带 dt）
        state = state.replace(reward=np.zeros(self._num_envs, dtype=np.float32))
        for k, v in scaled_terms.items():
            v = np.nan_to_num(v, nan=0.0, posinf=0.0, neginf=0.0)
            state.reward += v
        if self._cfg.reward_config.only_positive_rewards:
            state.reward = np.maximum(state.reward, 0.0)

        # TensorBoard 日志
        state.info["reward_terms"] = {k: float(np.mean(v)) for k, v in scaled_terms.items()}
        for k, v in scaled_terms.items():
            ek = f"ep_{k}"
            state.info[ek] = state.info.get(ek, np.zeros(self._num_envs, dtype=np.float32)) + v

        # 对即将结束的 env 存 ep_report
        steps = state.info.get("steps", np.zeros(self._num_envs, dtype=np.int32))
        will_end = state.terminated | (steps >= self._cfg.max_episode_steps - 1)
        if will_end.any():
            # 除以 max_episode_length_s，与上游对齐（上游: ep_sum / max_episode_length_s）
            ep_len_s = self._cfg.max_episode_steps * self._cfg.ctrl_dt
            ep_report = state.info.get("ep_report", {})
            for k in scaled_terms:
                ek = f"ep_{k}"
                vals = state.info.get(ek, np.zeros(self._num_envs, dtype=np.float32))[will_end]
                ep_report[f"rew_{k}"] = float(np.mean(vals) / ep_len_s)
            tl = state.info.get("terrain_level", np.zeros(self._num_envs, dtype=np.int32))[will_end]
            ep_report["terrain_level"] = float(np.mean(tl))
            state.info["ep_report"] = ep_report

        return state

    # ── 额外奖励函数 ──

    def _reward_base_height(self, data):
        """惩罚偏离目标基础高度——与上游 Go1RoughCfg.base_height_target=0.30 一致。"""
        pose = self._body.get_pose(data)
        base_z = pose[:, 2]
        heights = self._get_heights(data)
        ground_level = np.mean(heights, axis=1)
        target = 0.30
        return np.square(base_z - ground_level - target)

    def _reward_joint_power(self, data):
        torque = np.nan_to_num(data.actuator_ctrls, nan=0.0)
        vel = np.nan_to_num(self.get_dof_vel(data), nan=0.0)
        torque = np.clip(torque, -100, 100)
        vel = np.clip(vel, -100, 100)
        return np.clip(np.sum(np.abs(torque * vel), axis=1), 0, 1e6)

    def _reward_smoothness(self, info):
        scale = self.cfg.control_config.action_scale
        da = self.default_angles
        a0 = info["current_actions"] * scale + da
        a1 = info["last_actions"] * scale + da
        a2 = info.get("last_last_actions", info["last_actions"]) * scale + da
        diff = np.square(np.clip(a0 - 2.0 * a1 + a2, -1e4, 1e4))
        diff = diff * (info["last_actions"] != 0)
        diff = diff * (info.get("last_last_actions", info["last_actions"]) != 0)
        return np.sum(diff, axis=1)

    def _reward_power_distribution(self, data):
        torque = np.nan_to_num(data.actuator_ctrls, nan=0.0)
        vel = np.nan_to_num(self.get_dof_vel(data), nan=0.0)
        torque = np.clip(torque, -100, 100)
        vel = np.clip(vel, -100, 100)
        power = torque * vel
        return np.nan_to_num(np.var(np.abs(power), axis=1), nan=0.0)
