tinysim 0.0.2__py3-none-any.whl → 0.0.3__py3-none-any.whl

tinysim/frogger/tk.py

@@ -21,6 +21,8 @@ class FroggerTkFrontend(_tk_base.TkBaseFrontend):
         self.sim_env = sim_env
         self._viewport_size = viewport_size
 
+        self.keys = set()
+
     async def step(self, action, dt=0.01):
         state = self.sim_env.step(action, dt=dt)
         if self._root:
@@ -42,6 +44,10 @@ class FroggerTkFrontend(_tk_base.TkBaseFrontend):
         canvas.pack(fill="both", expand=True)
         self._root = root
         self._canvas = canvas
+
+        root.bind("<KeyPress>", lambda e: self.keys.add(e.keysym))
+        root.bind("<KeyRelease>", lambda e: self.keys.discard(e.keysym))
+
         self.bring_to_front(root)
         self._draw_state(self.sim_env)
         self._pump()

tinysim/topdown_driving/widget.py

@@ -14,7 +14,7 @@ class TopDownDrivingWidget(anywidget.AnyWidget):
 
     sim_state = traitlets.Dict(default_value={}).tag(sync=True)
     wall_positions = traitlets.List(default_value=LOCAL_WALLS).tag(sync=True)
-    
+
     _viewport_size = traitlets.Tuple(
         traitlets.Int(), traitlets.Int(), default_value=(800, 600)
     ).tag(sync=True)
@@ -38,8 +38,7 @@ class TopDownDrivingWidget(anywidget.AnyWidget):
                     ui_poll(100)
         except Exception:
             pass
