tinysim 0.0.4__py3-none-any.whl → 0.0.5__py3-none-any.whl

tinysim_warp/simple_quadruped/__init__.py

@@ -1,34 +1,17 @@
-import math
-import os
 import pathlib
 
 import numpy as np
+import warp as wp
 
-try:
-    import warp as wp
-    import warp.sim.render
-except ImportError:
-    raise ImportError(
-        "Warp is not installed. Install using `pip install tinysim[warp]`"
-    )
-
-
-def compute_env_offsets(num_envs, env_offset=(5.0, 0.0, 0.0)):
-    env_offset = np.array(env_offset, dtype=float)
-    axis = np.nonzero(env_offset)[0]
-    axis = axis[0] if len(axis) else 0
-    env_offsets = np.zeros((num_envs, 3))
-    env_offsets[:, axis] = np.arange(num_envs) * env_offset[axis]
-    env_offsets -= env_offsets.mean(axis=0)
-    env_offsets[:, 1] = 0.0
-    return env_offsets
+from tinysim_warp import WarpBaseEnv
 
 
-class SimpleRobotDogBaseEnv:
-    def __init__(self, use_cuda_graph=False, headless=False, num_envs=8):
+class SimpleRobotDogBaseEnv(WarpBaseEnv):
+    def build_model(self):
+        builder = wp.sim.ModelBuilder()
         articulation_builder = wp.sim.ModelBuilder()
-        rot_x = wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)
-        rot_y = wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), math.pi * 0.5)
+        rot_x = wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -np.pi * 0.5)
+        rot_y = wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), np.pi * 0.5)
         xform = wp.transform(wp.vec3(0.0, 0.35, 0.0), rot_y * rot_x)
         wp.sim.parse_urdf(
             # os.path.join(warp.examples.get_asset_directory(), "quadruped.urdf"),
@@ -48,21 +31,10 @@ class SimpleRobotDogBaseEnv:
             limit_kd=1.0e1,
         )
 
-        builder = wp.sim.ModelBuilder()
-
-        self.sim_time = 0.0
-        fps = 100
-        self.frame_dt = 1.0 / fps
-
-        self.sim_substeps = 10
-        self.sim_dt = self.frame_dt / self.sim_substeps
-
-        self.num_envs = num_envs
-        offsets = compute_env_offsets(self.num_envs)
-
         for i in range(self.num_envs):
             builder.add_builder(
-                articulation_builder, xform=wp.transform(offsets[i], wp.quat_identity())
+                articulation_builder,
+                xform=wp.transform(self.env_offsets[i], wp.quat_identity()),
             )
             builder.joint_q[-8:] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] # fmt: skip
             builder.joint_act[-8:] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] # fmt: skip
@@ -70,50 +42,18 @@ class SimpleRobotDogBaseEnv:
                 builder.joint_axis_mode
             )
 
-        # finalize model
-        self.model = builder.finalize()
-        self.model.ground = True
-        self.model.joint_attach_ke = 16000.0
-        self.model.joint_attach_kd = 200.0
-        self.use_tile_gemm = False
-        self.fuse_cholesky = False
-
-        self.integrator = wp.sim.FeatherstoneIntegrator(
-            self.model,
-            use_tile_gemm=self.use_tile_gemm,
-            fuse_cholesky=self.fuse_cholesky,
-        )
-
-        self.renderer = wp.sim.render.SimRendererOpenGL(
-            self.model, "example", headless=headless
-        )
-        self.state_0 = self.model.state()
-        self.state_1 = self.model.state()
-
-        wp.sim.eval_fk(
-            self.model, self.model.joint_q, self.model.joint_qd, None, self.state_0
-        )
+        model = builder.finalize()
+        model.joint_attach_ke = 16000.0
+        model.joint_attach_kd = 200.0
+        return model
 
