# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

import math
from re import A, T
from tkinter import N
import torch
from isaaclab.actuators import ImplicitActuatorCfg
import isaaclab.sim as sim_utils
from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg
from isaaclab.envs import ManagerBasedRLEnvCfg
from isaaclab.managers import EventTermCfg as EventTerm
from isaaclab.managers import ObservationGroupCfg as ObsGroup
from isaaclab.managers import ObservationTermCfg as ObsTerm
from isaaclab.managers import RewardTermCfg as RewTerm
from isaaclab.managers import SceneEntityCfg
from isaaclab.managers import TerminationTermCfg as DoneTerm
from isaaclab.managers import CurriculumTermCfg as CurrTerm
from isaaclab.scene import InteractiveSceneCfg
from isaaclab.utils import configclass
from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg
from isaaclab.envs.mdp.actions.actions_cfg import (
    DifferentialInverseKinematicsActionCfg,
    BinaryJointPositionActionCfg,
)
from isaaclab.sensors import CameraCfg

from . import mdp

##
# Pre-defined configs
##

from mindbot.robot.mindbot import MINDBOT_CFG

from mindbot.utils.assets import MINDBOT_ASSETS_DIR

# ====== 其他物体配置 ======
from isaaclab.sim.spawners.from_files.from_files_cfg import UsdFileCfg
from isaaclab.sim.schemas.schemas_cfg import (
    RigidBodyPropertiesCfg,
    CollisionPropertiesCfg,
)


TABLE_CFG = RigidObjectCfg(
    prim_path="{ENV_REGEX_NS}/Table",
    init_state=RigidObjectCfg.InitialStateCfg(
        pos=[0.95, -0.3, 0.005],
        rot=[0.7071, 0.0, 0.0, 0.7071],
        lin_vel=[0.0, 0.0, 0.0],
        ang_vel=[0.0, 0.0, 0.0],
    ),
    spawn=UsdFileCfg(
        usd_path=f"{MINDBOT_ASSETS_DIR}/equipment/table/Table_C.usd",
        rigid_props=RigidBodyPropertiesCfg(
            rigid_body_enabled=True,
            solver_position_iteration_count=32,
            solver_velocity_iteration_count=16,
            disable_gravity=False,
        ),
        collision_props=CollisionPropertiesCfg(
            collision_enabled=True,
            contact_offset=0.0005,  # original 0.02
            rest_offset=0,
        ),
    ),
)

LID_CFG = RigidObjectCfg(
    prim_path="{ENV_REGEX_NS}/Lid",
    init_state=RigidObjectCfg.InitialStateCfg(
        # pos=[0.9523,-0.2512,1.0923],
        pos=[0.803, -0.25, 1.0282],  # (0.9988, -0.2977, 1.0498633)
        # rot=[-0.7071, 0.0, 0.0, 0.7071],
        rot=[0.0, 0.0, 0.0, 1.0],
        lin_vel=[0.0, 0.0, 0.0],
        ang_vel=[0.0, 0.0, 0.0],
    ),
    spawn=UsdFileCfg(
        usd_path=f"{MINDBOT_ASSETS_DIR}/assets/Collected_equipment001/Equipment/tube1.usd",
        # scale=(0.2, 0.2, 0.2),
        rigid_props=RigidBodyPropertiesCfg(
            rigid_body_enabled=True,
            solver_position_iteration_count=32,
            solver_velocity_iteration_count=16,
            max_angular_velocity=1000.0,
            max_linear_velocity=1000.0,
            max_depenetration_velocity=0.5,  # original 5.0
            linear_damping=5.0,  # original 0.5
            angular_damping=5.0,  # original 0.5
            disable_gravity=False,
        ),
        collision_props=CollisionPropertiesCfg(
            collision_enabled=True,
            contact_offset=0.0005,  # original 0.02
            rest_offset=0,
        ),
    ),
)

