Source code for protomotions.robot_configs.amp

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from protomotions.robot_configs.base import (
    RobotConfig,
    RobotAssetConfig,
    ControlConfig,
    ControlType,
    SimulatorParams,
)
from protomotions.simulator.isaacgym.config import IsaacGymSimParams
from protomotions.simulator.isaaclab.config import IsaacLabSimParams
from protomotions.simulator.genesis.config import GenesisSimParams
from protomotions.simulator.newton.config import NewtonSimParams
from protomotions.simulator.mujoco.config import MujocoSimParams
from typing import List, Dict, Tuple
from dataclasses import dataclass, field




@dataclass
class AMPRobotConfig(RobotConfig):
    semantic_forward_axis_xy: Tuple[float, float] = (1.0, 0.0)

    common_naming_to_robot_body_names: Dict[str, List[str]] = field(
        default_factory=lambda: {
            "all_right_foot_bodies": ["right_foot"],
            "all_left_foot_bodies": ["left_foot"],
            "all_left_hand_bodies": ["left_hand"],
            "all_right_hand_bodies": ["right_hand"],
            "head_body_name": ["head"],
            "torso_body_name": ["torso"],
        }
    )

    trackable_bodies_subset: List[str] = field(
        default_factory=lambda: [
            "torso",
            "head",
            "left_foot",
            "right_foot",
            "left_hand",
            "right_hand",
        ]
    )

    default_root_height: float = 0.8

    asset: RobotAssetConfig = field(
        default_factory=lambda: RobotAssetConfig(
            asset_file_name="mjcf/amp_humanoid.xml",
        )
    )

    control: ControlConfig = field(
        default_factory=lambda: ControlConfig(control_type=ControlType.BUILT_IN_PD)
    )

    simulation_params: SimulatorParams = field(
        default_factory=lambda: SimulatorParams(
            isaacgym=IsaacGymSimParams(
                fps=60,
                decimation=2,
                substeps=2,
            ),
            isaaclab=IsaacLabSimParams(
                fps=120,
                decimation=4,
            ),
            genesis=GenesisSimParams(
                fps=60,
                decimation=2,
                substeps=2,
            ),
            newton=NewtonSimParams(
                fps=120,
                decimation=4,
            ),
            # Same 30 Hz control rate as the other simulators. implicitfast
            # also treats the PD damping implicitly (Euler only does that for
            # joint damping), which keeps stiff gains on light links stable.
            mujoco=MujocoSimParams(
                fps=960,
                decimation=32,
                integrator="implicitfast",
            ),
        )
    )