mani-skill-nightly 2025.4.5.813__py3-none-any.whl → 2025.4.5.2036__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.
Files changed (91) hide show
  1. mani_skill/agents/base_real_agent.py +202 -0
  2. mani_skill/agents/controllers/base_controller.py +13 -2
  3. mani_skill/agents/controllers/passive_controller.py +2 -0
  4. mani_skill/agents/controllers/pd_joint_pos.py +2 -0
  5. mani_skill/agents/controllers/pd_joint_pos_vel.py +2 -0
  6. mani_skill/agents/controllers/pd_joint_vel.py +2 -0
  7. mani_skill/agents/robots/__init__.py +2 -0
  8. mani_skill/agents/robots/koch/__init__.py +1 -0
  9. mani_skill/agents/robots/koch/koch.py +168 -0
  10. mani_skill/agents/robots/koch/koch_real.py +5 -0
  11. mani_skill/agents/robots/so100/__init__.py +1 -0
  12. mani_skill/agents/robots/so100/so_100.py +118 -0
  13. mani_skill/agents/robots/so100/so_100_real.py +5 -0
  14. mani_skill/assets/robots/koch/LICENSE +507 -0
  15. mani_skill/assets/robots/koch/README.md +8 -0
  16. mani_skill/assets/robots/koch/follower_arm_v1.1.srdf +9 -0
  17. mani_skill/assets/robots/koch/follower_arm_v1.1.urdf +635 -0
  18. mani_skill/assets/robots/koch/meshes/base_link.glb +0 -0
  19. mani_skill/assets/robots/koch/meshes/base_link.stl +0 -0
  20. mani_skill/assets/robots/koch/meshes/centered_base_link.stl +0 -0
  21. mani_skill/assets/robots/koch/meshes/gripper.glb +0 -0
  22. mani_skill/assets/robots/koch/meshes/gripper.stl +0 -0
  23. mani_skill/assets/robots/koch/meshes/gripper_collision_part_1.glb +0 -0
  24. mani_skill/assets/robots/koch/meshes/gripper_collision_part_2.glb +0 -0
  25. mani_skill/assets/robots/koch/meshes/gripper_collision_part_3.glb +0 -0
  26. mani_skill/assets/robots/koch/meshes/link_1.stl +0 -0
  27. mani_skill/assets/robots/koch/meshes/link_1_motor.glb +0 -0
  28. mani_skill/assets/robots/koch/meshes/link_2.stl +0 -0
  29. mani_skill/assets/robots/koch/meshes/link_2_collision_chassis.glb +0 -0
  30. mani_skill/assets/robots/koch/meshes/link_2_collision_motor.glb +0 -0
  31. mani_skill/assets/robots/koch/meshes/link_2_motor.glb +0 -0
  32. mani_skill/assets/robots/koch/meshes/link_2_rotation_connector.glb +0 -0
  33. mani_skill/assets/robots/koch/meshes/link_2_rotation_connector.stl +0 -0
  34. mani_skill/assets/robots/koch/meshes/link_3.stl +0 -0
  35. mani_skill/assets/robots/koch/meshes/link_3_collision_chassis_part_1.glb +0 -0
  36. mani_skill/assets/robots/koch/meshes/link_3_collision_chassis_part_2.glb +0 -0
  37. mani_skill/assets/robots/koch/meshes/link_3_collision_chassis_part_3.glb +0 -0
  38. mani_skill/assets/robots/koch/meshes/link_3_collision_chassis_part_4.glb +0 -0
  39. mani_skill/assets/robots/koch/meshes/link_3_collision_chassis_part_5.glb +0 -0
  40. mani_skill/assets/robots/koch/meshes/link_3_collision_motor.glb +0 -0
  41. mani_skill/assets/robots/koch/meshes/link_3_motor.glb +0 -0
  42. mani_skill/assets/robots/koch/meshes/link_3_motor.stl +0 -0
  43. mani_skill/assets/robots/koch/meshes/link_3_part.glb +0 -0
  44. mani_skill/assets/robots/koch/meshes/link_3_part.stl +0 -0
  45. mani_skill/assets/robots/koch/meshes/link_4.stl +0 -0
  46. mani_skill/assets/robots/koch/meshes/link_4_collision_chassis_part_1.glb +0 -0
  47. mani_skill/assets/robots/koch/meshes/link_4_collision_chassis_part_2.glb +0 -0
  48. mani_skill/assets/robots/koch/meshes/link_4_collision_chassis_part_3.glb +0 -0
  49. mani_skill/assets/robots/koch/meshes/link_4_collision_motor.glb +0 -0
  50. mani_skill/assets/robots/koch/meshes/link_4_motor.glb +0 -0
  51. mani_skill/assets/robots/koch/meshes/link_4_part.glb +0 -0
  52. mani_skill/assets/robots/koch/meshes/link_5.stl +0 -0
  53. mani_skill/assets/robots/koch/meshes/link_5_motor.glb +0 -0
  54. mani_skill/assets/robots/koch/meshes/link_5_part.glb +0 -0
  55. mani_skill/assets/robots/koch/meshes/link_6.stl +0 -0
  56. mani_skill/assets/robots/koch/meshes/link_6_collision_part_2.glb +0 -0
  57. mani_skill/assets/robots/koch/meshes/link_6_collision_part_3.glb +0 -0
  58. mani_skill/assets/robots/koch/meshes/link_6_collision_part_4.glb +0 -0
  59. mani_skill/assets/robots/koch/meshes/link_6_motor.glb +0 -0
  60. mani_skill/assets/robots/koch/meshes/link_6_part.glb +0 -0
  61. mani_skill/assets/robots/so100/LICENSE +201 -0
  62. mani_skill/assets/robots/so100/README.md +10 -0
  63. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Base.STL +0 -0
  64. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Fixed_Jaw.STL +0 -0
  65. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Fixed_Jaw_part1.ply +0 -0
  66. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Fixed_Jaw_part2.ply +0 -0
  67. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Lower_Arm.STL +0 -0
  68. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Moving Jaw.STL +0 -0
  69. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Moving_Jaw_part1.ply +0 -0
  70. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Moving_Jaw_part2.ply +0 -0
  71. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Moving_Jaw_part3.ply +0 -0
  72. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Rotation_Pitch.STL +0 -0
  73. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Upper_Arm.STL +0 -0
  74. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/meshes/Wrist_Pitch_Roll.STL +0 -0
  75. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/original.srdf +8 -0
  76. mani_skill/assets/robots/so100/SO_5DOF_ARM100_8j/original.urdf +470 -0
  77. mani_skill/envs/sapien_env.py +70 -9
  78. mani_skill/envs/sim2real_env.py +381 -0
  79. mani_skill/envs/tasks/digital_twins/base_env.py +74 -74
  80. mani_skill/envs/tasks/digital_twins/bridge_dataset_eval/base_env.py +6 -0
  81. mani_skill/envs/tasks/digital_twins/bridge_dataset_eval/put_on_in_scene.py +14 -1
  82. mani_skill/envs/utils/randomization/__init__.py +1 -0
  83. mani_skill/envs/utils/randomization/camera.py +60 -0
  84. mani_skill/examples/demo_robot.py +1 -0
  85. mani_skill/utils/sapien_utils.py +7 -6
  86. mani_skill/utils/structs/articulation.py +44 -18
  87. {mani_skill_nightly-2025.4.5.813.dist-info → mani_skill_nightly-2025.4.5.2036.dist-info}/METADATA +1 -1
  88. {mani_skill_nightly-2025.4.5.813.dist-info → mani_skill_nightly-2025.4.5.2036.dist-info}/RECORD +91 -19
  89. {mani_skill_nightly-2025.4.5.813.dist-info → mani_skill_nightly-2025.4.5.2036.dist-info}/LICENSE +0 -0
  90. {mani_skill_nightly-2025.4.5.813.dist-info → mani_skill_nightly-2025.4.5.2036.dist-info}/WHEEL +0 -0
  91. {mani_skill_nightly-2025.4.5.813.dist-info → mani_skill_nightly-2025.4.5.2036.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,470 @@
1
+ <?xml version="1.0" encoding="utf-8"?>
2
+ <!-- This URDF was automatically created by SolidWorks to URDF Exporter! Originally created by Stephen Brawner (brawner@gmail.com)
3
+ Commit Version: 1.6.0-1-g15f4949 Build Version: 1.6.7594.29634
4
+ For more information, please see http://wiki.ros.org/sw_urdf_exporter -->
5
+ <robot
6
+ name="SO_5DOF_ARM100_8j_URDF.SLDASM">
7
+ <link
8
+ name="Base">
9
+ <inertial>
10
+ <origin
11
+ xyz="-2.45960666746703E-07 0.0311418169687909 0.0175746661003382"
12
+ rpy="0 0 0" />
13
+ <mass
14
+ value="0.193184127927598" />
15
+ <inertia
16
+ ixx="0.000137030709467877"
17
+ ixy="2.10136126944992E-08"
18
+ ixz="4.24087422551286E-09"
19
+ iyy="0.000169089551209259"
20
+ iyz="2.26514711036514E-05"
21
+ izz="0.000145097720857224" />
22
+ </inertial>
23
+ <visual>
24
+ <origin
25
+ xyz="0 0 0"
26
+ rpy="0 0 0" />
27
+ <geometry>
28
+ <mesh
29
+ filename="meshes/Base.STL" />
30
+ </geometry>
31
+ <material
32
+ name="">
33
+ <color
34
+ rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
35
+ </material>
36
+ </visual>
37
+ <collision>
38
+ <origin
39
+ xyz="0 0 0"
40
+ rpy="0 0 0" />
41
+ <geometry>
42
+ <mesh
43
+ filename="meshes/Base.STL" />
44
+ </geometry>
45
+ </collision>
46
+ </link>
47
+ <link
48
+ name="Rotation_Pitch">
49
+ <inertial>
50
+ <origin
51
+ xyz="-9.07886224712597E-05 0.0590971820568318 0.031089016892169"
52
+ rpy="0 0 0" />
53
+ <mass
54
+ value="0.119226314127197" />
55
+ <inertia
56
+ ixx="5.90408775624429E-05"
57
+ ixy="4.90800532852998E-07"
58
+ ixz="-5.90451772654387E-08"
59
+ iyy="3.21498601038881E-05"
60
+ iyz="-4.58026206663885E-06"
61
+ izz="5.86058514263952E-05" />
62
+ </inertial>
63
+ <visual>
64
+ <origin
65
+ xyz="0 0 0"
66
+ rpy="0 0 0" />
67
+ <geometry>
68
+ <mesh
69
+ filename="meshes/Rotation_Pitch.STL" />
70
+ </geometry>
71
+ <material
72
+ name="">
73
+ <color
74
+ rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
75
+ </material>
76
+ </visual>
77
+ <collision>
78
+ <origin
79
+ xyz="0 0 0"
80
+ rpy="0 0 0" />
81
+ <geometry>
82
+ <mesh
83
+ filename="meshes/Rotation_Pitch.STL" />
84
+ </geometry>
85
+ </collision>
86
+ </link>
87
+ <joint
88
+ name="Rotation"
89
+ type="revolute">
90
+ <origin
91
+ xyz="0 -0.0452 0.0165"
92
+ rpy="1.5708 0 0" />
93
+ <parent
94
+ link="Base" />
95
+ <child
96
+ link="Rotation_Pitch" />
97
+ <axis
98
+ xyz="0 -1 0" />
99
+ <!-- note for the so100 arm there is no well defined effort/velocity limits at the moment -->
100
+ <limit
101
+ lower="-2.1"
102
+ upper="2.1"
103
+ effort="0"
104
+ velocity="0" />
105
+ </joint>
106
+ <link
107
+ name="Upper_Arm">
108
+ <inertial>
109
+ <origin
110
+ xyz="-1.7205170190925E-05 0.0701802156327694 0.00310545118155671"
111
+ rpy="0 0 0" />
112
+ <mass
113
+ value="0.162409284599177" />
114
+ <inertia
115
+ ixx="0.000167153146617081"
116
+ ixy="1.03902689187701E-06"
117
+ ixz="-1.20161820645189E-08"
118
+ iyy="7.01946992214245E-05"
119
+ iyz="2.11884806298698E-06"
120
+ izz="0.000213280241160769" />
121
+ </inertial>
122
+ <visual>
123
+ <origin
124
+ xyz="0 0 0"
125
+ rpy="0 0 0" />
126
+ <geometry>
127
+ <mesh
128
+ filename="meshes/Upper_Arm.STL" />
129
+ </geometry>
130
+ <material
131
+ name="">
132
+ <color
133
+ rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
134
+ </material>
135
+ </visual>
136
+ <collision>
137
+ <origin
138
+ xyz="0 0 0"
139
+ rpy="0 0 0" />
140
+ <geometry>
141
+ <mesh
142
+ filename="meshes/Upper_Arm.STL" />
143
+ </geometry>
144
+ </collision>
145
+ </link>
146
+ <joint
147
+ name="Pitch"
148
+ type="revolute">
149
+ <origin
150
+ xyz="0 0.1025 0.0306"
151
+ rpy="1.5708 0 0" />
152
+ <parent
153
+ link="Rotation_Pitch" />
154
+ <child
155
+ link="Upper_Arm" />
156
+ <axis
157
+ xyz="-1 0 0" />
158
+ <limit
159
+ lower="-0.1"
160
+ upper="3.45"
161
+ effort="0"
162
+ velocity="0" />
163
+ </joint>
164
+ <link
165
+ name="Lower_Arm">
166
+ <inertial>
167
+ <origin
168
+ xyz="-0.00339603710186651 0.00137796353960074 0.0768006751156044"
169
+ rpy="0 0 0" />
170
+ <mass
171
+ value="0.147967774582291" />
172
+ <inertia
173
+ ixx="0.000105333995841409"
174
+ ixy="1.73059237226499E-07"
175
+ ixz="-1.1720305455211E-05"
176
+ iyy="0.000138766654485212"
177
+ iyz="1.77429964684103E-06"
178
+ izz="5.08741652515214E-05" />
179
+ </inertial>
180
+ <visual>
181
+ <origin
182
+ xyz="0 0 0"
183
+ rpy="0 0 0" />
184
+ <geometry>
185
+ <mesh
186
+ filename="meshes/Lower_Arm.STL" />
187
+ </geometry>
188
+ <material
189
+ name="">
190
+ <color
191
+ rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
192
+ </material>
193
+ </visual>
194
+ <collision>
195
+ <origin
196
+ xyz="0 0 0"
197
+ rpy="0 0 0" />
198
+ <geometry>
199
+ <mesh
200
+ filename="meshes/Lower_Arm.STL" />
201
+ </geometry>
202
+ </collision>
203
+ </link>
204
+ <joint
205
+ name="Elbow"
206
+ type="revolute">
207
+ <origin
208
+ xyz="0 0.11257 0.028"
209
+ rpy="-1.5708 0 0" />
210
+ <parent
211
+ link="Upper_Arm" />
212
+ <child
213
+ link="Lower_Arm" />
214
+ <axis
215
+ xyz="1 0 0" />
216
+ <limit
217
+ lower="-0.2"
218
+ upper="3.14159"
219
+ effort="0"
220
+ velocity="0" />
221
+ </joint>
222
+ <link
223
+ name="Wrist_Pitch_Roll">
224
+ <inertial>
225
+ <origin
226
+ xyz="-0.00852653127372418 -0.0352278997897927 -2.34622481569413E-05"
227
+ rpy="0 0 0" />
228
+ <mass
229
+ value="0.066132067097723" />
230
+ <inertia
231
+ ixx="1.95717492443445E-05"
232
+ ixy="-6.62714374412293E-07"
233
+ ixz="5.20089016442066E-09"
234
+ iyy="2.38028417569933E-05"
235
+ iyz="4.09549055863776E-08"
236
+ izz="3.4540143384536E-05" />
237
+ </inertial>
238
+ <visual>
239
+ <origin
240
+ xyz="0 0 0"
241
+ rpy="0 0 0" />
242
+ <geometry>
243
+ <mesh
244
+ filename="meshes/Wrist_Pitch_Roll.STL" />
245
+ </geometry>
246
+ <material
247
+ name="">
248
+ <color
249
+ rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
250
+ </material>
251
+ </visual>
252
+ <collision>
253
+ <origin
254
+ xyz="0 0 0"
255
+ rpy="0 0 0" />
256
+ <geometry>
257
+ <mesh
258
+ filename="meshes/Wrist_Pitch_Roll.STL" />
259
+ </geometry>
260
+ </collision>
261
+ </link>
262
+ <joint
263
+ name="Wrist_Pitch"
264
+ type="revolute">
265
+ <origin
266
+ xyz="0 0.0052 0.1349"
267
+ rpy="-1.5708 0 0" />
268
+ <parent
269
+ link="Lower_Arm" />
270
+ <child
271
+ link="Wrist_Pitch_Roll" />
272
+ <axis
273
+ xyz="1 0 0" />
274
+ <limit
275
+ lower="-1.8"
276
+ upper="1.8"
277
+ effort="0"
278
+ velocity="0" />
279
+ </joint>
280
+ <link
281
+ name="Fixed_Jaw">
282
+ <inertial>
283
+ <origin
284
+ xyz="0.00552376906426563 -0.0280167153359021 0.000483582592841092"
285
+ rpy="0 0 0" />
286
+ <mass
287
+ value="0.0929859131176897" />
288
+ <inertia
289
+ ixx="4.3328249304211E-05"
290
+ ixy="7.09654328670947E-06"
291
+ ixz="5.99838530879484E-07"
292
+ iyy="3.04451747368212E-05"
293
+ iyz="-1.58743247545413E-07"
294
+ izz="5.02460913506734E-05" />
295
+ </inertial>
296
+ <visual>
297
+ <origin
298
+ xyz="0 0 0"
299
+ rpy="0 0 0" />
300
+ <geometry>
301
+ <mesh
302
+ filename="meshes/Fixed_Jaw.STL" />
303
+ </geometry>
304
+ <material
305
+ name="">
306
+ <color
307
+ rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
308
+ </material>
309
+ </visual>
310
+ <!-- <collision>
311
+ <origin
312
+ xyz="0 0 0"
313
+ rpy="0 0 0" />
314
+ <geometry>
315
+ <mesh
316
+ filename="meshes/Fixed_Jaw.STL" />
317
+ </geometry>
318
+ </collision> -->
319
+ <collision>
320
+ <origin
321
+ xyz="0 0 0"
322
+ rpy="0 0 0" />
323
+ <geometry>
324
+ <mesh
325
+ filename="meshes/Fixed_Jaw_part1.ply" />
326
+ </geometry>
327
+ </collision>
328
+ <collision>
329
+ <origin
330
+ xyz="0 0 0"
331
+ rpy="0 0 0" />
332
+ <geometry>
333
+ <mesh
334
+ filename="meshes/Fixed_Jaw_part2.ply" />
335
+ </geometry>
336
+ </collision>
337
+ </link>
338
+ <joint
339
+ name="Wrist_Roll"
340
+ type="revolute">
341
+ <origin
342
+ xyz="0 -0.0601 0"
343
+ rpy="0 1.5708 0" />
344
+ <parent
345
+ link="Wrist_Pitch_Roll" />
346
+ <child
347
+ link="Fixed_Jaw" />
348
+ <axis
349
+ xyz="0 -1 0" />
350
+ <limit
351
+ lower="-3.14159"
352
+ upper="3.14159"
353
+ effort="0"
354
+ velocity="0" />
355
+ </joint>
356
+ <link
357
+ name="Moving_Jaw">
358
+ <inertial>
359
+ <origin
360
+ xyz="-0.00161744605468241 -0.0303472584046471 0.000449645961853651"
361
+ rpy="0 0 0" />
362
+ <mass
363
+ value="0.0202443794940372" />
364
+ <inertia
365
+ ixx="1.10911325081525E-05"
366
+ ixy="-5.35076503033314E-07"
367
+ ixz="-9.46105662101403E-09"
368
+ iyy="3.03576451001973E-06"
369
+ iyz="-1.71146075110632E-07"
370
+ izz="8.9916083370498E-06" />
371
+ </inertial>
372
+ <visual>
373
+ <origin
374
+ xyz="0 0 0"
375
+ rpy="0 0 0" />
376
+ <geometry>
377
+ <mesh
378
+ filename="meshes/Moving Jaw.STL" />
379
+ </geometry>
380
+ <material
381
+ name="">
382
+ <color
383
+ rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
384
+ </material>
385
+ </visual>
386
+ <collision>
387
+ <origin
388
+ xyz="0 0 0"
389
+ rpy="0 0 0" />
390
+ <geometry>
391
+ <mesh
392
+ filename="meshes/Moving_Jaw_part1.ply" />
393
+ </geometry>
394
+ </collision>
395
+ <collision>
396
+ <origin
397
+ xyz="0 0 0"
398
+ rpy="0 0 0" />
399
+ <geometry>
400
+ <mesh
401
+ filename="meshes/Moving_Jaw_part2.ply" />
402
+ </geometry>
403
+ </collision>
404
+ <collision>
405
+ <origin
406
+ xyz="0 0 0"
407
+ rpy="0 0 0" />
408
+ <geometry>
409
+ <mesh
410
+ filename="meshes/Moving_Jaw_part3.ply" />
411
+ </geometry>
412
+ </collision>
413
+ </link>
414
+ <joint
415
+ name="Jaw"
416
+ type="revolute">
417
+ <origin
418
+ xyz="-0.0202 -0.0244 0"
419
+ rpy="3.1416 0 3.33" />
420
+ <parent
421
+ link="Fixed_Jaw" />
422
+ <child
423
+ link="Moving_Jaw" />
424
+ <axis
425
+ xyz="0 0 1" />
426
+ <limit
427
+ lower="0"
428
+ upper="1.7"
429
+ effort="0"
430
+ velocity="0" />
431
+ </joint>
432
+
433
+ <!-- fake joints for useful data -->
434
+ <joint name="Fixed_Jaw_tip_joint" type="fixed">
435
+ <origin rpy="0 0 0" xyz="0.01 -0.097 0"/>
436
+ <parent link="Fixed_Jaw"/>
437
+ <child link="Fixed_Jaw_tip"/>
438
+ </joint>
439
+ <link name="Fixed_Jaw_tip">
440
+ <!-- <visual>
441
+ <origin
442
+ xyz="0 0 0"
443
+ rpy="0 0 0" />
444
+ <geometry>
445
+ <sphere radius="0.01" />
446
+ </geometry>
447
+ <material name="red">
448
+ <color rgba="1 0 0 0.5"/>
449
+ </material>
450
+ </visual> -->
451
+ </link>
452
+ <joint name="Moving_Jaw_tip_joint" type="fixed">
453
+ <origin rpy="0 0 0" xyz="-0.01 -0.073 0"/>
454
+ <parent link="Moving_Jaw"/>
455
+ <child link="Moving_Jaw_tip"/>
456
+ </joint>
457
+ <link name="Moving_Jaw_tip">
458
+ <!-- <visual>
459
+ <origin
460
+ xyz="0 0 0"
461
+ rpy="0 0 0" />
462
+ <geometry>
463
+ <sphere radius="0.01" />
464
+ </geometry>
465
+ <material name="red">
466
+ <color rgba="1 0 0 0.5"/>
467
+ </material>
468
+ </visual> -->
469
+ </link>
470
+ </robot>
@@ -480,7 +480,7 @@ class BaseEnv(gym.Env):
480
480
  """The current observation mode. This affects the observation returned by env.get_obs()"""
481
481
  return self._obs_mode
482
482
 
483
- def get_obs(self, info: Optional[Dict] = None):
483
+ def get_obs(self, info: Optional[Dict] = None, unflattened: bool = False):
484
484
  """
485
485
  Return the current observation of the environment. User may call this directly to get the current observation
486
486
  as opposed to taking a step with actions in the environment.
@@ -493,6 +493,7 @@ class BaseEnv(gym.Env):
493
493
  Args:
494
494
  info (Dict): The info object of the environment. Generally should always be the result of `self.get_info()`.
495
495
  If this is None (the default), this function will call `self.get_info()` itself
496
+ unflattened (bool): Whether to return the observation without flattening even if the observation mode (`self.obs_mode`) asserts to return a flattened observation.
496
497
  """
497
498
  if info is None:
498
499
  info = self.get_info()
@@ -500,8 +501,7 @@ class BaseEnv(gym.Env):
500
501
  # Some cases do not need observations, e.g., MPC
501
502
  return dict()
502
503
  elif self._obs_mode == "state":
503
- state_dict = self._get_obs_state_dict(info)
504
- obs = common.flatten_state_dict(state_dict, use_torch=True, device=self.device)
504
+ obs = self._get_obs_state_dict(info)
505
505
  elif self._obs_mode == "state_dict":
506
506
  obs = self._get_obs_state_dict(info)
507
507
  elif self._obs_mode == "pointcloud":
@@ -512,8 +512,13 @@ class BaseEnv(gym.Env):
512
512
  obs = self._get_obs_with_sensor_data(info, apply_texture_transforms=False)
513
513
  else:
514
514
  obs = self._get_obs_with_sensor_data(info)
515
+ return obs if unflattened else self._flatten_raw_obs(obs)
515
516
 
517
+ def _flatten_raw_obs(self, obs: Any):
516
518
  # flatten parts of the state observation if requested
519
+ if self._obs_mode == "state":
520
+ return common.flatten_state_dict(obs, use_torch=True, device=self.device)
521
+
517
522
  if self.obs_mode_struct.state:
518
523
  if isinstance(obs, dict):
519
524
  data = dict(agent=obs.pop("agent"), extra=obs.pop("extra"))
@@ -553,7 +558,28 @@ class BaseEnv(gym.Env):
553
558
  return params
554
559
 
555
560
  def _get_obs_sensor_data(self, apply_texture_transforms: bool = True) -> dict:
556
- """get only data from sensors. Auto hides any objects that are designated to be hidden"""
561
+ """
562
+ Get data from all registered sensors. Auto hides any objects that are designated to be hidden
563
+
564
+ Args:
565
+ apply_texture_transforms (bool): Whether to apply texture transforms to the simulated sensor data to map to standard texture formats. Default is True.
566
+
567
+ Returns:
568
+ dict: A dictionary containing the sensor data mapping sensor name to its respective dictionary of data. The dictionary maps texture names to the data. For example the return could look like
569
+
570
+ .. code-block:: python
571
+
572
+ {
573
+ "sensor_1": {
574
+ "rgb": torch.Tensor,
575
+ "depth": torch.Tensor
576
+ },
577
+ "sensor_2": {
578
+ "rgb": torch.Tensor,
579
+ "depth": torch.Tensor
580
+ }
581
+ }
582
+ """
557
583
  for obj in self._hidden_objects:
558
584
  obj.hide_visual()
559
585
  self.scene.update_render(update_sensors=True, update_human_render_cameras=False)
@@ -579,6 +605,7 @@ class BaseEnv(gym.Env):
579
605
  if self.backend.render_device.is_cuda():
580
606
  torch.cuda.synchronize()
581
607
  return sensor_obs
608
+
582
609
  def _get_obs_with_sensor_data(self, info: Dict, apply_texture_transforms: bool = True) -> dict:
583
610
  """Get the observation with sensor data"""
584
611
  return dict(
@@ -601,6 +628,15 @@ class BaseEnv(gym.Env):
601
628
  return self._reward_mode
602
629
 
603
630
  def get_reward(self, obs: Any, action: torch.Tensor, info: Dict):
631
+ """
632
+ Compute the reward for environment at its current state. observation data, the most recent action, and the info dictionary (generated by the self.evaluate() function)
633
+ are provided as inputs. By default the observation data will be in its most raw form, a dictionary (no flattening, wrappers etc.)
634
+
635
+ Args:
636
+ obs (Any): The observation data.
637
+ action (torch.Tensor): The most recent action.
638
+ info (Dict): The info dictionary.
639
+ """
604
640
  if self._reward_mode == "sparse":
605
641
  reward = self.compute_sparse_reward(obs=obs, action=action, info=info)
606
642
  elif self._reward_mode == "dense":
@@ -617,8 +653,15 @@ class BaseEnv(gym.Env):
617
653
 
618
654
  def compute_sparse_reward(self, obs: Any, action: torch.Tensor, info: Dict):
619
655
  """
656
+
620
657
  Computes the sparse reward. By default this function tries to use the success/fail information in
621
- returned by the evaluate function and gives +1 if success, -1 if fail, 0 otherwise"""
658
+ returned by the evaluate function and gives +1 if success, -1 if fail, 0 otherwise.
659
+
660
+ Args:
661
+ obs (Any): The observation data. By default the observation data will be in its most raw form, a dictionary (no flattening, wrappers etc.)
662
+ action (torch.Tensor): The most recent action.
663
+ info (Dict): The info dictionary.
664
+ """
622
665
  if "success" in info:
623
666
  if "fail" in info:
624
667
  if isinstance(info["success"], torch.Tensor):
@@ -635,11 +678,27 @@ class BaseEnv(gym.Env):
635
678
  return reward
636
679
 
637
680
  def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: Dict):
681
+ """
682
+ Compute the dense reward.
683
+
684
+ Args:
685
+ obs (Any): The observation data. By default the observation data will be in its most raw form, a dictionary (no flattening, wrappers etc.)
686
+ action (torch.Tensor): The most recent action.
687
+ info (Dict): The info dictionary.
688
+ """
638
689
  raise NotImplementedError()
639
690
 
640
691
  def compute_normalized_dense_reward(
641
692
  self, obs: Any, action: torch.Tensor, info: Dict
642
693
  ):
694
+ """
695
+ Compute the normalized dense reward.
696
+
697
+ Args:
698
+ obs (Any): The observation data. By default the observation data will be in its most raw form, a dictionary (no flattening, wrappers etc.)
699
+ action (torch.Tensor): The most recent action.
700
+ info (Dict): The info dictionary.
701
+ """
643
702
  raise NotImplementedError()
644
703
 
645
704
  # -------------------------------------------------------------------------- #
@@ -949,10 +1008,10 @@ class BaseEnv(gym.Env):
949
1008
  action = self._step_action(action)
950
1009
  self._elapsed_steps += 1
951
1010
  info = self.get_info()
952
- obs = self.get_obs(info)
1011
+ obs = self.get_obs(info, unflattened=True)
953
1012
  reward = self.get_reward(obs=obs, action=action, info=info)
1013
+ obs = self._flatten_raw_obs(obs)
954
1014
  if "success" in info:
955
-
956
1015
  if "fail" in info:
957
1016
  terminated = torch.logical_or(info["success"], info["fail"])
958
1017
  else:
@@ -1010,8 +1069,10 @@ class BaseEnv(gym.Env):
1010
1069
  action = common.batch(action)
1011
1070
  self.agent.set_action(action)
1012
1071
  if self._sim_device.is_cuda():
1013
- self.scene.px.gpu_apply_articulation_target_position()
1014
- self.scene.px.gpu_apply_articulation_target_velocity()
1072
+ if self.agent.controller.sets_target_qpos:
1073
+ self.scene.px.gpu_apply_articulation_target_position()
1074
+ if self.agent.controller.sets_target_qvel:
1075
+ self.scene.px.gpu_apply_articulation_target_velocity()
1015
1076
  self._before_control_step()
1016
1077
  for _ in range(self._sim_steps_per_control):
1017
1078
  if self.agent is not None: