trossen_arm_sim 0.2.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- trossen_arm_sim/__init__.py +11 -0
- trossen_arm_sim/__main__.py +70 -0
- trossen_arm_sim/assets/base_link.stl +0 -0
- trossen_arm_sim/assets/camera_mount_d405.stl +0 -0
- trossen_arm_sim/assets/carriage_left.stl +0 -0
- trossen_arm_sim/assets/carriage_right.stl +0 -0
- trossen_arm_sim/assets/d405.stl +0 -0
- trossen_arm_sim/assets/gripper_left.stl +0 -0
- trossen_arm_sim/assets/gripper_right.stl +0 -0
- trossen_arm_sim/assets/link_1.stl +0 -0
- trossen_arm_sim/assets/link_2.stl +0 -0
- trossen_arm_sim/assets/link_3.stl +0 -0
- trossen_arm_sim/assets/link_4.stl +0 -0
- trossen_arm_sim/assets/link_5.stl +0 -0
- trossen_arm_sim/assets/link_6.stl +0 -0
- trossen_arm_sim/assets/scene_empty.xml +21 -0
- trossen_arm_sim/assets/wxai_follower.xml +118 -0
- trossen_arm_sim/collision.py +208 -0
- trossen_arm_sim/collision_server.py +74 -0
- trossen_arm_sim/protocol.py +165 -0
- trossen_arm_sim/server.py +240 -0
- trossen_arm_sim/sim.py +170 -0
- trossen_arm_sim-0.2.0.dist-info/METADATA +142 -0
- trossen_arm_sim-0.2.0.dist-info/RECORD +28 -0
- trossen_arm_sim-0.2.0.dist-info/WHEEL +5 -0
- trossen_arm_sim-0.2.0.dist-info/entry_points.txt +2 -0
- trossen_arm_sim-0.2.0.dist-info/licenses/LICENSE +201 -0
- trossen_arm_sim-0.2.0.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
"""Simulated Trossen arm controller: the trossen_arm driver connects to this
|
|
2
|
+
server exactly as it would to real hardware, and the arm runs in MuJoCo."""
|
|
3
|
+
|
|
4
|
+
from .collision import CollisionEvent, CollisionMonitor, Contact
|
|
5
|
+
from .collision_server import CollisionWebSocketServer
|
|
6
|
+
from .server import TrossenArmSimServer
|
|
7
|
+
from .sim import ARM_MODEL_PATH, DEFAULT_SCENE, WidowXSim, load_scene
|
|
8
|
+
|
|
9
|
+
__all__ = ["ARM_MODEL_PATH", "DEFAULT_SCENE", "CollisionEvent",
|
|
10
|
+
"CollisionMonitor", "CollisionWebSocketServer", "Contact",
|
|
11
|
+
"TrossenArmSimServer", "WidowXSim", "load_scene"]
|
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
"""Run the simulated Trossen arm controller.
|
|
2
|
+
|
|
3
|
+
Usage:
|
|
4
|
+
python -m trossen_arm_sim [SCENE] [--viewer] [--verbose] [--ws-port PORT]
|
|
5
|
+
|
|
6
|
+
SCENE is a MuJoCo MJCF file that includes the packaged arm model via
|
|
7
|
+
`<include file="WIDOWX_XML"/>`; the WIDOWX_XML and TROSSEN_ARM_ASSET_DIR
|
|
8
|
+
placeholders are replaced with the installed asset paths at load time. It
|
|
9
|
+
defaults to a built-in empty scene with just the arm on a ground plane.
|
|
10
|
+
With --viewer, a MuJoCo viewer window shows the scene (on macOS this
|
|
11
|
+
requires running under mjpython).
|
|
12
|
+
|
|
13
|
+
Contacts in the scene are streamed as JSON over a WebSocket on --ws-port;
|
|
14
|
+
see collision_server.py for the messages.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
import argparse
|
|
18
|
+
import logging
|
|
19
|
+
import time
|
|
20
|
+
|
|
21
|
+
from .collision import CollisionMonitor
|
|
22
|
+
from .collision_server import WS_PORT, CollisionWebSocketServer
|
|
23
|
+
from .server import TrossenArmSimServer
|
|
24
|
+
from .sim import DEFAULT_SCENE, WidowXSim
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def main() -> None:
|
|
28
|
+
parser = argparse.ArgumentParser(description=__doc__)
|
|
29
|
+
parser.add_argument("scene", nargs="?", default=DEFAULT_SCENE,
|
|
30
|
+
help="MuJoCo scene XML including the arm model "
|
|
31
|
+
"(default: built-in empty scene)")
|
|
32
|
+
parser.add_argument("--host", default="0.0.0.0", help="bind address")
|
|
33
|
+
parser.add_argument("--viewer", action="store_true",
|
|
34
|
+
help="show the MuJoCo viewer (macOS: run with mjpython)")
|
|
35
|
+
parser.add_argument("--verbose", action="store_true",
|
|
36
|
+
help="log every protocol request")
|
|
37
|
+
parser.add_argument("--ws-port", type=int, default=WS_PORT,
|
|
38
|
+
help="WebSocket port for the collision stream")
|
|
39
|
+
args = parser.parse_args()
|
|
40
|
+
|
|
41
|
+
logging.basicConfig(
|
|
42
|
+
level=logging.DEBUG if args.verbose else logging.INFO,
|
|
43
|
+
format="%(asctime)s %(name)s %(levelname)s: %(message)s",
|
|
44
|
+
)
|
|
45
|
+
|
|
46
|
+
sim = WidowXSim(args.scene)
|
|
47
|
+
monitor = CollisionMonitor(sim)
|
|
48
|
+
sim.start()
|
|
49
|
+
server = TrossenArmSimServer(sim, host=args.host)
|
|
50
|
+
server.start()
|
|
51
|
+
collision_server = CollisionWebSocketServer(monitor, host=args.host,
|
|
52
|
+
port=args.ws_port)
|
|
53
|
+
collision_server.start()
|
|
54
|
+
|
|
55
|
+
try:
|
|
56
|
+
if args.viewer:
|
|
57
|
+
sim.sync_viewer_loop()
|
|
58
|
+
else:
|
|
59
|
+
while True:
|
|
60
|
+
time.sleep(1)
|
|
61
|
+
except KeyboardInterrupt:
|
|
62
|
+
pass
|
|
63
|
+
finally:
|
|
64
|
+
collision_server.stop()
|
|
65
|
+
server.stop()
|
|
66
|
+
sim.stop()
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
if __name__ == "__main__":
|
|
70
|
+
main()
|
|
Binary file
|
|
Binary file
|
|
Binary file
|
|
Binary file
|
|
Binary file
|
|
Binary file
|
|
Binary file
|
|
Binary file
|
|
Binary file
|
|
Binary file
|
|
Binary file
|
|
Binary file
|
|
Binary file
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
<mujoco model="wxai empty scene">
|
|
2
|
+
<include file="wxai_follower.xml"/>
|
|
3
|
+
|
|
4
|
+
<statistic center="0.2 0 0.2" extent="0.6"/>
|
|
5
|
+
|
|
6
|
+
<visual>
|
|
7
|
+
<headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
|
|
8
|
+
<global azimuth="160" elevation="-20"/>
|
|
9
|
+
</visual>
|
|
10
|
+
|
|
11
|
+
<asset>
|
|
12
|
+
<texture type="2d" name="groundplane" builtin="checker" mark="edge"
|
|
13
|
+
rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
|
|
14
|
+
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
|
|
15
|
+
</asset>
|
|
16
|
+
|
|
17
|
+
<worldbody>
|
|
18
|
+
<light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
|
|
19
|
+
<geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
|
|
20
|
+
</worldbody>
|
|
21
|
+
</mujoco>
|
|
@@ -0,0 +1,118 @@
|
|
|
1
|
+
<mujoco model="wxai_follower">
|
|
2
|
+
<compiler angle="radian" autolimits="true"/>
|
|
3
|
+
|
|
4
|
+
<option integrator="implicitfast"/>
|
|
5
|
+
|
|
6
|
+
<default>
|
|
7
|
+
<default class="visual">
|
|
8
|
+
<geom type="mesh" group="2" contype="0" conaffinity="0" material="trossen_black"/>
|
|
9
|
+
</default>
|
|
10
|
+
<default class="collision">
|
|
11
|
+
<geom type="cylinder" group="3" mass="0" density="0"/>
|
|
12
|
+
</default>
|
|
13
|
+
</default>
|
|
14
|
+
|
|
15
|
+
<asset>
|
|
16
|
+
<material name="trossen_black" rgba="0.1 0.1 0.1 1"/>
|
|
17
|
+
|
|
18
|
+
<mesh name="base_link" file="base_link.stl"/>
|
|
19
|
+
<mesh name="link_1" file="link_1.stl"/>
|
|
20
|
+
<mesh name="link_2" file="link_2.stl"/>
|
|
21
|
+
<mesh name="link_3" file="link_3.stl"/>
|
|
22
|
+
<mesh name="link_4" file="link_4.stl"/>
|
|
23
|
+
<mesh name="link_5" file="link_5.stl"/>
|
|
24
|
+
<mesh name="link_6" file="link_6.stl"/>
|
|
25
|
+
<mesh name="carriage_right" file="carriage_right.stl"/>
|
|
26
|
+
<mesh name="carriage_left" file="carriage_left.stl"/>
|
|
27
|
+
<mesh name="camera_mount_d405" file="camera_mount_d405.stl"/>
|
|
28
|
+
<mesh name="d405" file="d405.stl" scale="0.001 0.001 0.001"/>
|
|
29
|
+
<mesh name="gripper_left" file="gripper_left.stl"/>
|
|
30
|
+
<mesh name="gripper_right" file="gripper_right.stl"/>
|
|
31
|
+
</asset>
|
|
32
|
+
|
|
33
|
+
<worldbody>
|
|
34
|
+
<body name="base_link" pos="0 0 0" gravcomp="1">
|
|
35
|
+
<inertial pos="-0.00014175 -5.228e-05 0.03175177" quat="0.503748 -0.507380 -0.494230 -0.494510" mass="0.45969858" diaginertia="0.000213926 0.000416701 0.000448963"/>
|
|
36
|
+
<geom class="visual" mesh="base_link"/>
|
|
37
|
+
<geom class="collision" type="cylinder" size="0.031 0.0325" pos="0 0 0.0325" quat="1 0 0 0"/>
|
|
38
|
+
<body name="link_1" pos="0 0 0.05725" gravcomp="1">
|
|
39
|
+
<inertial pos="-0.00011075 0.00171176 0.0204459" quat="0.313955 0.305238 -0.643575 0.627752" mass="0.152704" diaginertia="0.000203432 0.000173907 8.3521e-05"/>
|
|
40
|
+
<joint name="joint_0" pos="0 0 0" axis="0 0 1" range="-3.05433 3.05433" armature="0.032" frictionloss="0.1" actuatorgravcomp="true"/>
|
|
41
|
+
<geom class="visual" mesh="link_1"/>
|
|
42
|
+
<body name="link_2" pos="0.02 0 0.04625" gravcomp="1">
|
|
43
|
+
<inertial pos="-0.131215 -0.00292583 0.00021345" quat="-0.00123582 0.706546 -0.00109575 0.707665" mass="1.15316" diaginertia="0.0179449 0.0173612 0.00102628"/>
|
|
44
|
+
<joint name="joint_1" pos="0 0 0" axis="0 1 0" range="0 3.14159" armature="0.032" frictionloss="0.1" actuatorgravcomp="true"/>
|
|
45
|
+
<geom class="visual" mesh="link_2"/>
|
|
46
|
+
<geom class="collision" type="capsule" size="0.031 0.015" pos="0 0 0" quat="1 1 0 0"/>
|
|
47
|
+
<geom class="collision" type="box" size="0.0975 0.02 0.02" pos="-0.131 0 0" quat="1 0 0 0"/>
|
|
48
|
+
<body name="link_3" pos="-0.264 0 0" gravcomp="1">
|
|
49
|
+
<inertial pos="0.180836 -0.0009409 0.0555494" quat="0.0103901 0.669684 -0.0033175 0.742566" mass="0.686666" diaginertia="0.00566898 0.00547578 0.00059316"/>
|
|
50
|
+
<joint name="joint_2" pos="0 0 0" axis="0 -1 0" range="0 2.35619" armature="0.032" frictionloss="0.1" actuatorgravcomp="true"/>
|
|
51
|
+
<geom class="visual" mesh="link_3"/>
|
|
52
|
+
<geom class="collision" type="capsule" size="0.031 0.015" pos="0 0 0" quat="1 1 0 0"/>
|
|
53
|
+
<geom class="collision" type="box" size="0.08 0.02 0.02" pos="0.13 0 0.06" quat="1 0 0 0"/>
|
|
54
|
+
<body name="link_4" pos="0.245 0 0.06" gravcomp="1">
|
|
55
|
+
<inertial pos="0.0579784 0.00027145 0.0588445" quat="0.948688 0.0346237 0.307828 0.0635164" mass="0.457899" diaginertia="0.000670827 0.000661455 0.000288648"/>
|
|
56
|
+
<joint name="joint_3" pos="0 0 0" axis="0 -1 0" range="-1.5708 1.5708" armature="0.0018" frictionloss="0.1" actuatorgravcomp="true"/>
|
|
57
|
+
<geom class="visual" mesh="link_4"/>
|
|
58
|
+
<geom class="collision" type="capsule" size="0.031 0.015" pos="0 0 0" quat="1 1 0 0"/>
|
|
59
|
+
<body name="link_5" pos="0.06775 0 0.0455" gravcomp="1">
|
|
60
|
+
<inertial pos="0.00412447 -1.138e-05 -0.0428318" quat="0.496144 0.502128 0.503595 0.498098" mass="0.366939" diaginertia="0.000320622 0.00026751 0.000189108"/>
|
|
61
|
+
<joint name="joint_4" pos="0 0 0" axis="0 0 -1" range="-1.5708 1.5708" armature="0.0018" frictionloss="0.1" actuatorgravcomp="true"/>
|
|
62
|
+
<geom class="visual" mesh="link_5"/>
|
|
63
|
+
<geom class="collision" type="cylinder" size="0.031 0.028" pos="0 0 0.03" quat="1 0 0 0"/>
|
|
64
|
+
<body name="link_6" pos="0.02895 0 -0.0455" gravcomp="1">
|
|
65
|
+
<inertial pos="0.046841 0.00145133 0.00895674" quat="0.275854 0.291255 -0.691957 0.600225" mass="0.671874" diaginertia="0.00118316 0.00110863 0.00078309"/>
|
|
66
|
+
<joint name="joint_5" pos="0 0 0" axis="1 0 0" range="-3.14159 3.14159" armature="0.0018" frictionloss="0.1" actuatorgravcomp="true"/>
|
|
67
|
+
<geom class="visual" mesh="link_6"/>
|
|
68
|
+
<geom class="collision" type="cylinder" size="0.029 0.05" pos="0 0 0" quat="1 0 1 0"/>
|
|
69
|
+
<geom class="collision" type="box" size="0.0125 0.094 0.031" pos="0.0685 0 0" quat="1 0 0 0"/>
|
|
70
|
+
<!-- End-effector site at gripper tip -->
|
|
71
|
+
<site name="ee_site" pos="0.156 0 0" size="0.001"/>
|
|
72
|
+
<geom class="visual" pos="0.012 0 0" quat="1 0 0 0" mesh="camera_mount_d405"/>
|
|
73
|
+
<geom class="collision" type="box" size="0.0155 0.025 0.025" pos="0.0513206 0 0.0569399" quat="0.984808 0 0.173648 0"/>
|
|
74
|
+
<geom class="visual" type="box" size="0.0115 0.021 0.021" pos="0.0513206 0 0.0569399" quat="0.984808 0 0.173648 0"/>
|
|
75
|
+
<site name="camera_color_frame" pos="0.05864 0.009 0.05427" quat="0.9848078 0 0.1736482 0" size="0.005"/>
|
|
76
|
+
<camera name="cam" pos="0.05864 0.009 0.05427" quat="0.5793 0.4056 -0.4056 -0.5793" mode="fixed" fovy="87"/>
|
|
77
|
+
<body name="carriage_right" pos="0.0865 -0.023 0" gravcomp="1">
|
|
78
|
+
<inertial pos="0.00169015 0.00592793 0.00201817" quat="-0.370973 0.740886 -0.0847721 0.553427" mass="0.081271" diaginertia="4.29565e-05 3.87081e-05 2.21528e-05"/>
|
|
79
|
+
<joint name="right_carriage_joint" pos="0 0 0" axis="0 -1 0" type="slide" range="0 0.044" armature="0.1" frictionloss="0.1" actuatorgravcomp="true"/>
|
|
80
|
+
<geom class="visual" mesh="carriage_right"/>
|
|
81
|
+
<geom class="visual" mesh="gripper_right"/>
|
|
82
|
+
<geom class="collision" type="box" size="0.0275 0.005 0.01" pos="0.0286 0.01436 0.0037" quat="0.99257 -0.00451 0.10246 0.06547"/>
|
|
83
|
+
<geom class="collision" type="box" size="0.0275 0.005 0.01" pos="0.0285 0.01429 -0.00973" quat="0.98344 0.0133 -0.16893 0.06427"/>
|
|
84
|
+
<geom class="collision" type="box" size="0.00732 0.00214 0.00767" pos="0.0622 0.02083 -0.00105" quat="1 0 0 0"/>
|
|
85
|
+
</body>
|
|
86
|
+
<body name="carriage_left" pos="0.0865 0.023 0" gravcomp="1">
|
|
87
|
+
<inertial pos="0.00169017 -0.00592796 -0.00365701" quat="-0.0382691 0.628022 -0.255101 0.734198" mass="0.081271" diaginertia="4.21112e-05 3.61943e-05 2.40003e-05"/>
|
|
88
|
+
<joint name="left_carriage_joint" pos="0 0 0" axis="0 1 0" type="slide" range="0 0.044" armature="0.1" frictionloss="0.1" actuatorgravcomp="true"/>
|
|
89
|
+
<geom class="visual" mesh="carriage_left"/>
|
|
90
|
+
<geom class="visual" mesh="gripper_left"/>
|
|
91
|
+
<geom class="collision" type="box" size="0.0275 0.005 0.01" pos="0.0286 -0.01436 0.0037" quat="0.99257 0.00451 0.10246 -0.06547"/>
|
|
92
|
+
<geom class="collision" type="box" size="0.0275 0.005 0.01" pos="0.0285 -0.01429 -0.00973" quat="0.98344 -0.0133 -0.16893 -0.06427"/>
|
|
93
|
+
<geom class="collision" type="box" size="0.00732 0.00214 0.00767" pos="0.0622 -0.02083 -0.00105" quat="1 0 0 0"/>
|
|
94
|
+
</body>
|
|
95
|
+
</body>
|
|
96
|
+
</body>
|
|
97
|
+
</body>
|
|
98
|
+
</body>
|
|
99
|
+
</body>
|
|
100
|
+
</body>
|
|
101
|
+
</body>
|
|
102
|
+
</worldbody>
|
|
103
|
+
|
|
104
|
+
<equality>
|
|
105
|
+
<joint joint1="right_carriage_joint" joint2="left_carriage_joint" polycoef="0 1 0 0 0"/>
|
|
106
|
+
</equality>
|
|
107
|
+
|
|
108
|
+
<actuator>
|
|
109
|
+
<position name="joint_0" joint="joint_0" kp="200" kv="10" ctrlrange="-3.05433 3.05433" forcerange="-27 27"/>
|
|
110
|
+
<position name="joint_1" joint="joint_1" kp="200" kv="10" ctrlrange="0 3.14159" forcerange="-27 27"/>
|
|
111
|
+
<position name="joint_2" joint="joint_2" kp="200" kv="10" ctrlrange="0 2.35619" forcerange="-27 27"/>
|
|
112
|
+
<position name="joint_3" joint="joint_3" kp="100" kv="5" ctrlrange="-1.5708 1.5708" forcerange="-7 7"/>
|
|
113
|
+
<position name="joint_4" joint="joint_4" kp="50" kv="5" ctrlrange="-1.5708 1.5708" forcerange="-7 7"/>
|
|
114
|
+
<position name="joint_5" joint="joint_5" kp="50" kv="5" ctrlrange="-3.14159 3.14159" forcerange="-7 7"/>
|
|
115
|
+
<position name="left_gripper" joint="left_carriage_joint" kp="1000" kv="50" ctrlrange="0 0.044" forcerange="-400 400"/>
|
|
116
|
+
</actuator>
|
|
117
|
+
|
|
118
|
+
</mujoco>
|
|
@@ -0,0 +1,208 @@
|
|
|
1
|
+
"""Collision detection for the simulation.
|
|
2
|
+
|
|
3
|
+
MuJoCo reports every contact it resolves in ``data.contact`` but attaches no
|
|
4
|
+
meaning to them: a fingertip pressing a key and an elbow hitting the floor
|
|
5
|
+
look the same. This module reads the contacts after each physics step and
|
|
6
|
+
aggregates them per pair of touching geoms, so that a client can decide
|
|
7
|
+
which pairs matter by name and by force.
|
|
8
|
+
|
|
9
|
+
Pairs are tracked over time so that subscribers receive one ``begin`` event
|
|
10
|
+
when a contact starts and one ``end`` event, with duration and peak normal
|
|
11
|
+
force, when it has been absent for ``release_time`` seconds.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
import logging
|
|
15
|
+
import queue
|
|
16
|
+
import threading
|
|
17
|
+
from dataclasses import dataclass
|
|
18
|
+
|
|
19
|
+
import mujoco
|
|
20
|
+
import numpy as np
|
|
21
|
+
|
|
22
|
+
from .sim import WidowXSim
|
|
23
|
+
|
|
24
|
+
logger = logging.getLogger("trossen_arm_sim.collision")
|
|
25
|
+
|
|
26
|
+
ARM_ROOT_BODY = "base_link"
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
@dataclass
|
|
30
|
+
class Contact:
|
|
31
|
+
"""Aggregate of all contact points between two geoms.
|
|
32
|
+
|
|
33
|
+
An unnamed geom is labelled by its body's name, an unnamed body by
|
|
34
|
+
"body<id>", and a deformable flex by its name or "flex<id>". When the
|
|
35
|
+
arm is involved, its side is reported first.
|
|
36
|
+
"""
|
|
37
|
+
geom1: str
|
|
38
|
+
body1: str
|
|
39
|
+
geom2: str
|
|
40
|
+
body2: str
|
|
41
|
+
force: float # total normal force over the contact points [N]
|
|
42
|
+
pos: tuple[float, float, float] # mean contact position [m]
|
|
43
|
+
|
|
44
|
+
def to_dict(self) -> dict:
|
|
45
|
+
return {
|
|
46
|
+
"geom1": self.geom1,
|
|
47
|
+
"body1": self.body1,
|
|
48
|
+
"geom2": self.geom2,
|
|
49
|
+
"body2": self.body2,
|
|
50
|
+
"force": round(self.force, 4),
|
|
51
|
+
"pos": [round(x, 4) for x in self.pos],
|
|
52
|
+
}
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
@dataclass
|
|
56
|
+
class CollisionEvent:
|
|
57
|
+
event: str # "begin" or "end"
|
|
58
|
+
time: float # simulation time [s]
|
|
59
|
+
contact: Contact
|
|
60
|
+
duration: float # time the pair has been in contact [s]
|
|
61
|
+
peak_force: float # largest normal force seen over that time [N]
|
|
62
|
+
|
|
63
|
+
def to_dict(self) -> dict:
|
|
64
|
+
return {
|
|
65
|
+
"type": self.event,
|
|
66
|
+
"time": round(self.time, 4),
|
|
67
|
+
**self.contact.to_dict(),
|
|
68
|
+
"duration": round(self.duration, 4),
|
|
69
|
+
"peak_force": round(self.peak_force, 4),
|
|
70
|
+
}
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
@dataclass
|
|
74
|
+
class _TrackedContact:
|
|
75
|
+
contact: Contact
|
|
76
|
+
start_time: float
|
|
77
|
+
last_seen: float
|
|
78
|
+
peak_force: float
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
class CollisionMonitor:
|
|
82
|
+
"""Tracks contacts in the scene and publishes begin/end events.
|
|
83
|
+
|
|
84
|
+
Installs itself as a step hook on the simulation. Subscribers get a
|
|
85
|
+
queue of CollisionEvent; active contacts can also be polled.
|
|
86
|
+
"""
|
|
87
|
+
|
|
88
|
+
def __init__(self, sim: WidowXSim, release_time: float = 0.02):
|
|
89
|
+
self.sim = sim
|
|
90
|
+
self.release_time = release_time
|
|
91
|
+
model = sim.model
|
|
92
|
+
|
|
93
|
+
root = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, ARM_ROOT_BODY)
|
|
94
|
+
assert root >= 0, f"body '{ARM_ROOT_BODY}' not found in scene"
|
|
95
|
+
self._is_robot_body = np.zeros(model.nbody, dtype=bool)
|
|
96
|
+
for body in range(model.nbody):
|
|
97
|
+
ancestor = body
|
|
98
|
+
while ancestor > 0 and ancestor != root:
|
|
99
|
+
ancestor = model.body_parentid[ancestor]
|
|
100
|
+
self._is_robot_body[body] = ancestor == root
|
|
101
|
+
|
|
102
|
+
self._body_labels = [model.body(i).name or f"body{i}"
|
|
103
|
+
for i in range(model.nbody)]
|
|
104
|
+
self._geom_labels = [
|
|
105
|
+
model.geom(i).name or self._body_labels[model.geom_bodyid[i]]
|
|
106
|
+
for i in range(model.ngeom)
|
|
107
|
+
]
|
|
108
|
+
self._flex_labels = [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_FLEX, i) or f"flex{i}"
|
|
109
|
+
for i in range(model.nflex)]
|
|
110
|
+
|
|
111
|
+
self._lock = threading.Lock()
|
|
112
|
+
self._time = 0.0
|
|
113
|
+
self._tracked: dict[tuple[str, str], _TrackedContact] = {}
|
|
114
|
+
self._subscribers: list[queue.Queue] = []
|
|
115
|
+
self._force = np.zeros(6)
|
|
116
|
+
sim.add_step_hook(self._on_step)
|
|
117
|
+
|
|
118
|
+
@property
|
|
119
|
+
def time(self) -> float:
|
|
120
|
+
"""Simulation time of the last processed step [s]."""
|
|
121
|
+
with self._lock:
|
|
122
|
+
return self._time
|
|
123
|
+
|
|
124
|
+
def active_contacts(self) -> list[Contact]:
|
|
125
|
+
with self._lock:
|
|
126
|
+
return [tracked.contact for tracked in self._tracked.values()]
|
|
127
|
+
|
|
128
|
+
def subscribe(self) -> tuple[list[Contact], queue.Queue]:
|
|
129
|
+
"""Start receiving events; returns the active contacts at that moment
|
|
130
|
+
and the queue the following events will be delivered to."""
|
|
131
|
+
with self._lock:
|
|
132
|
+
events: queue.Queue = queue.Queue()
|
|
133
|
+
self._subscribers.append(events)
|
|
134
|
+
return [tracked.contact for tracked in self._tracked.values()], events
|
|
135
|
+
|
|
136
|
+
def unsubscribe(self, events: queue.Queue) -> None:
|
|
137
|
+
with self._lock:
|
|
138
|
+
self._subscribers.remove(events)
|
|
139
|
+
|
|
140
|
+
# ------------------------------------------------------------ step hook
|
|
141
|
+
|
|
142
|
+
def _on_step(self) -> None:
|
|
143
|
+
"""Runs on the stepping thread with the simulation lock held."""
|
|
144
|
+
data = self.sim.data
|
|
145
|
+
current = self._read_contacts() if data.ncon > 0 else {}
|
|
146
|
+
with self._lock:
|
|
147
|
+
self._time = data.time
|
|
148
|
+
if not current and not self._tracked:
|
|
149
|
+
return
|
|
150
|
+
for key, contact in current.items():
|
|
151
|
+
tracked = self._tracked.get(key)
|
|
152
|
+
if tracked is None:
|
|
153
|
+
self._tracked[key] = _TrackedContact(
|
|
154
|
+
contact, data.time, data.time, contact.force)
|
|
155
|
+
self._publish("begin", contact, 0.0, contact.force)
|
|
156
|
+
else:
|
|
157
|
+
tracked.contact = contact
|
|
158
|
+
tracked.last_seen = data.time
|
|
159
|
+
tracked.peak_force = max(tracked.peak_force, contact.force)
|
|
160
|
+
for key in [k for k, t in self._tracked.items()
|
|
161
|
+
if k not in current
|
|
162
|
+
and data.time - t.last_seen >= self.release_time]:
|
|
163
|
+
tracked = self._tracked.pop(key)
|
|
164
|
+
self._publish("end", tracked.contact,
|
|
165
|
+
tracked.last_seen - tracked.start_time,
|
|
166
|
+
tracked.peak_force)
|
|
167
|
+
|
|
168
|
+
def _read_contacts(self) -> dict[tuple[str, str], Contact]:
|
|
169
|
+
"""Aggregate this step's contact points per geom pair."""
|
|
170
|
+
model, data = self.sim.model, self.sim.data
|
|
171
|
+
ncon = data.ncon
|
|
172
|
+
geoms = data.contact.geom[:ncon]
|
|
173
|
+
flexes = data.contact.flex[:ncon]
|
|
174
|
+
# A flex side has geom -1; it is never part of the arm.
|
|
175
|
+
is_robot = (geoms >= 0) & self._is_robot_body[model.geom_bodyid[geoms]]
|
|
176
|
+
positions = data.contact.pos[:ncon]
|
|
177
|
+
|
|
178
|
+
sums: dict[tuple[str, str], list] = {}
|
|
179
|
+
for i in range(ncon):
|
|
180
|
+
mujoco.mj_contactForce(model, data, i, self._force)
|
|
181
|
+
# Arm side first; otherwise keep MuJoCo's order.
|
|
182
|
+
first = 1 if is_robot[i, 1] and not is_robot[i, 0] else 0
|
|
183
|
+
geom1, body1 = self._labels(geoms[i, first], flexes[i, first])
|
|
184
|
+
geom2, body2 = self._labels(geoms[i, 1 - first], flexes[i, 1 - first])
|
|
185
|
+
entry = sums.setdefault((geom1, geom2), [body1, body2, 0.0, np.zeros(3), 0])
|
|
186
|
+
entry[2] += self._force[0]
|
|
187
|
+
entry[3] += positions[i]
|
|
188
|
+
entry[4] += 1
|
|
189
|
+
return {
|
|
190
|
+
key: Contact(key[0], body1, key[1], body2, float(force),
|
|
191
|
+
tuple(map(float, pos_sum / count)))
|
|
192
|
+
for key, (body1, body2, force, pos_sum, count) in sums.items()
|
|
193
|
+
}
|
|
194
|
+
|
|
195
|
+
def _labels(self, geom: int, flex: int) -> tuple[str, str]:
|
|
196
|
+
"""(geom label, body label) of one side of a contact."""
|
|
197
|
+
if geom >= 0:
|
|
198
|
+
return self._geom_labels[geom], self._body_labels[self.sim.model.geom_bodyid[geom]]
|
|
199
|
+
return self._flex_labels[flex], self._flex_labels[flex]
|
|
200
|
+
|
|
201
|
+
def _publish(self, event: str, contact: Contact, duration: float,
|
|
202
|
+
peak_force: float) -> None:
|
|
203
|
+
"""Called with self._lock held."""
|
|
204
|
+
logger.debug("contact %s: %s <-> %s (%.2f N peak, %.2f s)", event,
|
|
205
|
+
contact.geom1, contact.geom2, peak_force, duration)
|
|
206
|
+
collision = CollisionEvent(event, self._time, contact, duration, peak_force)
|
|
207
|
+
for events in self._subscribers:
|
|
208
|
+
events.put(collision)
|
|
@@ -0,0 +1,74 @@
|
|
|
1
|
+
"""WebSocket server streaming collision events as JSON."""
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import logging
|
|
5
|
+
import queue
|
|
6
|
+
import threading
|
|
7
|
+
|
|
8
|
+
from websockets.exceptions import ConnectionClosed
|
|
9
|
+
from websockets.sync.server import Server, ServerConnection, serve
|
|
10
|
+
|
|
11
|
+
from .collision import CollisionMonitor
|
|
12
|
+
|
|
13
|
+
logger = logging.getLogger("trossen_arm_sim.collision")
|
|
14
|
+
|
|
15
|
+
WS_PORT = 50002
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class CollisionWebSocketServer:
|
|
19
|
+
"""Streams CollisionMonitor events to every connected WebSocket client.
|
|
20
|
+
|
|
21
|
+
On connect a client receives one ``snapshot`` message with the contacts
|
|
22
|
+
active at that moment, then a ``begin`` or ``end`` message per collision
|
|
23
|
+
event (see CollisionEvent.to_dict). Messages sent by clients are ignored.
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
def __init__(self, monitor: CollisionMonitor, host: str = "0.0.0.0",
|
|
27
|
+
port: int = WS_PORT):
|
|
28
|
+
self.monitor = monitor
|
|
29
|
+
self.host = host
|
|
30
|
+
self.port = port
|
|
31
|
+
self._server: Server | None = None
|
|
32
|
+
self._thread: threading.Thread | None = None
|
|
33
|
+
self._lock = threading.Lock()
|
|
34
|
+
self._queues: set[queue.Queue] = set()
|
|
35
|
+
|
|
36
|
+
def start(self) -> None:
|
|
37
|
+
self._server = serve(self._handle, self.host, self.port)
|
|
38
|
+
self._thread = threading.Thread(target=self._server.serve_forever,
|
|
39
|
+
daemon=True, name="collision-ws")
|
|
40
|
+
self._thread.start()
|
|
41
|
+
logger.info("Collision stream listening on ws://%s:%d", self.host, self.port)
|
|
42
|
+
|
|
43
|
+
def stop(self) -> None:
|
|
44
|
+
if self._server is None:
|
|
45
|
+
return
|
|
46
|
+
self._server.shutdown()
|
|
47
|
+
with self._lock:
|
|
48
|
+
for events in self._queues:
|
|
49
|
+
events.put(None) # wakes the handler so it closes its connection
|
|
50
|
+
self._thread.join()
|
|
51
|
+
self._server = None
|
|
52
|
+
self._thread = None
|
|
53
|
+
|
|
54
|
+
def _handle(self, websocket: ServerConnection) -> None:
|
|
55
|
+
logger.info("Collision stream client connected: %s", websocket.remote_address)
|
|
56
|
+
snapshot, events = self.monitor.subscribe()
|
|
57
|
+
with self._lock:
|
|
58
|
+
self._queues.add(events)
|
|
59
|
+
try:
|
|
60
|
+
websocket.send(json.dumps({
|
|
61
|
+
"type": "snapshot",
|
|
62
|
+
"time": round(self.monitor.time, 4),
|
|
63
|
+
"contacts": [contact.to_dict() for contact in snapshot],
|
|
64
|
+
}))
|
|
65
|
+
while (event := events.get()) is not None:
|
|
66
|
+
websocket.send(json.dumps(event.to_dict()))
|
|
67
|
+
websocket.close()
|
|
68
|
+
except ConnectionClosed:
|
|
69
|
+
pass
|
|
70
|
+
finally:
|
|
71
|
+
with self._lock:
|
|
72
|
+
self._queues.discard(events)
|
|
73
|
+
self.monitor.unsubscribe(events)
|
|
74
|
+
logger.info("Collision stream client disconnected")
|