# LID_CFG = RigidObjectCfg(
#     prim_path="{ENV_REGEX_NS}/Lid",
#     init_state=RigidObjectCfg.InitialStateCfg(
#         # pos=[0.9523,-0.2512,1.0923],
#         pos=[0.8488, -0.2477, 1.0498633],#(0.9988, -0.2977, 1.0498633)
#         # rot=[-0.7071, 0.0, 0.0, 0.7071],
#         rot=[0.0, 0.0, 0.0, 1.0],
#         lin_vel=[0.0, 0.0, 0.0],
#         ang_vel=[0.0, 0.0, 0.0],
#     ),
#     spawn=UsdFileCfg(
#         usd_path="/home/maic/LYT/maic_usd_assets/equipment/centrifuge/Lid_B.usd",
#         # scale=(0.2, 0.2, 0.2),
#         rigid_props=RigidBodyPropertiesCfg(
#             rigid_body_enabled= True,
#             solver_position_iteration_count=32,
#             solver_velocity_iteration_count=16,
#             max_angular_velocity=1000.0,
#             max_linear_velocity=1000.0,
#             max_depenetration_velocity=0.5,#original 5.0
#             linear_damping=5.0, #original 0.5
#             angular_damping=5.0, #original 0.5
#             disable_gravity=False,
#         ),
#         collision_props=CollisionPropertiesCfg(
#             collision_enabled=True,
#             contact_offset=0.0005,#original 0.02
#             rest_offset=0
#         ),
#     ),
# )

CENTRIFUGE_CFG = ArticulationCfg(
    spawn=sim_utils.UsdFileCfg(
        usd_path=f"{MINDBOT_ASSETS_DIR}/equipment/centrifuge/centrifuge.usd",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(
            disable_gravity=False,
            max_depenetration_velocity=1.0,
            linear_damping=0.5,
            angular_damping=0.5,
        ),
        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
            enabled_self_collisions=False,
            solver_position_iteration_count=32,
            solver_velocity_iteration_count=16,
            stabilization_threshold=1e-6,
            # 【重要】必须要固定底座，否则机器人按盖子时，离心机会翻倒
            # fix_root_link=True,
        ),
        collision_props=CollisionPropertiesCfg(
            collision_enabled=True,
            contact_offset=0.0005,  # original 0.02
            rest_offset=0,
        ),
    ),
    init_state=ArticulationCfg.InitialStateCfg(
        # 1. 参照机器人配置，在这里定义初始关节角度
        joint_pos={
            "r1": math.radians(0.0),  # 转子位置（无关紧要）
            # 【关键修改】：初始让盖子处于打开状态
            # 您的 USD 限位是 (-100, 0)，-100度为最大开启
            "r2": math.radians(-100.0),
        },
        pos=(0.80, -0.25, 0.8085),  # (0.95, -0.3,  0.8085)
        rot=[-0.7071, 0.0, 0.0, 0.7071],
        # rot=[0.0, 0.0, 0.0, 1.0],
    ),
    actuators={
        "lid_passive_mechanism": ImplicitActuatorCfg(
            joint_names_expr=["r2"],
            # 【关键差异说明】
            # 机器人的手臂 stiffness 是 10000.0，因为机器人需要精准定位且不晃动。
            #
            # 但是！离心机盖子如果设为 10000，它就会像焊死在空中一样，机器人根本按不动。
            # 这里设为 200.0 左右：
            # 1. 力度足以克服重力，让盖子初始保持打开。
            # 2. 力度又足够软，当机器人用力下压时，盖子会顺从地闭合。
            stiffness=200.0,
            # 阻尼：给一点阻力，模拟液压杆手感，防止像弹簧一样乱晃
            damping=20.0,
            # 确保力矩上限不要太小，否则托不住盖子
            effort_limit_sim=1000.0,
            velocity_limit_sim=100.0,
        ),
        # 转子可以硬一点，不需要被机器人按动
        "rotor_control": ImplicitActuatorCfg(
            joint_names_expr=["r1"],
            stiffness=0.0,
            damping=10.0,
        ),
    },
)


ROOM_CFG = AssetBaseCfg(
    prim_path="{ENV_REGEX_NS}/Room",
    spawn=UsdFileCfg(
        usd_path=f"{MINDBOT_ASSETS_DIR}/twinlab/MIC_sim-3.0/244_140/room.usd",
    ),
)

##
# Scene definition
##


