Skip to content

Commit e55e200

Browse files
oleflboleflb
authored
Remove Ball Position Control (#2146)
* Add ball body name * Improve position control implementation * Fix broadcasting bug * low_command.py -> position_control.py
1 parent 4916357 commit e55e200

6 files changed

Lines changed: 111 additions & 79 deletions

File tree

‎tools/mujoco-simulator/mujoco-simulator/K1/ball.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
</asset>
2323

2424
<worldbody>
25-
<body pos="2 -1 2.0">
25+
<body pos="2 -1 2.0" name="ball">
2626
<joint
2727
name="ball_free_joint"
2828
type="free"

‎tools/mujoco-simulator/mujoco-simulator/main.py‎

Lines changed: 5 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,9 @@
1010
from rich.logging import RichHandler
1111

1212
from mujoco_simulator.exceptions import UnknownTaskException
13-
from mujoco_simulator.low_command import get_control_input
14-
from mujoco_simulator.low_state import JointActuatorInfo, generate_low_state
13+
from mujoco_simulator.joint_actuator_info import joint_actuator_info_list
14+
from mujoco_simulator.low_state import generate_low_state
15+
from mujoco_simulator.position_control import RobotPositionControl
1516
from mujoco_simulator.rate_logger import SimulationRateLogger
1617
from mujoco_simulator.render import CameraRenderer
1718
from mujoco_simulator.scene import (
@@ -37,34 +38,6 @@ def request_rgbd_sensors(renderer: CameraRenderer, data: MjData) -> RGBDSensors:
3738
)
3839

3940

40-
def joint_actuator_info_list(model: MjModel) -> list:
41-
joints = [
42-
"AAHead_yaw",
43-
"Head_pitch",
44-
"ALeft_Shoulder_Pitch",
45-
"Left_Shoulder_Roll",
46-
"Left_Elbow_Pitch",
47-
"Left_Elbow_Yaw",
48-
"ARight_Shoulder_Pitch",
49-
"Right_Shoulder_Roll",
50-
"Right_Elbow_Pitch",
51-
"Right_Elbow_Yaw",
52-
"Left_Hip_Pitch",
53-
"Left_Hip_Roll",
54-
"Left_Hip_Yaw",
55-
"Left_Knee_Pitch",
56-
"Left_Ankle_Pitch",
57-
"Left_Ankle_Roll",
58-
"Right_Hip_Pitch",
59-
"Right_Hip_Roll",
60-
"Right_Hip_Yaw",
61-
"Right_Knee_Pitch",
62-
"Right_Ankle_Pitch",
63-
"Right_Ankle_Roll",
64-
]
65-
return [JointActuatorInfo(name, model) for name in joints]
66-
67-
6841
async def run_simulation(
6942
server: SimulationServer, model: MjModel, data: MjData
7043
) -> None:
@@ -78,6 +51,7 @@ async def run_simulation(
7851
model=model, camera_name="camera", height=480, width=640
7952
)
8053
actuator_info_list = joint_actuator_info_list(model)
54+
position_control = RobotPositionControl(model, actuator_info_list)
8155

8256
last_tick = time.time()
8357
while True:
@@ -91,7 +65,7 @@ async def run_simulation(
9165
await task.respond(data.time, low_state)
9266
case TaskName.ApplyLowCommand:
9367
if low_command := await task.receive():
94-
data.ctrl[:] = get_control_input(model, data, low_command)
68+
position_control.apply_control(data, low_command)
9569
case TaskName.Reset:
9670
reset_simulation(model, data)
9771
case TaskName.StepSimulation:

‎tools/mujoco-simulator/mujoco-simulator/mujoco_simulator/joint_actuator_info.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
from mujoco import MjModel
2+
3+
4+
class JointActuatorInfo:
5+
name: str
6+
qpos_addr: int
7+
qvel_addr: int
8+
qacc_addr: int
9+
qfrc_actuator_addr: int
10+
11+
def __init__(self, name: str, model: MjModel) -> None:
12+
self.name = name
13+
self.qpos_addr = model.joint(name).qposadr.item()
14+
self.qvel_addr = model.joint(name).dofadr.item()
15+
self.qacc_addr = model.joint(name).dofadr.item()
16+
self.qfrc_actuator_addr = model.actuator(name).id
17+
18+
19+
def joint_actuator_info_list(model: MjModel) -> list[JointActuatorInfo]:
20+
joints = [
21+
"AAHead_yaw",
22+
"Head_pitch",
23+
"ALeft_Shoulder_Pitch",
24+
"Left_Shoulder_Roll",
25+
"Left_Elbow_Pitch",
26+
"Left_Elbow_Yaw",
27+
"ARight_Shoulder_Pitch",
28+
"Right_Shoulder_Roll",
29+
"Right_Elbow_Pitch",
30+
"Right_Elbow_Yaw",
31+
"Left_Hip_Pitch",
32+
"Left_Hip_Roll",
33+
"Left_Hip_Yaw",
34+
"Left_Knee_Pitch",
35+
"Left_Ankle_Pitch",
36+
"Left_Ankle_Roll",
37+
"Right_Hip_Pitch",
38+
"Right_Hip_Roll",
39+
"Right_Hip_Yaw",
40+
"Right_Knee_Pitch",
41+
"Right_Ankle_Pitch",
42+
"Right_Ankle_Roll",
43+
]
44+
return [JointActuatorInfo(name, model) for name in joints]

‎tools/mujoco-simulator/mujoco-simulator/mujoco_simulator/low_command.py‎

Lines changed: 0 additions & 31 deletions
This file was deleted.

‎tools/mujoco-simulator/mujoco-simulator/mujoco_simulator/low_state.py‎

Lines changed: 2 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,8 @@
1-
from mujoco import MjData, MjModel
1+
from mujoco import MjData
22
from mujoco_rust_server.booster_types import ImuState, LowState, MotorState
33

44
from mujoco_simulator._utils import mj_quaternion_to_rpy
5-
6-
7-
class JointActuatorInfo:
8-
name: str
9-
qpos_addr: int
10-
qvel_addr: int
11-
qacc_addr: int
12-
qfrc_actuator_addr: int
13-
14-
def __init__(self, name: str, model: MjModel) -> None:
15-
self.name = name
16-
self.qpos_addr = model.joint(name).qposadr
17-
self.qvel_addr = model.joint(name).dofadr
18-
self.qacc_addr = model.joint(name).dofadr
19-
self.qfrc_actuator_addr = model.actuator(name).id
5+
from mujoco_simulator.joint_actuator_info import JointActuatorInfo
206

217

228
def generate_low_state(

‎tools/mujoco-simulator/mujoco-simulator/mujoco_simulator/position_control.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
import numpy as np
2+
from mujoco import MjData, MjModel
3+
from mujoco_rust_server.booster_types import LowCommand
4+
from numpy.typing import NDArray
5+
6+
from mujoco_simulator.joint_actuator_info import JointActuatorInfo
7+
8+
9+
class RobotPositionControl:
10+
model: MjModel
11+
qpos_indices: NDArray
12+
qvel_indices: NDArray
13+
actuator_indices: NDArray
14+
15+
def __init__(
16+
self, model: MjModel, actuator_info: list[JointActuatorInfo]
17+
) -> None:
18+
self.model = model
19+
self.qpos_indices = np.array(
20+
[info.qpos_addr for info in actuator_info], dtype=np.uint32
21+
)
22+
self.qvel_indices = np.array(
23+
[info.qvel_addr for info in actuator_info], dtype=np.uint32
24+
)
25+
self.actuator_indices = np.array(
26+
[info.qfrc_actuator_addr for info in actuator_info], dtype=np.uint32
27+
)
28+
29+
def apply_control(
30+
self, data: MjData, low_command: LowCommand | None
31+
) -> None:
32+
if low_command is None:
33+
return
34+
35+
q = np.array([cmd.position for cmd in low_command.motor_commands])
36+
dq = np.array([cmd.velocity for cmd in low_command.motor_commands])
37+
tau = np.array([cmd.torque for cmd in low_command.motor_commands])
38+
kp = np.array([cmd.kp for cmd in low_command.motor_commands])
39+
kd = np.array([cmd.kd for cmd in low_command.motor_commands])
40+
weight = np.array([cmd.weight for cmd in low_command.motor_commands])
41+
42+
current_position = data.qpos[self.qpos_indices]
43+
current_velocity = data.qvel[self.qvel_indices]
44+
45+
# TODO(oleflb): booster supposedly clips position to joint limits first
46+
desired = (
47+
kp * (q - current_position) + kd * (dq - current_velocity) + tau
48+
)
49+
50+
ctrl_min = self.model.actuator_ctrlrange[self.actuator_indices, 0]
51+
ctrl_max = self.model.actuator_ctrlrange[self.actuator_indices, 1]
52+
control_torque = np.clip(desired, ctrl_min, ctrl_max)
53+
54+
# Ensure existing control values are used as floats when smoothing
55+
current_ctrl = data.ctrl[self.actuator_indices]
56+
smoothed_control_torque = (
57+
weight * control_torque + (1 - weight) * current_ctrl
58+
)
59+
data.ctrl[self.actuator_indices] = smoothed_control_torque

0 commit comments

Comments
 (0)