-        
-    
+
     def copy_py_state(self, sim_state: dict):
         self.sim_state = {
             "x": sim_state["x"].tolist(),
@@ -57,4 +56,4 @@ class TopDownDrivingWidget(anywidget.AnyWidget):
     async def reset(self) -> dict:
         sim_state = self.sim_env.reset()
         self.copy_py_state(sim_state)
-        return sim_state
+        return sim_state

tinysim-0.0.3.dist-info/METADATA

@@ -0,0 +1,80 @@
+Metadata-Version: 2.4
+Name: tinysim
+Version: 0.0.3
+Summary: small modular simulation environments
+Author-email: Matthew Taylor <matthew.taylor.andre@gmail.com>
+Project-URL: Homepage, https://github.com/MatthewAndreTaylor/TinySim
+Keywords: jupyter,simulation
+Classifier: Framework :: Jupyter
+Classifier: Framework :: Jupyter :: JupyterLab
+Classifier: Framework :: Jupyter :: JupyterLab :: 4
+Classifier: Framework :: Jupyter :: JupyterLab :: Extensions
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy
+Requires-Dist: anywidget
+Requires-Dist: jupyter-ui-poll
+Provides-Extra: warp
+Requires-Dist: warp-lang==1.8.1; extra == "warp"
+Provides-Extra: mujoco
+Requires-Dist: mujoco; extra == "mujoco"
+
+
+# TinySim 
+
+[![PyPI](https://img.shields.io/pypi/v/tinysim.svg)](https://pypi.org/project/tinysim)
+[![Jupyterlite](https://jupyterlite.rtfd.io/en/latest/_static/badge.svg)](https://matthewandretaylor.github.io/TinySim/lab?path=example.ipynb)
+
+The goal of this project is to create more minimal simulation environments for robotics and machine learning.
+
+Existing solutions like [ROS 2](https://github.com/ros2) and [Gymnasium](https://gymnasium.farama.org/) come with many heavy dependencies.
+
+Additionally they are usually difficult to get setup and may require domain specific knowledge to use.
+Simulation environments should be able to run in any notebook environment and require fewer dependencies.
+
+
+## Get Started 🚀
+
+To use the baseline Python only simulation environments.
+
+```bash
+pip install tinysim
+```
+
+| [MountainCarEnv](https://github.com/MatthewAndreTaylor/TinySim/blob/main/tinysim/mountain_car) | [TopDownDrivingEnv](https://github.com/MatthewAndreTaylor/TinySim/blob/main/tinysim/topdown_driving) |
+|:-------------------------:|:-------------------------:|
+| <img width="312" alt="mountain_car" src="https://github.com/user-attachments/assets/adbde7f4-6b4e-4c06-8dd3-ea278d5576ad" /> | <img width="312" alt="topdown_driving" src="https://github.com/user-attachments/assets/79dcbf45-055e-4a4e-b6ff-38e938e99a4e" /> |
+
+| [FroggerEnv](https://github.com/MatthewAndreTaylor/TinySim/blob/main/tinysim/frogger) | [FlappyEnv](https://github.com/MatthewAndreTaylor/TinySim/blob/main/tinysim/flappy) |
+|:-------------------------:|:-------------------------:|
+| <img width="312" alt="frogger" src="https://github.com/user-attachments/assets/bb9a8382-b9a0-4280-ae23-cbd02e2d78f6" /> | <img width="312" alt="flappy" src="https://github.com/user-attachments/assets/b5e2a502-d9fe-4ff1-9e4e-a47f7e50ebd9" /> |
+
+
+To use the [warp](https://github.com/NVIDIA/warp) simulation environments.
+
+```bash
+pip install tinysim[warp]
+```
+
+| [CartPoleBaseEnv](https://github.com/MatthewAndreTaylor/TinySim/tree/main/tinysim_warp/cart_pole) | [SimpleRobotDogBaseEnv](https://github.com/MatthewAndreTaylor/TinySim/tree/main/tinysim_warp/simple_quadruped) |
+|:-------------------------:|:-------------------------:|
+| <img width="312" alt="warp_cart_pole" src="https://github.com/user-attachments/assets/f29f3749-d836-4afa-9f27-8578aca020dc" /> | <img width="312" alt="warp_simple_quadruped" src="https://github.com/user-attachments/assets/a966e881-358e-47ee-b0e1-b67eba903987" /> |
+
+
+To use the [mujoco](https://github.com/google-deepmind/mujoco) simulation environments.
+
+```bash
+pip install tinysim[mujoco]
+```
+| [UnitreeA1WalkEnv](https://github.com/MatthewAndreTaylor/TinySim/tree/main/tinysim_mujoco/unitree_a1) |
+|:-------------------------:|
+| <img width="312" alt="mujoco_sim" src="https://github.com/user-attachments/assets/1109488e-0563-4694-8926-682edd0c8278" /> |
+
+

{tinysim-0.0.2.dist-info → tinysim-0.0.3.dist-info}/RECORD

@@ -6,7 +6,7 @@ tinysim/flappy/tk.py,sha256=DUJqy0uPn-AvNE4tuA3a4Dh6xaj0twsUOP8leqGgBX4,2458
 tinysim/flappy/widget.py,sha256=zyaoW-rWgAkmbCY7yrSYyAZDUQaAPnuQEzldSawkOnI,1597
 tinysim/frogger/__init__.py,sha256=gA2jvG_WDZRpZY0Zg_aJmWrdsk2siIL49wjFTRfOSxo,5227
 tinysim/frogger/sim.js,sha256=b1HfEtFZFbbsrZNtltuWO4h-QM4Fw3XjMRydwm5l2NU,3081
-tinysim/frogger/tk.py,sha256=3A1wnWcOBCoRNSI4-oHFbMpRwf0iwdV54GBgtAxVufo,2823
+tinysim/frogger/tk.py,sha256=gb2V8bVCVuVazRrfTbAWq46bK20d5ugHxau_rlNABXE,2992
 tinysim/frogger/widget.py,sha256=Bvb6aUiVqJTHTdd6EzpCp0FqpRa0bl9RNcK12_weiPM,1810
 tinysim/mountain_car/__init__.py,sha256=V3Qxj_ydWAqOJiQUqmIRegBe9GNjU82BfbG2XvV0OQg,1986
 tinysim/mountain_car/sim.js,sha256=z7kxl7LbkRglfzXxmrBbwz2vMYdQV84QP3WhE0hoPVQ,4581
@@ -22,12 +22,24 @@ tinysim/topdown_driving/__init__.py,sha256=HiRpxQNeu4IIOPFraGHLTO2F0VNISAk6bUiK1
 tinysim/topdown_driving/sim.js,sha256=5jmyQk4zp4KFhaG-uTsW_Cq1gHoIkR9d2Cv_ezDKarM,3076
 tinysim/topdown_driving/tk.py,sha256=8UB95ZjIJiUQeXC5w3pXkjukOJXfisLKrfIn0bVFzjU,4751
 tinysim/topdown_driving/track_0.json,sha256=3zPfPX72s5zxs4znoO8rXBqGNciai10Jbl2rFCaACow,15338
-tinysim/topdown_driving/widget.py,sha256=UqhV9Dps8BQTwTFP1g2wlmzql2ooAhAU8-YF8OmpoD8,1734
-tinysim_warp/cart_pole/__init__.py,sha256=TGtvBvQP6lh7ItOwdDyC-tMV-9N3GhgtjAcGvSPpgrk,8344
-tinysim_warp/quadruped/__init__.py,sha256=HlC_DYrpvbTUwvlmBY8FUUQPfmqbPjS1hfNP_bXaljk,6393
-tinysim_warp/simple_quadruped/__init__.py,sha256=sVyQoFVu_EiBNw1P1U9UfSLXNqyIjIks_XF1tyq3Rtk,5329
+tinysim/topdown_driving/widget.py,sha256=Eh7tvtmMbgmzFEvQZ10D6npfF2s1k7QD5AHLYX5Jaqw,1718
+tinysim_mujoco/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinysim_mujoco/unitree_a1/__init__.py,sha256=lxN_94WK_9JdZd3pjZfS_BpvtCY7lmPqQ8KLvaKoLmE,12684
+tinysim_mujoco/unitree_a1/unitree_a1/CHANGELOG.md,sha256=FIdhX4_Q21xamWMS3-2677HhlETlCnfaAzGemqs4Lpo,143
+tinysim_mujoco/unitree_a1/unitree_a1/LICENSE,sha256=K3lKzx4lD1VFrlpH2B5WhepUiM4HfRCH7Yvl-B6bpcA,1559
+tinysim_mujoco/unitree_a1/unitree_a1/README.md,sha256=tHCjy_A1jcGD7bkFy9Qw9z249Kb6HTDQQBlCk8buqlM,1606
+tinysim_mujoco/unitree_a1/unitree_a1/a1.xml,sha256=gJidzvTdw5WwD0X_qrJWjsLwCYmHNqK19fExTtkws4s,9674
+tinysim_mujoco/unitree_a1/unitree_a1/scene.xml,sha256=kvMcNRf1mXFg8CQBp9jrM8B8GCm_TcnQgW68k_iVKis,887
+tinysim_mujoco/unitree_a1/unitree_a1/assets/calf.obj,sha256=duTEMtWXoV4C0q_TFBCxQ_QhodZ7yG1_OALcbzkfCh4,5327177
+tinysim_mujoco/unitree_a1/unitree_a1/assets/hip.obj,sha256=2xx78ZCRU3bnj31Q6iqztiXRtu3QLsLtDQLXweGVqrs,1974009
+tinysim_mujoco/unitree_a1/unitree_a1/assets/thigh.obj,sha256=kUajiusjGyU_X8EyFormznSFUCvGtvrT2K8IhyTGILk,3245342
+tinysim_mujoco/unitree_a1/unitree_a1/assets/thigh_mirror.obj,sha256=A8ynmhTjB9AKowPWTcNk9OYsoWE41Y-XifThe-PdL1g,3354472
+tinysim_mujoco/unitree_a1/unitree_a1/assets/trunk.obj,sha256=wp6C0K3bTH4PArWQV_zMU-m28rFJ1BQMwmRHEqObiHg,5410512
+tinysim_warp/cart_pole/__init__.py,sha256=ltL0cgySOSbeqBvks-UqV1OpT-K-kxo3GylNdvfCB1w,8344
+tinysim_warp/quadruped/__init__.py,sha256=aIOlF9vO9GIYn1y4dth4q-fpUdvB9KZq2rf25YcOcqc,6393
+tinysim_warp/simple_quadruped/__init__.py,sha256=sXL7EOtlB4qm_l4NsYTOFKvpt438voq7OaPgzciPOU0,5329
 tinysim_warp/simple_quadruped/simple_quadruped.urdf,sha256=UfyGROUD5Qlf_t5ISipE2L7uQPBl-hwFnmRsjNL9yH4,8311
-tinysim-0.0.2.dist-info/METADATA,sha256=D0Gt-BPOUwtiYvDosx0Zq_9zusl0Narr2tw9j_UHL80,1942
-tinysim-0.0.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tinysim-0.0.2.dist-info/top_level.txt,sha256=JsBSgxajxod2mlNodBvwkBsn3vMHse0jLqqni2Q5-9Y,21
-tinysim-0.0.2.dist-info/RECORD,,
+tinysim-0.0.3.dist-info/METADATA,sha256=C8c2FG_mLRBcUv3j8Rqd3Z14zHif-BI8G4LUjfvQw2A,3941
+tinysim-0.0.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tinysim-0.0.3.dist-info/top_level.txt,sha256=v-9NmTLfZukVjPwJwEdtCU7yY4TdLsBf1H8s25tR3V4,36
+tinysim-0.0.3.dist-info/RECORD,,

{tinysim-0.0.2.dist-info → tinysim-0.0.3.dist-info}/top_level.txt

@@ -1,2 +1,3 @@
 tinysim
+tinysim_mujoco
 tinysim_warp

tinysim_mujoco/unitree_a1/__init__.py

@@ -0,0 +1,354 @@
+import pathlib
+
+import numpy as np
+
+try:
+    import mujoco
+    import glfw
+except ImportError:
+    raise ImportError(
+        "Mujoco is not properly installed. Install using `pip install tinysim[mujoco]`"
+    )
+
+
+class Viewer:
+    def __init__(self, model, data, width=1200, height=900):
+        self.model = model
+        self.data = data
+        self.width = width
+        self.height = height
+
+        self.window = None
+        self.scene = None
+        self.context = None
+        if not glfw.init():
+            raise RuntimeError("Failed to initialize GLFW")
+
+        self._create_window()
+
+    def _create_window(self):
+        glfw.window_hint(glfw.VISIBLE, glfw.TRUE)
+        self.window = glfw.create_window(self.width, self.height, "Viewer", None, None)
+        if not self.window:
+            raise RuntimeError("Failed to create GLFW window")
+
+        glfw.make_context_current(self.window)
+        glfw.swap_interval(1)
+        glfw.set_window_close_callback(self.window, self._on_close)
+
+        self.cam = mujoco.MjvCamera()
+        self.opt = mujoco.MjvOption()
+        self.scene = mujoco.MjvScene(self.model, maxgeom=500)
+        self.context = mujoco.MjrContext(
+            self.model, mujoco.mjtFontScale.mjFONTSCALE_150
+        )
+
+        self.cam.azimuth = 90.0
+        self.cam.elevation = -25.0
+        self.cam.distance = 4.0
+        self.cam.lookat[:] = np.array([0.0, 0.0, 0.0])
+        self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE
+
+    def render(self):
+        if self.window is None:
+            return
+
+        glfw.make_context_current(self.window)
+        width, height = glfw.get_framebuffer_size(self.window)
+        viewport = mujoco.MjrRect(0, 0, width, height)
+        mujoco.mjv_updateScene(
+            self.model,
+            self.data,
+            self.opt,
+            None,
+            self.cam,
+            mujoco.mjtCatBit.mjCAT_ALL,
+            self.scene,
+        )
+
+        mujoco.mjr_render(viewport, self.scene, self.context)
+        glfw.swap_buffers(self.window)
+        glfw.poll_events()
+
+    def _on_close(self, window):
+        self.close()
+
+    def close(self):
+        if self.window is not None:
+            glfw.destroy_window(self.window)
+            self.window = None
+            self.scene = None
+            self.context = None
+
+
+class UnitreeA1BaseEnv:
+
+    def __init__(self, headless=False):
+        model_path = str(pathlib.Path(__file__).parent / "unitree_a1/scene.xml")
+        self.model = mujoco.MjModel.from_xml_path(model_path)
+        self.data = mujoco.MjData(self.model)
+        self.action_space_size = self.model.nu
+
+        self.default_joint_position = np.array(
+            [-0.05, 0.8, -1.5, 0.05, 0.8, -1.5, -0.03, 0.9, -1.4, 0.03, 0.9, -1.4],
+            dtype=np.float32,
+        )
+
+        if not headless:
+            self.viewer = Viewer(self.model, self.data)
+        else:
+            self.viewer = None
+
+    def step(self, actions, n_frames=20):
+        if self.viewer is not None:
+            self.render()
+
+        if len(actions) != self.action_space_size:
+            raise ValueError(
+                f"Expected {self.action_space_size} actions, but got {len(actions)}"
+            )
+
+        self.data.ctrl[:] = actions
+        mujoco.mj_step(self.model, self.data, nstep=n_frames)
+        mujoco.mj_rnePostConstraint(self.model, self.data)
+
+    def render(self):
+        if self.viewer is None:
+            return
+        self.viewer.render()
+
+    def set_state(self, qpos, qvel):
+        self.data.qpos[:] = qpos
+        self.data.qvel[:] = qvel
+        mujoco.mj_forward(self.model, self.data)
+
+    def close(self):
+        self.viewer.close()
+
+    def reset(self):
+        mujoco.mj_resetData(self.model, self.data)
+        self.data.qpos[0:3] = [0.0, 0.0, 0.3]
+        self.data.qpos[7:19] = self.default_joint_position.copy()
+        self.data.qvel[:] = 0
+        self.data.ctrl[:] = 0
+        mujoco.mj_forward(self.model, self.data)
+
+    def get_position_data(self):
+        return self.data.qpos[0:3].copy()
+
+    def get_joint_data(self):
+        return self.data.qpos[7:19].copy()
+
+    def get_joint_effort(self):
+        return self.data.ctrl[:12].copy()
+
+    def get_velocity_data(self):
+        return self.data.qvel.copy()
+
+
+class UnitreeA1WalkEnv:
+    def __init__(self, reward_weights, cost_weights, headless=False):
+        self.env = UnitreeA1BaseEnv(headless=headless)
+        self.frame_skip = 20
+        self.obs_dim = 45
+        self._previous_observation = np.zeros(self.obs_dim, dtype=np.float32)
+        self.joint_limits_low = np.array(
+            [-0.1, 0.25, -2.0, -0.1, 0.25, -2.0, -0.03, 0.25, -1.6, 0.03, 0.25, -1.6],
+            dtype=np.float32,
+        )
+        self.joint_limits_high = np.array(
+            [0.1, 1.2, -0.5, 0.1, 1.2, -0.5, 0.03, 1.2, -1.1, 0.03, 1.2, -1.1],
+            dtype=np.float32,
+        )
+
+        self._feet_air_time = np.zeros(4)
+        self._last_contacts = np.zeros(4)
+        self.step_counter = 0
+        self.tracking_sigma = 0.25
+        self.maximum_episode_steps = 1024
+        self.prev_action = np.zeros_like(self.env.default_joint_position)
+        self.target_lin_vel = np.array([0.5, 0.0])
+        self.target_ang_vel = 0.0
+        self._healthy_z_range = (0.3, 0.335)
+        self._healthy_rotation = np.array(
+            [np.deg2rad(6), np.deg2rad(6), np.deg2rad(10)]
+        )
+
+        self._cfrc_ext_feet_indices = [4, 7, 10, 13]
+        self.reward_weights = reward_weights
+        self.cost_weights = cost_weights
+        # TODO: requires keys
+
+    def reset(self, **kwargs):
+        self.env.reset()
+        self.step_counter = 0
+        self._feet_air_time = np.zeros(4)
+        self._last_contacts = np.zeros(4)
+        self.prev_action = self.env.default_joint_position.copy()
+        self.target_lin_vel = np.array([0.5, 0.0])
+        current_obs = self._get_obs()
+        self._previous_observation = np.zeros_like(current_obs)
+        full_obs = np.concatenate([current_obs, self._previous_observation])
+        self._previous_observation = current_obs.copy()
+        return full_obs, {}
+
+    def step(self, action):
+        self.step_counter += 1
+        a = 0.7
+        action_filtered = a * self.prev_action + (1 - a) * action
+        if self.step_counter < 1:
+            action_filtered = self.env.default_joint_position.copy()
+        self.env.step(action_filtered, self.frame_skip)
+
+        current_obs = self._get_obs()
+        obs = np.concatenate([current_obs, self._previous_observation])
+        reward = self._compute_reward(action)
+        terminated = not self.is_healthy
+        truncated = self.step_counter >= self.maximum_episode_steps
+        self.prev_action = action_filtered
+        self._previous_observation = current_obs.copy()
+        return obs, reward, terminated, truncated, {}
+
+    def _get_obs(self):
+        base_ang_vel = self.euler_from_quaternion(*self.env.data.qpos[3:7])
+        joint_positions = self.env.get_joint_data()
+        joint_velocities = self.env.get_velocity_data()[6:]
+        curr_contact = self.feet_contact_forces > 1.0
+        return np.concatenate(
+            (
+                base_ang_vel,
+                joint_positions,
+                joint_velocities,
+                curr_contact,
+                self.target_lin_vel,
+                self.prev_action,
+            ),
+            dtype=np.float32,
+        )
+
+    def _compute_reward(self, action):
+        vel = self.env.get_velocity_data()
+        lin_vel = np.array(vel[:3])
+        ang_vel = np.array(vel[3:6])
+
+        tracking_lin_vel_reward, tracking_ang_vel_reward, lin_vel_z_penalty = (
+            self._tracking_velocity_penalty(lin_vel, ang_vel)
+        )
+        action_diff_penalty = np.sum(np.abs(action - self.prev_action))
+        action_sym = self.action_sym()
+
+        Positive_rewards = (
+            tracking_lin_vel_reward * self.reward_weights["linear_vel_tracking"]
+            + tracking_ang_vel_reward * self.reward_weights["angular_vel_tracking"]
+            + self.is_healthy * self.reward_weights["healthy"]
+            + self.feet_air_time_reward * self.reward_weights["feet_airtime"]
+        )
+
+        Negative_rewards = (
+            self.torque_cost * self.cost_weights["torque"]
+            + action_diff_penalty * self.cost_weights["action_rate"]
+            + lin_vel_z_penalty * self.cost_weights["vertical_vel"]
+            + self.xy_angular_velocity_cost * self.cost_weights["xy_angular_vel"]
+            + action_sym * self.cost_weights["action_sym"]
+            + self.acceleration_cost * self.cost_weights["joint_acceleration"]
+            + self.orientation_cost * self.cost_weights["orientation"]
+            + self.default_joint_position_cost
+            * self.cost_weights["default_joint_position"]
+        )
+        reward = Positive_rewards - Negative_rewards
+        return reward
+
+    def _tracking_velocity_penalty(self, lin_vel, ang_vel):
+        lin_vel_error = np.sum(np.abs(self.target_lin_vel[0] - lin_vel[0]))
+        tracking_lin_vel_reward = np.exp(-lin_vel_error / self.tracking_sigma)
+        ang_vel_error = np.sum(np.square(self.target_ang_vel - ang_vel[2]))
+        tracking_ang_vel_reward = np.exp(-ang_vel_error / self.tracking_sigma)
+        lin_vel_z_penalty = np.square((lin_vel[2]))
+        return tracking_lin_vel_reward, tracking_ang_vel_reward, lin_vel_z_penalty
+
+    @property
+    def feet_air_time_reward(self):
+        curr_contact = self.feet_contact_forces > 1.0
+        contact_filter = np.logical_or(curr_contact, self._last_contacts)
+        self._last_contacts = curr_contact
+        first_contact = (self._feet_air_time > 0.0) * contact_filter
+        self._feet_air_time += self.env.model.opt.timestep * self.frame_skip
+        air_time_reward = np.sum((self._feet_air_time) * first_contact)
+        air_time_reward *= np.linalg.norm(self.target_lin_vel) > 0.1
+        self._feet_air_time *= ~contact_filter
+        return air_time_reward
+
+    @property
+    def torque_cost(self):
+        return np.sum(np.abs(self.env.get_joint_effort()))
+
+    @property
+    def xy_angular_velocity_cost(self):
+        return np.sum(np.square(self.env.data.qvel[5]))
+
+    @property
+    def is_healthy(self):
+        _, _, z = self.env.get_position_data()
+        min_z, max_z = self._healthy_z_range
+        if not (min_z <= z <= max_z):
+            return False
+
+        r = self.euler_from_quaternion(*self.env.data.qpos[3:7])
+        return np.all(np.abs(r) <= self._healthy_rotation)
+
+    @property
+    def feet_contact_forces(self):
+        contact_forces = self.env.data.cfrc_ext[self._cfrc_ext_feet_indices].copy()
+        return np.linalg.norm(contact_forces, axis=1)
+
+    @property
+    def orientation_cost(self):
+        roll, pitch, _ = self.euler_from_quaternion(*self.env.data.qpos[3:7])
+        return np.square(roll) + np.square(pitch)
+
+    @property
+    def default_joint_position_cost(self):
+        joint_pos = self.env.get_joint_data()
+        soft_joint_limits_low = self.joint_limits_low * 0.9
+        soft_joint_limits_high = self.joint_limits_high * 0.9
+
+        lower_violation = np.maximum(soft_joint_limits_low - joint_pos, 0)
+        upper_violation = np.maximum(joint_pos - soft_joint_limits_high, 0)
+        return np.sum(np.square(lower_violation + upper_violation))
+
+    def action_sym(self):
+        in_phase_pairs_thigh = [(1, 7), (4, 10)]
+        out_of_phase_pairs_thigh = [(1, 4), (7, 10)]
+        out_of_phase_pairs_calf = [(2, 5), (8, 11)]
+        in_phase_pairs_calf = [(2, 8), (5, 11)]
+
+        jointpositions = self.env.get_joint_data()
+        loss_in_thigh = sum(
+            (jointpositions[i] - jointpositions[j]) ** 2
+            for i, j in in_phase_pairs_thigh
+        )
+        loss_out_thigh = sum(
+            (jointpositions[i] + jointpositions[j] - 1.5) ** 2
+            for i, j in out_of_phase_pairs_thigh
+        )
+        loss_in_calf = sum(
+            (jointpositions[i] - jointpositions[j]) ** 2 for i, j in in_phase_pairs_calf
+        )
+        loss_out_calf = sum(
+            (jointpositions[i] + jointpositions[j] + 3.6) ** 2
+            for i, j in out_of_phase_pairs_calf
+        )
+        loss_in = loss_in_thigh + 0.1 * loss_in_calf
+        loss_out = loss_out_thigh + 0.1 * loss_out_calf
+        return loss_in + loss_out
+
+    @property
+    def acceleration_cost(self):
+        return np.sum(np.square(self.env.data.qacc[6:]))
+
+    @staticmethod
+    def euler_from_quaternion(w, x, y, z):
+        roll = np.arctan2(2 * (w * x + y * z), 1 - 2 * (x * x + y * y))
+        pitch = np.arcsin(np.clip(2 * (w * y - z * x), -1.0, 1.0))
+        yaw = np.arctan2(2 * (w * z + x * y), 1 - 2 * (y * y + z * z))
+        return roll, pitch, yaw

tinysim_mujoco/unitree_a1/unitree_a1/CHANGELOG.md

@@ -0,0 +1,6 @@
+# Changelog – Unitree A1 Description
+
+All notable changes to this model will be documented in this file.
+
+## [2022-09-07]
+- Initial release.

tinysim_mujoco/unitree_a1/unitree_a1/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2016-2022 HangZhou YuShu TECHNOLOGY CO.,LTD. ("Unitree Robotics")
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

tinysim_mujoco/unitree_a1/unitree_a1/README.md

@@ -0,0 +1,41 @@
+# Unitree A1 Description (MJCF)
+
+> [!IMPORTANT]
+> Requires MuJoCo 2.2.2 or later.
+
+## Changelog
+
+See [CHANGELOG.md](./CHANGELOG.md) for a full history of changes.
+
+## Overview
+
+This package contains a simplified robot description (MJCF) of the [A1 Quadruped
+Robot](https://www.unitree.com/a1/) developed by [Unitree
+Robotics](https://www.unitree.com/). It is derived from the [publicly available
+URDF
+description](https://github.com/unitreerobotics/unitree_mujoco/tree/main/data/a1/urdf).
+
+<p float="left">
+  <img src="a1.png" width="400">
+</p>
+
+## URDF → MJCF derivation steps
+
+1. Converted the DAE [mesh
+   files](https://github.com/unitreerobotics/unitree_mujoco/tree/main/data/a1/meshes)
+to OBJ format using [Blender](https://www.blender.org/).
+2. Processed `.obj` files with [`obj2mjcf`](https://github.com/kevinzakka/obj2mjcf).
+3. Added `<mujoco> <compiler discardvisual="false"/> </mujoco>` to the URDF's
+   `<robot>` clause in order to preserve visual geometries.
+4. Loaded the URDF into MuJoCo and saved a corresponding MJCF.
+5. Added a `<freejoint/>` to the base, and a tracking light.
+6. Manually edited the MJCF to extract common properties into the `<default>` section.
+7. Manually designed collision geometries.
+8. Shifted joint reference values and ranges so that 0 corresponds to standing pose.
+9. Softened the contacts of the feet to approximate the effect of rubber and
+   increased `impratio` to reduce slippage.
+10. Added `scene.xml` which includes the robot, with a textured groundplane, skybox, and haze.
+
+## License
+
+This model is released under a [BSD-3-Clause License](LICENSE).