@configclass
class MindbotSceneCfg(InteractiveSceneCfg):
    """Configuration for a cart-pole scene."""

    # ground plane
    ground = AssetBaseCfg(
        prim_path="/World/ground",
        spawn=sim_utils.GroundPlaneCfg(size=(100.0, 100.0)),
    )

    # robot
    # Room: AssetBaseCfg = ROOM_CFG.replace(prim_path="{ENV_REGEX_NS}/Room")
    Mindbot: ArticulationCfg = MINDBOT_CFG.replace(prim_path="{ENV_REGEX_NS}/Mindbot")
    table: RigidObjectCfg = TABLE_CFG.replace(prim_path="{ENV_REGEX_NS}/Table")
    centrifuge: ArticulationCfg = CENTRIFUGE_CFG.replace(
        prim_path="{ENV_REGEX_NS}/centrifuge"
    )
    lid: RigidObjectCfg = LID_CFG.replace(prim_path="{ENV_REGEX_NS}/Lid")
    # lights
    dome_light = AssetBaseCfg(
        prim_path="/World/DomeLight",
        spawn=sim_utils.DomeLightCfg(color=(0.9, 0.9, 0.9), intensity=500.0),
    )
    """
    添加相机定义
    # head_camera=CameraCfg(
    #     prim_path="{ENV_REGEX_NS}/Mindbot/robot_head/cam_head", 
    #     update_period=1/120,
    #     height=480,
    #     width=640,
    #     data_type=["rgb"],
    #     spawn=None,
    # )
    """
    # left_hand_camera=CameraCfg(
    #     prim_path="{ENV_REGEX_NS}/Mindbot/rm_65_fb_left/Link_6/robot___left_8_02/cam_left_hand",
    #     update_period=1/120,
    #     height=480,
    #     width=640,
    #     data_types=["rgb"],
    #     spawn=None,
    # )
    # right_hand_camera=CameraCfg(
    #     prim_path="{ENV_REGEX_NS}/Mindbot/rm_65_b_right/Link6/robot___right_8_2/cam_right_hand",
    #     update_period=1/120,
    #     height=480,
    #     width=640,
    #     data_types=["rgb"],
    #     spawn=None,
    # )
    # head_camera=CameraCfg(
    #     prim_path="{ENV_REGEX_NS}/Mindbot/robot_head/cam_head",
    #     update_period=1/120,
    #     height=480,
    #     width=640,
    #     data_types=["rgb"],
    #     spawn=None,
    # )
    # chest_camera=CameraCfg(
    #     prim_path="{ENV_REGEX_NS}/Mindbot/robot_trunk/cam_chest",
    #     update_period=1/120,
    #     height=480,
    #     width=640,
    #     data_types=["rgb"],
    #     spawn=None,
    # )


##
# MDP settings
##


def left_gripper_pos_w(
    env, asset_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot")
) -> torch.Tensor:
    asset = env.scene[asset_cfg.name]
    body_ids, _ = asset.find_bodies(["left_hand_body"], preserve_order=True)
    pos_w = asset.data.body_pos_w[:, body_ids[0]] - env.scene.env_origins
    return pos_w


@configclass
class ActionsCfg:
    """Action specifications for the MDP."""

    # 使用任务空间控制器：策略输出末端执行器位置+姿态增量（6维）
    # left_arm_ee = DifferentialInverseKinematicsActionCfg(
    #     asset_name="Mindbot",
    #     joint_names=["r_joint[1-6]"],  # 左臂的6个关节
    #     body_name="right_hand_body",     # 末端执行器body名称
    #     controller=DifferentialIKControllerCfg(
    #         command_type="pose",        # 控制位置+姿态
    #         use_relative_mode=True,     # 相对模式：动作是增量
    #         ik_method="dls",            # Damped Least Squares方法
    #     ),
    #     scale=(0.01, 0.01, 0.01, 0.025, 0.025, 0.025),
    # )

    right_arm_fixed = mdp.JointPositionActionCfg(
        asset_name="Mindbot",
        joint_names=["r_joint[1-6]"],  # 对应 l_joint1 到 l_joint6
        # 1. 缩放设为 0：这让 RL 策略无法控制它，它不会动
        scale=0.0,
        # 2. 偏移设为你的目标角度（弧度）：这告诉电机“永远停在这里”
        # 对应 (135, -70, -90, 90, 90, -70)
        offset={
            "r_joint1": 2.3562,
            "r_joint2": -1.2217,
            "r_joint3": -1.5708,
            "r_joint4": 1.5708,
            "r_joint5": 1.5708,
            "r_joint6": -1.2217,
        },
    )

    grippers_position = mdp.BinaryJointPositionActionCfg(
        asset_name="Mindbot",
        joint_names=["left_hand_joint_left", "left_hand_joint_right"],
        # 修正：使用字典格式
        # open_command_expr={"关节名正则": 值}
        open_command_expr={"left_hand_joint_left": 0.0, "left_hand_joint_right": 0.0},
        # close_command_expr={"关节名正则": 值}
        close_command_expr={
            "left_hand_joint_left": -0.03,
            "left_hand_joint_right": 0.03,
        },
    )

    trunk_position = mdp.JointPositionActionCfg(
        asset_name="Mindbot",
        joint_names=["PrismaticJoint"],
        scale=0.00,
        offset=0.24,  # 0.3
        clip=None,
    )

    centrifuge_lid_passive = mdp.JointPositionActionCfg(
        asset_name="centrifuge",  # 对应场景中的名称
        joint_names=["r2"],
        # 将 scale 设为 0，意味着 RL 算法输出的任何值都会被乘 0，即无法干扰它
        scale=0.00,
        # 将 offset 设为目标角度，这会成为 PD 控制器的恒定 Target
        offset=-1.7453,
        clip=None,
    )