-        self.use_cuda_graph = (
-            use_cuda_graph
-            and wp.get_device().is_cuda
-            and wp.is_mempool_enabled(wp.get_device())
+    def solver(self):
+        return wp.sim.FeatherstoneIntegrator(
+            self.model, use_tile_gemm=False, fuse_cholesky=False
         )
-        if self.use_cuda_graph:
-            with wp.ScopedCapture() as capture:
-                self.simulate()
-            self.graph = capture.graph
-
-    def simulate(self):
-        for _ in range(self.sim_substeps):
-            self.state_0.clear_forces()
-            wp.sim.collide(self.model, self.state_0)
-            self.integrator.simulate(
-                self.model, self.state_0, self.state_1, self.sim_dt
-            )
-            self.state_0, self.state_1 = self.state_1, self.state_0
 
-    def set_leg_poses(self, actions):
+    def apply_actions(self, actions):
+        # Set the target positions for all legs
         action_space_size = 8 * self.num_envs
         if len(actions) != action_space_size:
             raise ValueError(
@@ -122,19 +62,8 @@ class SimpleRobotDogBaseEnv:
 
         wp.copy(self.model.joint_act, wp.array(actions, dtype=wp.float32))
 
-    def step(self, actions):
-        self.set_leg_poses(actions)
-
-        if self.use_cuda_graph:
-            wp.capture_launch(self.graph)
-        else:
-            self.simulate()
-        self.sim_time += self.frame_dt
-
-    def render(self):
-        self.renderer.begin_frame(self.sim_time)
-        self.renderer.render(self.state_0)
-        self.renderer.end_frame()
+    def get_obs(self):
+        return
 
 
 if __name__ == "__main__":

tinysim_warp/simple_quadruped/simple_quadruped.urdf

@@ -1,7 +1,5 @@
 <?xml version="1.0" encoding="utf-8"?>
 <robot name="quadruped">
-
-    <!-- Base -->
     <link name="base">
         <visual>
             <origin rpy="0 1.57079632679 0" xyz="0 0 0"/>
@@ -22,7 +20,6 @@
         </inertial>
     </link>
 
-    <!-- LEFT FRONT LEG -->
     <link name="LF_HAA">
         <visual>
             <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
@@ -77,9 +74,6 @@
         <limit effort="80.0" velocity="20."/>
     </joint>
 
-
-
-    <!-- RIGHT FRONT LEG -->
     <link name="RF_HAA">
         <visual>
             <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
@@ -134,7 +128,6 @@
         <limit effort="80.0" velocity="20."/>
     </joint>
 
-    <!-- LEFT HIND LEG -->
     <link name="LH_HAA">
         <visual>
             <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
@@ -189,7 +182,6 @@
         <limit effort="80.0" velocity="20."/>
     </joint>
 
-    <!-- RIGHT HIND LEG -->
     <link name="RH_HAA">
         <visual>
             <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
@@ -243,5 +235,4 @@
         <axis xyz="1 0 0"/>
         <limit effort="80.0" velocity="20."/>
     </joint>
-
 </robot>

tinysim/simple_amr/example_maps.py

@@ -1,121 +0,0 @@
-import random
-import math
-import copy
-
-empty_map = {
-    "name": "rect_triangle_map",
-    "robot": {
-        "pos": [500, 300],
-        "angle": 0,
-    },  # speed: 0.01, turn_speed: 0.04, lidar: { numBeams: 30, fov: 6.28, maxRange: 1000}
-    "map": [
-        {
-            "type": "rectangle",
-            "x": 400,
-            "y": 0,
-            "width": 800,
-            "height": 30,
-            "bodyInfo": {"angle": 0, "render": {"fillStyle": "#2f6f8f"}},
-            "label": "Wall",
-        },
-        {
-            "type": "rectangle",
-            "x": 400,
-            "y": 590,
-            "width": 800,
-            "height": 20,
-            "bodyInfo": {"angle": 0, "render": {"fillStyle": "#2f6f8f"}},
-            "label": "Wall",
-        },
-        {
-            "type": "rectangle",
-            "x": 10,
-            "y": 500,
-            "width": 20,
-            "height": 1000,
-            "bodyInfo": {"angle": 0, "render": {"fillStyle": "#2f6f8f"}},
-            "label": "Wall",
-        },
-        {
-            "type": "rectangle",
-            "x": 790,
-            "y": 500,
-            "width": 20,
-            "height": 1000,
-            "bodyInfo": {"angle": 0, "render": {"fillStyle": "#2f6f8f"}},
-            "label": "Wall",
-        },
-    ],
-}
-
-
-def gen_simple_map():
-    map_data = copy.deepcopy(empty_map)
-    width = 800
-    height = 600
-
-    # Generate static obstacles
-    for _ in range(6):
-        w = 60 + random.random() * 120
-        h = 40 + random.random() * 100
-        x = 120 + random.random() * (width - 240)
-        y = 120 + random.random() * (height - 240)
-        angle = random.random() * math.pi
-        rect = {
-            "type": "rectangle",
-            "x": x,
-            "y": y,
-            "width": w,
-            "height": h,
-            "bodyInfo": {
-                "angle": angle,
-                "render": {"fillStyle": "#2f6f8f"},
-                "label": "Obstacle",
-            },
-        }
-        map_data["map"].append(rect)
-
-    # Generate dynamic pushable boxes
-    for _ in range(3):
-        size = 40 + random.random() * 20
-        x = 100 + random.random() * (width - 200)
-        y = 100 + random.random() * (height - 200)
-        box = {
-            "type": "rectangle",
-            "x": x,
-            "y": y,
-            "width": size,
-            "height": size,
-            "bodyInfo": {
-                "isStatic": False,
-                "frictionAir": 0.1,
-                "friction": 0.5,
-                "restitution": 0.2,
-                "density": 0.002,
-                "render": {"fillStyle": "#b2df8a"},
-                "label": "Pushable Box",
-            },
-        }
-        map_data["map"].append(box)
-
-    for i in range(2):
-        size = 40 + random.random() * 20
-        x = 100 + random.random() * (width - 200)
-        y = 100 + random.random() * (height - 200)
-        angle = random.random() * math.pi
-        box = {
-            "type": "rectangle",
-            "x": x,
-            "y": y,
-            "width": size,
-            "height": size,
-            "bodyInfo": {
-                "isStatic": True,
-                "angle": angle,
-                "label": "Goal",
-                "render": {"fillStyle": "#ffffff"},
-            },
-        }
-        map_data["map"].append(box)
-
-    return map_data