Source code for mani_skill.agents.controllers.pd_joint_vel

from dataclasses import dataclass
from typing import Sequence, Union

import numpy as np
from gymnasium import spaces

from mani_skill.utils.structs.types import DriveMode

from .base_controller import BaseController, ControllerConfig


# TODO (stao): add GPU support here
[docs]class PDJointVelController(BaseController):
[docs]    config: "PDJointVelControllerConfig"
[docs]    sets_target_qpos = False
[docs]    sets_target_qvel = True

[docs]    def _initialize_action_space(self):
        n = len(self.joints)
        low = np.float32(np.broadcast_to(self.config.lower, n))
        high = np.float32(np.broadcast_to(self.config.upper, n))
        self.single_action_space = spaces.Box(low, high, dtype=np.float32)

[docs]    def set_drive_property(self):
        n = len(self.joints)
        damping = np.broadcast_to(self.config.damping, n)
        force_limit = np.broadcast_to(self.config.force_limit, n)
        friction = np.broadcast_to(self.config.friction, n)

        for i, joint in enumerate(self.joints):
            drive_mode = self.config.drive_mode
            if not isinstance(drive_mode, str):
                drive_mode = drive_mode[i]
            joint.set_drive_properties(
                0, damping[i], force_limit=force_limit[i], mode=drive_mode
            )
            joint.set_friction(friction[i])

[docs]    def set_action(self, action: np.ndarray):
        action = self._preprocess_action(action)
        self.articulation.set_joint_drive_velocity_targets(
            action, self.joints, self.active_joint_indices
        )


@dataclass
[docs]class PDJointVelControllerConfig(ControllerConfig):
[docs]    lower: Union[float, Sequence[float]]
[docs]    upper: Union[float, Sequence[float]]
[docs]    damping: Union[float, Sequence[float]]
[docs]    force_limit: Union[float, Sequence[float]] = 1e10
[docs]    friction: Union[float, Sequence[float]] = 0.0
[docs]    normalize_action: bool = True
[docs]    drive_mode: Union[Sequence[DriveMode], DriveMode] = "force"
[docs]    controller_cls = PDJointVelController