@configclass
class ObservationsCfg:
    """Observation specifications for the MDP."""

    @configclass
    class PolicyCfg(ObsGroup):
        """Observations for policy group."""

        mindbot_joint_pos_rel = ObsTerm(
            func=mdp.joint_pos_rel, params={"asset_cfg": SceneEntityCfg("Mindbot")}
        )
        mindbot_root_pos = ObsTerm(
            func=mdp.root_pos_w, params={"asset_cfg": SceneEntityCfg("Mindbot")}
        )
        mindbot_root_quat = ObsTerm(
            func=mdp.root_quat_w, params={"asset_cfg": SceneEntityCfg("Mindbot")}
        )

        left_gripper_pos = ObsTerm(
            func=left_gripper_pos_w,
            params={
                "asset_cfg": SceneEntityCfg("Mindbot", body_names=["left_hand_body"])
            },
        )
        lid_pos_abs = ObsTerm(
            func=mdp.root_pos_w, params={"asset_cfg": SceneEntityCfg("lid")}
        )
        centrifuge_pos_abs = ObsTerm(
            func=mdp.root_pos_w, params={"asset_cfg": SceneEntityCfg("centrifuge")}
        )

        def __post_init__(self) -> None:
            self.enable_corruption = False
            self.concatenate_terms = True

    # observation groups
    policy: PolicyCfg = PolicyCfg()


@configclass
class EventCfg:
    """Configuration for events."""

    # 重置所有关节到 init_state（无偏移）must be
    reset_mindbot_all_joints = EventTerm(
        func=mdp.reset_joints_by_offset,
        mode="reset",
        params={
            "asset_cfg": SceneEntityCfg("Mindbot"),
            "position_range": (0.0, 0.0),
            "velocity_range": (0.0, 0.0),
        },
    )
    # 2. 底座安装误差 (建议缩小到 2cm)
    reset_mindbot_base = EventTerm(
        func=mdp.reset_root_state_uniform,
        mode="reset",
        params={
            "asset_cfg": SceneEntityCfg("Mindbot"),
            "pose_range": {
                "x": (-0.05, 0.05),
                "y": (-0.05, 0.05),
                "z": (0.74, 0.75),
                "yaw": (-0.1, 0.1),
            },
            "velocity_range": {"x": (0.0, 0.0), "y": (0.0, 0.0)},
        },
    )
    # 1. 机械臂关节微小偏移 (0.01 弧度约 0.6 度，很合适)
    reset_mindbot_left_arm = EventTerm(
        func=mdp.reset_joints_by_offset,
        mode="reset",
        params={
            "asset_cfg": SceneEntityCfg("Mindbot", joint_names=["r_joint[1-6]"]),
            "position_range": (-0.01, 0.01),
            "velocity_range": (0.0, 0.0),
        },
    )

    reset_table = EventTerm(
        func=mdp.reset_root_state_uniform,
        mode="reset",
        params={
            "asset_cfg": SceneEntityCfg("table"),
            "pose_range": {"x": (0.0, 0.0)},
            "velocity_range": {"x": (0.0, 0.0)},
        },
    )
    reset_centrifuge = EventTerm(
        func=mdp.reset_root_state_uniform,
        mode="reset",
        params={
            "asset_cfg": SceneEntityCfg("centrifuge"),
            # "pose_range": {"x": (-0.07, 0.07), "y": (-0.07, 0.07), "yaw": (-0.1, 0.1)},
            "pose_range": {"x": (0.0, 0.0), "y": (0.0, 0.0), "z": (0.0, 0.0)},
            "velocity_range": {"x": (0.0, 0.0)},
        },
    )
    reset_lid = EventTerm(
        func=mdp.reset_root_state_uniform,
        mode="reset",
        params={
            "asset_cfg": SceneEntityCfg("lid"),
            # "pose_range": {"x": (-0.03, 0.03), "y": (-0.03, 0.03), "yaw": (-1.5, 1.5)},
            "pose_range": {"x": (0.0, 0.0), "y": (0.0, 0.0), "z": (0.0, 0.0)},
            "velocity_range": {"x": (0.0, 0.0)},
        },
    )


