tinysim 0.0.1__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.
@@ -0,0 +1,168 @@
1
+ import math
2
+ import os
3
+ import pathlib
4
+
5
+ import numpy as np
6
+
7
+ try:
8
+ import warp as wp
9
+ import warp.sim.render
10
+ except ImportError:
11
+ raise ImportError(
12
+ "Warp is not installed. Install using `pip install coderbot_sim[warp]`"
13
+ )
14
+
15
+
16
+ def compute_env_offsets(num_envs, env_offset=(5.0, 0.0, 0.0)):
17
+ env_offset = np.array(env_offset, dtype=float)
18
+ axis = np.nonzero(env_offset)[0]
19
+ axis = axis[0] if len(axis) else 0
20
+ env_offsets = np.zeros((num_envs, 3))
21
+ env_offsets[:, axis] = np.arange(num_envs) * env_offset[axis]
22
+ env_offsets -= env_offsets.mean(axis=0)
23
+ env_offsets[:, 1] = 0.0
24
+ return env_offsets
25
+
26
+
27
+ class SimpleRobotDogExample:
28
+ def __init__(self, use_cuda_graph=False, headless=False, num_envs=8):
29
+ articulation_builder = wp.sim.ModelBuilder()
30
+ rot_x = wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)
31
+ rot_y = wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), math.pi * 0.5)
32
+ xform = wp.transform(wp.vec3(0.0, 0.35, 0.0), rot_y * rot_x)
33
+ wp.sim.parse_urdf(
34
+ # os.path.join(warp.examples.get_asset_directory(), "quadruped.urdf"),
35
+ str(pathlib.Path(__file__).parent / "simple_quadruped.urdf"),
36
+ articulation_builder,
37
+ xform=xform,
38
+ floating=True,
39
+ density=900,
40
+ armature=0.01,
41
+ stiffness=200,
42
+ damping=1,
43
+ contact_ke=1.0e4,
44
+ contact_kd=1.0e2,
45
+ contact_kf=1.0e2,
46
+ contact_mu=1.0,
47
+ limit_ke=1.0e4,
48
+ limit_kd=1.0e1,
49
+ )
50
+
51
+ builder = wp.sim.ModelBuilder()
52
+
53
+ self.sim_time = 0.0
54
+ fps = 100
55
+ self.frame_dt = 1.0 / fps
56
+
57
+ self.sim_substeps = 10
58
+ self.sim_dt = self.frame_dt / self.sim_substeps
59
+
60
+ self.num_envs = num_envs
61
+ offsets = compute_env_offsets(self.num_envs)
62
+
63
+ for i in range(self.num_envs):
64
+ builder.add_builder(
65
+ articulation_builder, xform=wp.transform(offsets[i], wp.quat_identity())
66
+ )
67
+ builder.joint_q[-8:] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] # fmt: skip
68
+ builder.joint_act[-8:] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] # fmt: skip
69
+ builder.joint_axis_mode = [wp.sim.JOINT_MODE_TARGET_POSITION] * len(
70
+ builder.joint_axis_mode
71
+ )
72
+
73
+ # finalize model
74
+ self.model = builder.finalize()
75
+ self.model.ground = True
76
+ self.model.joint_attach_ke = 16000.0
77
+ self.model.joint_attach_kd = 200.0
78
+ self.use_tile_gemm = False
79
+ self.fuse_cholesky = False
80
+
81
+ self.integrator = wp.sim.FeatherstoneIntegrator(
82
+ self.model,
83
+ use_tile_gemm=self.use_tile_gemm,
84
+ fuse_cholesky=self.fuse_cholesky,
85
+ )
86
+
87
+ self.renderer = wp.sim.render.SimRendererOpenGL(
88
+ self.model, "example", headless=headless
89
+ )
90
+ self.state_0 = self.model.state()
91
+ self.state_1 = self.model.state()
92
+
93
+ wp.sim.eval_fk(
94
+ self.model, self.model.joint_q, self.model.joint_qd, None, self.state_0
95
+ )
96
+
97
+ self.use_cuda_graph = (
98
+ use_cuda_graph
99
+ and wp.get_device().is_cuda
100
+ and wp.is_mempool_enabled(wp.get_device())
101
+ )
102
+ if self.use_cuda_graph:
103
+ with wp.ScopedCapture() as capture:
104
+ self.simulate()
105
+ self.graph = capture.graph
106
+
107
+ def simulate(self):
108
+ for _ in range(self.sim_substeps):
109
+ self.state_0.clear_forces()
110
+ wp.sim.collide(self.model, self.state_0)
111
+ self.integrator.simulate(
112
+ self.model, self.state_0, self.state_1, self.sim_dt
113
+ )
114
+ self.state_0, self.state_1 = self.state_1, self.state_0
115
+
116
+ def set_leg_poses(self, actions):
117
+ action_space_size = 8 * self.num_envs
118
+ if len(actions) != action_space_size:
119
+ raise ValueError(
120
+ f"Expected {action_space_size} actions, but got {len(actions)}"
121
+ )
122
+
123
+ wp.copy(self.model.joint_act, wp.array(actions, dtype=wp.float32))
124
+
125
+ def step(self, actions):
126
+ self.set_leg_poses(actions)
127
+
128
+ if self.use_cuda_graph:
129
+ wp.capture_launch(self.graph)
130
+ else:
131
+ self.simulate()
132
+ self.sim_time += self.frame_dt
133
+
134
+ def render(self):
135
+ self.renderer.begin_frame(self.sim_time)
136
+ self.renderer.render(self.state_0)
137
+ self.renderer.end_frame()
138
+
139
+
140
+ if __name__ == "__main__":
141
+ import argparse
142
+
143
+ parser = argparse.ArgumentParser(
144
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
145
+ )
146
+ parser.add_argument(
147
+ "--device", type=str, default=None, help="Override the default Warp device."
148
+ )
149
+ parser.add_argument(
150
+ "--num-frames", type=int, default=1500, help="Total number of frames."
151
+ )
152
+ parser.add_argument(
153
+ "--num-envs",
154
+ type=int,
155
+ default=3,
156
+ help="Total number of simulated environments.",
157
+ )
158
+
159
+ args = parser.parse_known_args()[0]
160
+
161
+ with wp.ScopedDevice(args.device):
162
+ example = SimpleRobotDogExample(num_envs=args.num_envs)
163
+
164
+ for frame in range(args.num_frames):
165
+ example.step(np.zeros(8).repeat(args.num_envs))
166
+ example.render()
167
+
168
+ example.renderer.save()
@@ -0,0 +1,247 @@
1
+ <?xml version="1.0" encoding="utf-8"?>
2
+ <robot name="quadruped">
3
+
4
+ <!-- Base -->
5
+ <link name="base">
6
+ <visual>
7
+ <origin rpy="0 1.57079632679 0" xyz="0 0 0"/>
8
+ <geometry>
9
+ <cylinder length="0.75" radius="0.1"/>
10
+ </geometry>
11
+ </visual>
12
+ <collision>
13
+ <origin rpy="0 1.57079632679 0" xyz="0 0 0"/>
14
+ <geometry>
15
+ <cylinder length="0.75" radius="0.1"/>
16
+ </geometry>
17
+ </collision>
18
+ <inertial>
19
+ <origin rpy="0 0 0" xyz="0 0 0.024"/>
20
+ <mass value="6.222"/>
21
+ <inertia ixx="0.017938806" ixy="0.00387963" ixz="0.001500772" iyy="0.370887745" iyz="6.8963e-05" izz="0.372497653"/>
22
+ </inertial>
23
+ </link>
24
+
25
+ <!-- LEFT FRONT LEG -->
26
+ <link name="LF_HAA">
27
+ <visual>
28
+ <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
29
+ <geometry>
30
+ <cylinder length="0.05" radius="0.04"/>
31
+ </geometry>
32
+ </visual>
33
+ <collision>
34
+ <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
35
+ <geometry>
36
+ <cylinder length="0.05" radius="0.04"/>
37
+ </geometry>
38
+ </collision>
39
+ <inertial>
40
+ <origin rpy="0 0 0" xyz="0 0 0.00046"/>
41
+ <mass value="2.04"/>
42
+ <inertia ixx="0.001053013" ixy="4.527e-05" ixz="8.855e-05" iyy="0.001805509" iyz="9.909e-05" izz="0.001765827"/>
43
+ </inertial>
44
+ </link>
45
+ <joint name="LF_HAA" type="revolute">
46
+ <parent link="base"/>
47
+ <child link="LF_HAA"/>
48
+ <origin rpy="0 0 0" xyz="0.2999 0.104 0.0"/>
49
+ <axis xyz="1 0 0"/>
50
+ <limit effort="80.0" velocity="20."/>
51
+ </joint>
52
+
53
+ <link name="LF_THIGH">
54
+ <visual>
55
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
56
+ <geometry>
57
+ <cylinder length="0.25" radius="0.02"/>
58
+ </geometry>
59
+ </visual>
60
+ <collision>
61
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
62
+ <geometry>
63
+ <cylinder length="0.25" radius="0.02"/>
64
+ </geometry>
65
+ </collision>
66
+ <inertial>
67
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
68
+ <mass value="2.04"/>
69
+ <inertia ixx="0.001053013" ixy="4.527e-05" ixz="8.855e-05" iyy="0.001805509" iyz="9.909e-05" izz="0.001765827"/>
70
+ </inertial>
71
+ </link>
72
+ <joint name="LF_HFE" type="revolute">
73
+ <parent link="LF_HAA"/>
74
+ <child link="LF_THIGH"/>
75
+ <origin rpy="0 0 1.57079632679" xyz="0 0.05 0"/>
76
+ <axis xyz="1 0 0"/>
77
+ <limit effort="80.0" velocity="20."/>
78
+ </joint>
79
+
80
+
81
+
82
+ <!-- RIGHT FRONT LEG -->
83
+ <link name="RF_HAA">
84
+ <visual>
85
+ <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
86
+ <geometry>
87
+ <cylinder length="0.05" radius="0.04"/>
88
+ </geometry>
89
+ </visual>
90
+ <collision>
91
+ <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
92
+ <geometry>
93
+ <cylinder length="0.05" radius="0.04"/>
94
+ </geometry>
95
+ </collision>
96
+ <inertial>
97
+ <origin rpy="0 0 0" xyz="0 0 0.00046"/>
98
+ <mass value="2.04"/>
99
+ <inertia ixx="0.001053013" ixy="4.527e-05" ixz="8.855e-05" iyy="0.001805509" iyz="9.909e-05" izz="0.001765827"/>
100
+ </inertial>
101
+ </link>
102
+ <joint name="RF_HAA" type="revolute">
103
+ <parent link="base"/>
104
+ <child link="RF_HAA"/>
105
+ <origin rpy="0 0 0" xyz="0.2999 -0.104 0.0"/>
106
+ <axis xyz="1 0 0"/>
107
+ <limit effort="80.0" velocity="20."/>
108
+ </joint>
109
+
110
+ <link name="RF_THIGH">
111
+ <visual>
112
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
113
+ <geometry>
114
+ <cylinder length="0.25" radius="0.02"/>
115
+ </geometry>
116
+ </visual>
117
+ <collision>
118
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
119
+ <geometry>
120
+ <cylinder length="0.25" radius="0.02"/>
121
+ </geometry>
122
+ </collision>
123
+ <inertial>
124
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
125
+ <mass value="2.04"/>
126
+ <inertia ixx="0.001053013" ixy="4.527e-05" ixz="8.855e-05" iyy="0.001805509" iyz="9.909e-05" izz="0.001765827"/>
127
+ </inertial>
128
+ </link>
129
+ <joint name="RF_HFE" type="revolute">
130
+ <parent link="RF_HAA"/>
131
+ <child link="RF_THIGH"/>
132
+ <origin rpy="0 0 -1.57079632679" xyz="0 -0.05 0"/>
133
+ <axis xyz="1 0 0"/>
134
+ <limit effort="80.0" velocity="20."/>
135
+ </joint>
136
+
137
+ <!-- LEFT HIND LEG -->
138
+ <link name="LH_HAA">
139
+ <visual>
140
+ <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
141
+ <geometry>
142
+ <cylinder length="0.05" radius="0.04"/>
143
+ </geometry>
144
+ </visual>
145
+ <collision>
146
+ <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
147
+ <geometry>
148
+ <cylinder length="0.05" radius="0.04"/>
149
+ </geometry>
150
+ </collision>
151
+ <inertial>
152
+ <origin rpy="0 0 0" xyz="0 0 0.00046"/>
153
+ <mass value="2.04"/>
154
+ <inertia ixx="0.001053013" ixy="4.527e-05" ixz="8.855e-05" iyy="0.001805509" iyz="9.909e-05" izz="0.001765827"/>
155
+ </inertial>
156
+ </link>
157
+ <joint name="LH_HAA" type="revolute">
158
+ <parent link="base"/>
159
+ <child link="LH_HAA"/>
160
+ <origin rpy="0 0 3.1415" xyz="-0.2999 0.104 0.0"/>
161
+ <axis xyz="1 0 0"/>
162
+ <limit effort="80.0" velocity="20."/>
163
+ </joint>
164
+
165
+ <link name="LH_THIGH">
166
+ <visual>
167
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
168
+ <geometry>
169
+ <cylinder length="0.25" radius="0.02"/>
170
+ </geometry>
171
+ </visual>
172
+ <collision>
173
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
174
+ <geometry>
175
+ <cylinder length="0.25" radius="0.02"/>
176
+ </geometry>
177
+ </collision>
178
+ <inertial>
179
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
180
+ <mass value="2.04"/>
181
+ <inertia ixx="0.001053013" ixy="4.527e-05" ixz="8.855e-05" iyy="0.001805509" iyz="9.909e-05" izz="0.001765827"/>
182
+ </inertial>
183
+ </link>
184
+ <joint name="LH_HFE" type="revolute">
185
+ <parent link="LH_HAA"/>
186
+ <child link="LH_THIGH"/>
187
+ <origin rpy="0 0 1.57079632679" xyz="0 -0.05 0"/>
188
+ <axis xyz="1 0 0"/>
189
+ <limit effort="80.0" velocity="20."/>
190
+ </joint>
191
+
192
+ <!-- RIGHT HIND LEG -->
193
+ <link name="RH_HAA">
194
+ <visual>
195
+ <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
196
+ <geometry>
197
+ <cylinder length="0.05" radius="0.04"/>
198
+ </geometry>
199
+ </visual>
200
+ <collision>
201
+ <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
202
+ <geometry>
203
+ <cylinder length="0.05" radius="0.04"/>
204
+ </geometry>
205
+ </collision>
206
+ <inertial>
207
+ <origin rpy="0 0 0" xyz="0 0 0.00046"/>
208
+ <mass value="2.04"/>
209
+ <inertia ixx="0.001053013" ixy="4.527e-05" ixz="8.855e-05" iyy="0.001805509" iyz="9.909e-05" izz="0.001765827"/>
210
+ </inertial>
211
+ </link>
212
+ <joint name="RH_HAA" type="revolute">
213
+ <parent link="base"/>
214
+ <child link="RH_HAA"/>
215
+ <origin rpy="0 0 3.1415" xyz="-0.2999 -0.104 0.0"/>
216
+ <axis xyz="1 0 0"/>
217
+ <limit effort="80.0" velocity="20."/>
218
+ </joint>
219
+
220
+ <link name="RH_THIGH">
221
+ <visual>
222
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
223
+ <geometry>
224
+ <cylinder length="0.25" radius="0.02"/>
225
+ </geometry>
226
+ </visual>
227
+ <collision>
228
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
229
+ <geometry>
230
+ <cylinder length="0.25" radius="0.02"/>
231
+ </geometry>
232
+ </collision>
233
+ <inertial>
234
+ <origin rpy="0 0 0" xyz="0 0 -0.125"/>
235
+ <mass value="2.04"/>
236
+ <inertia ixx="0.001053013" ixy="4.527e-05" ixz="8.855e-05" iyy="0.001805509" iyz="9.909e-05" izz="0.001765827"/>
237
+ </inertial>
238
+ </link>
239
+ <joint name="RH_HFE" type="revolute">
240
+ <parent link="RH_HAA"/>
241
+ <child link="RH_THIGH"/>
242
+ <origin rpy="0 0 -1.57079632679" xyz="0 0.05 0"/>
243
+ <axis xyz="1 0 0"/>
244
+ <limit effort="80.0" velocity="20."/>
245
+ </joint>
246
+
247
+ </robot>