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

tinysim_warp/__init__.py

@@ -0,0 +1,108 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+
+try:
+    import warp as wp
+    import warp.sim.render
+except ImportError:
+    raise ImportError(
+        "Warp is not installed. Install using `pip install tinysim[warp]`"
+    )
+
+
+class WarpBaseEnv(ABC):
+    def __init__(self, use_cuda_graph=False, headless=False, num_envs=1, fps=60):
+        self.try_use_cuda_graph = use_cuda_graph
+        self.headless = headless
+        self.num_envs = num_envs
+
+        self.fps = fps
+        self.sim_time = 0.0
+        self.frame_dt = 1.0 / fps
+        self.sim_substeps = 10  # todo: make configurable
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.num_envs = num_envs
+        self.env_offsets = self.comp_env_offsets()
+        self.model = self.build_model()
+        self.integrator = self.solver()
+
+        self.renderer = wp.sim.render.SimRendererOpenGL(
+            self.model, "viewer", headless=headless
+        )
+        self.state_0 = self.model.state()
+        self.state_1 = self.model.state()
+        wp.sim.eval_fk(
+            self.model, self.state_0.joint_q, self.state_0.joint_qd, None, self.state_0
+        )
+        self.capture_graph()
+
+    @abstractmethod
+    def build_model(self):
+        raise NotImplementedError("build_model() must be implemented in subclass.")
+
+    @abstractmethod
+    def solver(self):
+        raise NotImplementedError("solver() must be implemented in subclass.")
+
+    @abstractmethod
+    def apply_actions(self, actions):
+        raise NotImplementedError("apply_actions() must be implemented in subclass.")
+
+    @abstractmethod
+    def get_obs(self):
+        raise NotImplementedError("get_obs() must be implemented in subclass.")
+
+    def comp_env_offsets(self, offset=(5.0, 0.0, 0.0)):
+        offset = np.array(offset, dtype=float)
+        axis = np.nonzero(offset)[0]
+        axis = axis[0] if len(axis) else 0
+        t = np.arange(self.num_envs) * offset[axis]
+        offsets = np.zeros((self.num_envs, 3))
+        offsets[:, axis] = t - t.mean()
+        return offsets
+
+    def capture_graph(self):
+        # CUDA graph
+        self.use_cuda_graph = (
+            self.try_use_cuda_graph
+            and wp.get_device().is_cuda
+            and wp.is_mempool_enabled(wp.get_device())
+        )
+        if self.try_use_cuda_graph and not self.use_cuda_graph:
+            raise Warning("""Unable to enable CUDA graph capture.
+                Please verify that Warp is using CUDA and that the memory pool is enabled.
+                """)
+
+        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 step(self, actions):
+        self.apply_actions(actions)
+
+        if self.use_cuda_graph:
+            wp.capture_launch(self.graph)
+        else:
+            self.simulate()
+
+        self.sim_time += self.frame_dt
+        return self.get_obs()
+
+    def render(self):
+        self.renderer.begin_frame(self.sim_time)
+        self.renderer.render(self.state_0)
+        self.renderer.end_frame()
+
+    def reset(self, **kwargs):
+        raise NotImplementedError("reset() must be implemented in subclass.")

tinysim_warp/cart_pole/__init__.py

@@ -1,94 +1,67 @@
-# written by: Matthew Taylor
-# source: https://github.com/MatthewAndreTaylor/WarpSimBalance
-
-try:
-    import warp as wp
-    import warp.sim.render
-except ImportError:
-    raise ImportError(
-        "Warp is not installed. Install using `pip install tinysim[warp]`"
-    )
-
 import numpy as np
+import warp as wp
 
+from tinysim_warp import WarpBaseEnv
 
-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
-
+quat_id = wp.quat_identity()
 
-class CartPoleBaseEnv:
-    def __init__(self, use_cuda_graph=False, headless=False, num_envs=3):
-        self.actions = np.array([0.0, 0.5, -0.5])
-        self.num_envs = num_envs
 
-        fps = 60
-        self.frame_dt = 1.0 / fps
-        self.sim_substeps = 10
-        self.sim_dt = self.frame_dt / self.sim_substeps
+class CartPoleBaseEnv(WarpBaseEnv):
 
-        self.initialize()
+    actions = np.array([0.0, 0.5, -0.5])
 
-        self.integrator = wp.sim.SemiImplicitIntegrator()
-        self.renderer = wp.sim.render.SimRendererOpenGL(
-            self.model, "example", headless=headless
-        )
+    def build_model(self):
+        builder = wp.sim.ModelBuilder()
+        for i in range(self.num_envs):
+            builder.add_builder(
+                self.create_cartpole(),
+                xform=wp.transform(self.env_offsets[i], quat_id),
+            )
 