@configclass
class RewardsCfg:
    """Reward terms for the MDP."""

    gripper_lid_orientation_alignment = RewTerm(
        func=mdp.gripper_lid_orientation_alignment,
        params={
            "lid_cfg": SceneEntityCfg("lid"),
            "robot_cfg": SceneEntityCfg("Mindbot"),
            "gripper_body_name": "right_hand_body",
            "gripper_forward_axis": (0.0, 0.0, 1.0),
            "gripper_width_axis": (0.0, 1.0, 0.0),
            "lid_handle_axis": (0.0, 1.0, 0.0),
            "max_angle_deg": 85.0,  # 允许60度内的偏差
        },
        weight=5.0,  # original 5.0
    )

    # stage 2
    # 位置对齐奖励：鼓励夹爪位于 lid 两侧，且 Z 方向在 lid 上方 0.1m
    gripper_lid_position_alignment = RewTerm(
        func=mdp.gripper_lid_position_alignment,
        params={
            "lid_cfg": SceneEntityCfg("lid"),
            "robot_cfg": SceneEntityCfg("Mindbot"),
            "left_gripper_name": "right_hand_l",
            "right_gripper_name": "right_hand__r",
            "height_offset": 0.115,  # Z方向：lid 上方 0.1m
            "std": 0.3,  # 位置对齐的容忍度
        },
        weight=3.0,  # original 3.0
    )
    approach_lid_penalty = RewTerm(
        func=mdp.gripper_lid_distance_penalty,
        params={
            "lid_cfg": SceneEntityCfg("lid"),
            "robot_cfg": SceneEntityCfg("Mindbot"),
            "left_gripper_name": "right_hand_l",
            "right_gripper_name": "right_hand__r",
            "height_offset": 0.115,
        },
        # weight 为负数表示惩罚。
        # 假设距离 0.5m，惩罚就是 -0.5 * 2.0 = -1.0 分。
        weight=-1.0,
    )
    # # 【新增】精细对齐 (引导进入 2cm 圈)
    # gripper_lid_fine_alignment = RewTerm(
    #     func=mdp.gripper_lid_position_alignment,
    #     params={
    #         "lid_cfg": SceneEntityCfg("lid"),
    #         "robot_cfg": SceneEntityCfg("Mindbot"),
    #         "left_gripper_name": "left_hand__l",
    #         "right_gripper_name": "left_hand_r",
    #         "height_offset": 0.07, # 【注意】这个高度必须非常准，和 interaction 判定高度一致
    #         "std": 0.05, # 非常陡峭的梯度，只在 5-10cm 内有效
    #     },
    #     weight=10.0 # 高权重
    # )

    # gripper_close_interaction = RewTerm(
    #     func=mdp.gripper_close_when_near,
    #     params={
    #         "lid_cfg": SceneEntityCfg("lid"),
    #         "robot_cfg": SceneEntityCfg("Mindbot"),
    #         "left_gripper_name": "left_hand__l",
    #         "right_gripper_name": "left_hand_r",
    #         "joint_names": ["left_hand_joint_left", "left_hand_joint_right"],
    #         "target_close_pos": 0.03,
    #         "height_offset": 0.04,
    #         "grasp_radius": 0.03,
    #     },
    #     weight=10.0
    # )

    # lid_lifted_reward = RewTerm(
    #     func=mdp.lid_is_lifted, # 使用 rewards.py 中已有的函数
    #     params={
    #         "lid_cfg": SceneEntityCfg("lid"),
    #         "minimal_height": 1.0, # 根据你的场景调整，比初始高度高 5cm 左右
    #     },
    #     weight=10.0 # 给一个足够大的权重
    # )

    # lid_lifting_reward = RewTerm(
    #     func=mdp.lid_lift_progress_reward,
    #     params={
    #         "lid_cfg": SceneEntityCfg("lid"),
    #         "initial_height": 0.8, # 请务必核实这个值！在 print 里看一下 lid 的 z 坐标
    #         "target_height_lift": 0.2,
    #         "height_offset": 0.07, # 与其他奖励保持一致
    #         "std": 0.1
    #     },
    #     # 权重设大一点，告诉它这是最终目标
    #     weight=10.0
    # )


