Files
walkersim/sim/mujoco/walker_sim/controller.py

63 lines
1.9 KiB
Python

from __future__ import annotations
from dataclasses import dataclass
from math import pi
import numpy as np
@dataclass
class PDGaitConfig:
frequency_hz: float = 0.9
kp: float = 24.0
kd: float = 2.5
near_phase_offset: float = 0.0
far_phase_offset: float = pi
class PDGaitController:
"""Phase-based crank PD controller for the tandem 4-leg linkage."""
JOINTS = (
"near_crank_joint",
"far_crank_joint",
)
def __init__(self, model, data, config: PDGaitConfig | None = None):
self.model = model
self.data = data
self.cfg = config or PDGaitConfig()
self._joint_info = []
for name in self.JOINTS:
jid = model.joint(name).id
actuator_name = "near_crank_motor" if name.startswith("near") else "far_crank_motor"
aid = model.actuator(actuator_name).id
qpos_adr = model.jnt_qposadr[jid]
qvel_adr = model.jnt_dofadr[jid]
ctrl_min, ctrl_max = model.actuator_ctrlrange[aid]
self._joint_info.append((aid, qpos_adr, qvel_adr, ctrl_min, ctrl_max))
def desired_positions(self, t: float) -> np.ndarray:
w = 2.0 * pi * self.cfg.frequency_hz
phase_n = w * t + self.cfg.near_phase_offset
phase_f = phase_n + self.cfg.far_phase_offset
q_des = np.array(
[
phase_n,
phase_f,
],
dtype=np.float64,
)
return q_des
def step(self, t: float) -> None:
q_des = self.desired_positions(t)
for i, (aid, qpos_adr, qvel_adr, ctrl_min, ctrl_max) in enumerate(self._joint_info):
q = self.data.qpos[qpos_adr]
qd = self.data.qvel[qvel_adr]
err = float(q_des[i] - q)
tau = self.cfg.kp * err - self.cfg.kd * qd
self.data.ctrl[aid] = np.clip(tau, ctrl_min, ctrl_max)