-        # CUDA graph
-        self.use_cuda_graph = (
-            use_cuda_graph
-            and wp.get_device().is_cuda
-            and wp.is_mempool_enabled(wp.get_device())
+        model = builder.finalize()
+        model.joint_attach_ke = 1000.0
+        model.joint_attach_kd = 1.0
+        self.joint_act_np = np.zeros(model.joint_dof_count, dtype=np.float32)
+        self.joint_act_wp = wp.array(
+            self.joint_act_np, dtype=float, device=wp.get_device()
         )
-        if self.use_cuda_graph:
-            with wp.ScopedCapture() as capture:
-                self.simulate()
-            self.graph = capture.graph
-
-        self.step(actions=[1] * self.num_envs)
+        model.joint_act = self.joint_act_wp
+        return model
 
     def create_cartpole(self):
         """Create cartpole system using pure Python/Warp API"""
         builder = wp.sim.ModelBuilder(gravity=-0.3)
-        pole_size = wp.vec3(0.04, 1.0, 0.06)
         cart_body = builder.add_body(
-            origin=wp.transform(wp.vec3(0.0, 2.0, 0.0), wp.quat_identity()), m=1.0
+            origin=wp.transform(wp.vec3(0.0, 2.0, 0.0), quat_id), m=1.0
         )
         builder.add_shape_sphere(body=cart_body, radius=0.1, density=100.0)
+
+        # Start angled pole
         pole_body = builder.add_body(
-            origin=wp.transform(wp.vec3(0.0, 2.5, 0.0), wp.quat_identity()), m=0.01
+            origin=wp.transform(wp.vec3(0.0, 2.5, 0.0), quat_id), m=0.01
         )
+        pole_size = wp.vec3(0.04, 1.0, 0.06)
         builder.add_shape_box(
             body=pole_body,
-            pos=wp.vec3(0.0, 0.0, 0.0),
             hx=pole_size[0] / 2.0,
             hy=pole_size[1] / 2.0,
             hz=pole_size[2] / 2.0,
             density=50.0,
         )
+        # Joint starts with a small initial angle
+        q = wp.quat_from_axis_angle(wp.vec3(0.0, 0.0, 1.0), 0.01)
         builder.add_joint_revolute(
             parent=cart_body,
             child=pole_body,
             axis=wp.vec3(0.0, 0.0, 1.0),  # rotation around Z-axis
-            parent_xform=wp.transform(wp.vec3(0.0, 0.0, 0.0), wp.quat_identity()),
-            child_xform=wp.transform(
-                wp.vec3(0.0, -0.5, 0.0), wp.quat_identity()
-            ),  # joint at pole base
+            child_xform=wp.transform(wp.vec3(0.0, -0.5, 0.0), q),
             limit_ke=1.0e4,
             limit_kd=1.0e1,
         )
         builder.add_joint_prismatic(
             parent=-1,
             child=cart_body,
-            axis=wp.vec3(1.0, 0.0, 0.0),
-            parent_xform=wp.transform(wp.vec3(0.0, 2.0, 0.0), wp.quat_identity()),
-            child_xform=wp.transform(wp.vec3(0.0, 0.0, 0.0), wp.quat_identity()),
+            parent_xform=wp.transform(wp.vec3(0.0, 2.0, 0.0), quat_id),
             limit_lower=-5.0,
             limit_upper=5.0,
             limit_ke=1.0e4,
@@ -99,7 +72,10 @@ class CartPoleBaseEnv:
         )
         return builder
 
-    def set_cart_trajectory(self, actions):
+    def solver(self):
+        return wp.sim.SemiImplicitIntegrator()
+
+    def apply_actions(self, actions):
         if len(actions) != self.num_envs:
             raise ValueError("Length of actions must match number of environments.")
 
@@ -108,152 +84,20 @@ class CartPoleBaseEnv:
         act[dof_indices] = self.actions[actions]
         wp.copy(self.model.joint_act, wp.array(act, dtype=wp.float32))
 
-    def is_fallen(self, pole_quat):
-        qw = float(np.clip(pole_quat[3], -1.0, 1.0))
-        tilt = 2.0 * np.arccos(qw)  # radians, in [0, pi]
-        max_tilt = np.deg2rad(90.0)
-        return tilt > max_tilt
-
-    def step(self, actions):
-        """Apply action and step the simulation for one environment timestep.
-
-        Returns (observation, reward, terminated)
-        """
-        self.set_cart_trajectory(actions)
-        if self.use_cuda_graph:
-            wp.capture_launch(self.graph)
-        else:
-            self.simulate()
-
-        self.sim_time += self.frame_dt
-        obs = self.get_state_vector()
-        pole_quat = obs[1:5]
-        terminated = self.is_fallen(pole_quat)
-        return obs, 1.0, terminated
-
-    def simulate(self):
-        for _ in range(self.sim_substeps):
-            self.state.clear_forces()
-            self.state = self.integrator.simulate(
-                self.model, self.state, self.state, self.sim_dt
-            )
-
-    def reset(self):
-        self.initialize()
-        return self.step(actions=[1] * self.num_envs)
-
-    def initialize(self):
-        self.sim_time = 0.0
-        builder = wp.sim.ModelBuilder()
-        offsets = compute_env_offsets(self.num_envs)
-
-        for i in range(self.num_envs):
-            builder.add_builder(
-                self.create_cartpole(),
-                xform=wp.transform(offsets[i], wp.quat_identity()),
-            )
-
-        self.model = builder.finalize()
-        self.model.joint_attach_ke = 1000.0
-        self.model.joint_attach_kd = 1.0
-        self.state = self.model.state()
-
-        wp.sim.eval_fk(
-            self.model, self.model.joint_q, self.model.joint_qd, None, self.state
-        )
-        self.joint_act_np = np.zeros(self.model.joint_dof_count, dtype=np.float32)
-        self.joint_act_wp = wp.array(
-            self.joint_act_np, dtype=float, device=wp.get_device()
-        )
-        self.model.joint_act = self.joint_act_wp
-
-    def render(self):
-        self.renderer.begin_frame(self.sim_time)
-        self.renderer.render(self.state)
-        self.renderer.end_frame()
-
     def get_state_vector(self):
-        cart_id = 0
-        pole_id = 1
-
-        # convert transforms to numpy
-        body_q = self.state.body_q.numpy()
-        body_qd = self.state.body_qd.numpy()
-
-        # --- Cart ---
+        cart_id, pole_id = 0, 1  # body indices
+        body_q = self.state_0.body_q.numpy()
+        body_qd = self.state_0.body_qd.numpy()
         cart_pos = body_q[cart_id]  # [px, py, pz, qx, qy, qz, qw]
-        # cart_pos = cart_pos[[0, 2]]
-        cart_pos = cart_pos[[0]]
-
-        # --- Pole ---
+        cart_pos = cart_pos[[0]]  # cart_pos = cart_pos[[0, 2]]
         pole_quat = body_q[pole_id][3:]  # quaternion part
         pole_vel = body_qd[pole_id][3:]  # angular velocity part
-
         return np.concatenate([cart_pos, pole_quat, pole_vel])
 
-
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser(
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter
-    )
-    parser.add_argument(
-        "--device", type=str, default=None, help="Override the default Warp device."
-    )
-    parser.add_argument(
-        "--num-frames", type=int, default=1200, help="Total number of frames."
-    )
-    parser.add_argument(
-        "--manual-control", action="store_true", help="Enable manual control."
-    )
-
-    args = parser.parse_known_args()[0]
-
-    if args.manual_control:
-        print("Manual control enabled")
-        import keyboard
-
-    with wp.ScopedDevice(args.device):
-        example = CartPoleBaseEnv(use_cuda_graph=True)
-
-        terminated = False
-        check_terminated = True
-
-        for i in range(args.num_frames):
-            if args.manual_control:
-                # Manual control logic
-                if keyboard.is_pressed("k"):
-                    action = 2
-                elif keyboard.is_pressed("l"):
-                    action = 1
-                else:
-                    action = 0
-
-            else:
-                # Example control signals
-                if i < 200:
-                    action = 0
-                elif i < 255:  # 260
-                    action = 2
-                else:
-                    action = 0
-
-            obs, reward, terminated = example.step(actions=[action] * example.num_envs)
-            # print(f"Step {i}")
-
-            if check_terminated and terminated:
-                print("Pole fallen")
-                check_terminated = False
-
-            example.render()
-
-            # print("State Vector:", obs)
-            # print("Reward:", reward)
-            # print("Terminated:", terminated)
-
-        state = example.reset()
-        print("Reset State Vector:", state)
-        example.render()
-
-        example.renderer.save()
+    def get_obs(self):
+        obs = self.get_state_vector()
+        pole_quat = obs[1:5]
+        qw = float(np.clip(pole_quat[3], -1.0, 1.0))
+        tilt = 2.0 * np.arccos(qw)  # radians, in [0, pi]
+        max_tilt = np.deg2rad(90.0)
+        return obs, 1.0, tilt > max_tilt

tinysim_warp/quadruped/__init__.py

@@ -1,37 +1,21 @@
-import math
-import os
+import pathlib
 
 import numpy as np
+import warp as wp
+import warp.examples
 
-try:
-    import warp as wp
-    import warp.examples
-    import warp.sim.render
-except ImportError:
-    raise ImportError(
-        "Warp is not installed. Install using `pip install tinysim[warp]`"
-    )
+from tinysim_warp import WarpBaseEnv
 
 
-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
-
-
-class RobotDogBaseEnv:
-    def __init__(self, use_cuda_graph=False, headless=False, num_envs=8):
+class RobotDogBaseEnv(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.65, 0.0), rot_y * rot_x)
         wp.sim.parse_urdf(
-            os.path.join(warp.examples.get_asset_directory(), "quadruped.urdf"),
+            str(pathlib.Path(warp.examples.get_asset_directory()) / "quadruped.urdf"),
             articulation_builder,
             xform=xform,
             floating=True,
@@ -47,21 +31,10 @@ class RobotDogBaseEnv:
             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[-12:] = [0.2, 0.4, -0.6, -0.2, -0.4, 0.6, -0.2, 0.4, -0.6, 0.2, -0.4, 0.6] # fmt: skip
             builder.joint_act[-12:] = [0.2, 0.4, -0.6, -0.2, -0.4, 0.6, -0.2, 0.4, -0.6, 0.2, -0.4, 0.6] # fmt: skip
@@ -69,50 +42,18 @@ class RobotDogBaseEnv:
                 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
+        model = builder.finalize()
+        model.joint_attach_ke = 16000.0
+        model.joint_attach_kd = 200.0
+        return model
 
-        self.integrator = wp.sim.FeatherstoneIntegrator(
-            self.model,
-            use_tile_gemm=self.use_tile_gemm,
-            fuse_cholesky=self.fuse_cholesky,
+    def solver(self):
+        return wp.sim.FeatherstoneIntegrator(
+            self.model, use_tile_gemm=False, fuse_cholesky=False
         )
 
-        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
-        )
-
-        self.use_cuda_graph = (
-            use_cuda_graph
-            and wp.get_device().is_cuda
-            and wp.is_mempool_enabled(wp.get_device())
-        )
-        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 = 12 * self.num_envs
         if len(actions) != action_space_size:
             raise ValueError(
@@ -121,83 +62,5 @@ class RobotDogBaseEnv:
 
         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()
-
-
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser(
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter
-    )
-    parser.add_argument(
-        "--device", type=str, default=None, help="Override the default Warp device."
-    )
-    parser.add_argument(
-        "--num-frames", type=int, default=300, help="Total number of frames."
-    )
-    parser.add_argument(
-        "--num-envs",
-        type=int,
-        default=3,
-        help="Total number of simulated environments.",
-    )
-
-    args = parser.parse_known_args()[0]
-
-    with wp.ScopedDevice(args.device):
-        example = RobotDogBaseEnv(num_envs=args.num_envs)
-        # [
-        #     (front right leg out),
-        #     (front right leg forward),
-        #     (front right forearm forward),
-
-        #     (front left leg out),
-        #     (front left leg forward),
-        #     (front left forearm forward),
-
-        #     (back right leg out),
-        #     (back right leg forward),
-        #     (back right forearm forward),
-
-        #     (back left leg out),
-        #     (back left leg forward),
-        #     (back left forearm forward),
-        # ]
-        # left and right directions are reversed
-
-        actions = [0.2, 0.4, -0.6, -0.2, -0.4, 0.6, -0.2, 0.4, -0.6, 0.2, -0.4, 0.6] * args.num_envs # fmt: skip
-
-        for i in range(args.num_frames):
-            example.step(actions)
-            example.render()
-
-            if i % 10 == 0:
-                # Stretch the right leg to the side
-                # actions[0] += 0.1
-
-                # Stretch the right leg forward
-                actions[1] += 0.1
-
-                # Strecth the forearm foward
-                # actions[2] += 0.1
-
-                # Stretch the left leg forward
-                actions[4] -= 0.1
-
-                # Back right leg backwards
-                # actions[7] += 0.1
-
-        example.renderer.save()
+    def get_obs(self):
+        return