@configclass
class TerminationsCfg:
    """Termination terms for the MDP."""

    time_out = DoneTerm(func=mdp.time_out, time_out=True)

    mindbot_fly_away = DoneTerm(
        func=mdp.base_height_failure,
        params={"asset_cfg": SceneEntityCfg("Mindbot"), "maximum_height": 2.0},
    )

    lid_fly_away = DoneTerm(
        func=mdp.base_height_failure,
        params={"asset_cfg": SceneEntityCfg("lid"), "maximum_height": 2.0},
    )

    # 新增：盖子掉落判定
    # 如果盖子掉回桌面（高度接近初始高度），且已经运行了一段时间（避开刚开始重置的瞬间）
    lid_dropped = DoneTerm(
        func=mdp.base_height_failure,  # 复用高度判定函数
        params={
            "asset_cfg": SceneEntityCfg("lid"),
            "minimum_height": 0.79,  # 假设初始高度是 0.9，低于 0.88 说明掉下去了或者被砸进去了
        },
    )


##
# Environment configuration
##


@configclass
class CurriculumCfg:
    pass
    # action_rate = CurrTerm(func=mdp.modify_reward_weight,
    #                        params={"term_name": "action_rate", "weight": -1e-1, "num_steps": 10000})

    # joint_vel = CurrTerm(func=mdp.modify_reward_weight,
    #                      params={"term_name": "joint_vel", "weight": -1e-1, "num_steps": 10000})

    # position_std_scheduler = CurrTerm(
    #     func=mdp.modify_term_cfg, # 直接使用 Isaac Lab 内置的类
    #     params={
    #         # 路径：rewards管理器 -> 你的奖励项名 -> params字典 -> std键
    #         "address": "rewards.gripper_lid_position_alignment.params.std",

    #         # 修改逻辑：使用我们刚才写的函数
    #         "modify_fn": mdp.annealing_std,

    #         # 传给函数的参数
    #         "modify_params": {
    #             "start_step": 00, # 约 600 轮
    #             "end_step":   5_000, # 约 1200 轮
    #             "start_std": 0.3,
    #             "end_std": 0.05,
    #         }
    #     }
    # )


@configclass
class MindbotEnvCfg(ManagerBasedRLEnvCfg):
    # Scene settings
    scene: MindbotSceneCfg = MindbotSceneCfg(num_envs=5, env_spacing=5.0)
    # Basic settings
    observations: ObservationsCfg = ObservationsCfg()
    actions: ActionsCfg = ActionsCfg()
    events: EventCfg = EventCfg()
    # MDP settings
    rewards: RewardsCfg = RewardsCfg()
    terminations: TerminationsCfg = TerminationsCfg()
    # curriculum: CurriculumCfg = CurriculumCfg()

    # Post initialization
    def __post_init__(self) -> None:
        """Post initialization."""
        # general settings
        self.decimation = 4  # original 2  # 从 2 增加到 4，控制频率从 25Hz 降到 12.5Hz
        self.episode_length_s = 10
        # viewer settings
        self.viewer.eye = (3.5, 0.0, 13.2)  # (3.5, 0.0, 3.2)
        # simulation settings
        self.sim.dt = 1 / 120  # original 1 / 60
        self.sim.render_interval = self.decimation
        # # 1. 基础堆内存
        self.sim.physx.gpu_heap_capacity = 256 * 1024 * 1024  # 256MB
        self.sim.physx.gpu_collision_stack_size = 1024 * 1024 * 1024
        self.sim.physx.gpu_temp_buffer_capacity = 128 * 1024 * 1024
        self.sim.physx.gpu_max_rigid_contact_count = 16 * 1024 * 1024  # 200万个接触点
        self.sim.physx.gpu_max_rigid_patch_count = 1024 * 1024  # 32万个 patch

        # # 5. 聚合对容量 (针对复杂的 Articulation)
        self.sim.physx.gpu_found_lost_pairs_capacity = 16 * 1024 * 1024
        self.sim.physx.gpu_found_lost_aggregate_pairs_capacity = 16 * 1024 * 1024
        self.sim.physx.gpu_total_aggregate_pairs_capacity = 16 * 1024 * 1024
