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

tinysim_mujoco/manipulation/xmls/panda.xml

@@ -0,0 +1,280 @@
+<mujoco model="panda">
+  <compiler angle="radian" meshdir="assets" autolimits="true"/>
+
+  <default>
+    <default class="panda">
+      <material specular="0.5" shininess="0.25"/>
+      <joint range="-2.8973 2.8973"/>
+      <default class="finger">
+        <joint axis="0 1 0" type="slide" range="0 0.04"/>
+      </default>
+
+      <default class="visual">
+        <geom type="mesh" contype="0" conaffinity="0" group="2"/>
+      </default>
+      <default class="collision">
+        <geom type="mesh" group="3"/>
+          <default class="fingertip_collision_1">
+            <geom type="box" size="0.011 0.0075 0.01" pos="0 0.0185 0.011"/>
+          </default>
+          <default class="fingertip_collision_2">
+            <geom type="box" size="0.011 0.0044 0.0019" pos="0 0.0068 0.0022"/>
+          </default>
+          <default class="fingertip_collision_3">
+            <geom type="box" size="0.00875 0.0035 0.01175" pos="0 0.0159 0.02835" quat="1 0.25 0 0"/>
+          </default>
+          <default class="fingertip_pad_collision_1">
+            <geom type="box" size="0.0085 0.004 0.0085" pos="0 0.0055 0.0445"/>
+          </default>
+          <default class="fingertip_pad_collision_2">
+            <geom type="box" size="0.003 0.002 0.003" pos="0.0055 0.002 0.05"/>
+          </default>
+          <default class="fingertip_pad_collision_3">
+            <geom type="box" size="0.003 0.002 0.003" pos="-0.0055 0.002 0.05"/>
+          </default>
+          <default class="fingertip_pad_collision_4">
+            <geom type="box" size="0.003 0.002 0.0035" pos="0.0055 0.002 0.0395"/>
+          </default>
+        <default class="fingertip_pad_collision_5">
+          <geom type="box" size="0.003 0.002 0.0035" pos="-0.0055 0.002 0.0395"/>
+        </default>
+        <default class="fingertip_pad_collision_6">
+          <geom type="box" size="0.00875 0.0076 0.00825" pos="0 0.00758 0.04525"/>
+        </default>
+      </default>
+    </default>
+  </default>
+
+  <asset>
+    <material class="panda" name="white" rgba="1 1 1 1"/>
+    <material class="panda" name="off_white" rgba="0.901961 0.921569 0.929412 1"/>
+    <material class="panda" name="black" rgba="0.25 0.25 0.25 1"/>
+    <material class="panda" name="green" rgba="0 1 0 1"/>
+    <material class="panda" name="light_blue" rgba="0.039216 0.541176 0.780392 1"/>
+
+    <!-- Collision meshes -->
+    <mesh name="link0_collision" file="link0.stl"/>
+    <mesh name="link1_collision" file="link1.stl"/>
+    <mesh name="link2_collision" file="link2.stl"/>
+    <mesh name="link3_collision" file="link3.stl"/>
+    <mesh name="link4_collision" file="link4.stl"/>
+    <mesh name="link5_collision_0" file="link5_collision_0.obj"/>
+    <mesh name="link5_collision_1" file="link5_collision_1.obj"/>
+    <mesh name="link5_collision_2" file="link5_collision_2.obj"/>
+    <mesh name="link6_collision" file="link6.stl"/>
+    <mesh name="link7_collision" file="link7.stl"/>
+    <mesh name="hand_collision" file="hand.stl"/>
+
+    <!-- Visual meshes -->
+    <mesh file="link0_0.obj"/>
+    <mesh file="link0_1.obj"/>
+    <mesh file="link0_2.obj"/>
+    <mesh file="link0_3.obj"/>
+    <mesh file="link0_4.obj"/>
+    <mesh file="link0_5.obj"/>
+    <mesh file="link0_7.obj"/>
+    <mesh file="link0_8.obj"/>
+    <mesh file="link0_9.obj"/>
+    <mesh file="link0_10.obj"/>
+    <mesh file="link0_11.obj"/>
+    <mesh file="link1.obj"/>
+    <mesh file="link2.obj"/>
+    <mesh file="link3_0.obj"/>
+    <mesh file="link3_1.obj"/>
+    <mesh file="link3_2.obj"/>
+    <mesh file="link3_3.obj"/>
+    <mesh file="link4_0.obj"/>
+    <mesh file="link4_1.obj"/>
+    <mesh file="link4_2.obj"/>
+    <mesh file="link4_3.obj"/>
+    <mesh file="link5_0.obj"/>
+    <mesh file="link5_1.obj"/>
+    <mesh file="link5_2.obj"/>
+    <mesh file="link6_0.obj"/>
+    <mesh file="link6_1.obj"/>
+    <mesh file="link6_2.obj"/>
+    <mesh file="link6_3.obj"/>
+    <mesh file="link6_4.obj"/>
+    <mesh file="link6_5.obj"/>
+    <mesh file="link6_6.obj"/>
+    <mesh file="link6_7.obj"/>
+    <mesh file="link6_8.obj"/>
+    <mesh file="link6_9.obj"/>
+    <mesh file="link6_10.obj"/>
+    <mesh file="link6_11.obj"/>
+    <mesh file="link6_12.obj"/>
+    <mesh file="link6_13.obj"/>
+    <mesh file="link6_14.obj"/>
+    <mesh file="link6_15.obj"/>
+    <mesh file="link6_16.obj"/>
+    <mesh file="link7_0.obj"/>
+    <mesh file="link7_1.obj"/>
+    <mesh file="link7_2.obj"/>
+    <mesh file="link7_3.obj"/>
+    <mesh file="link7_4.obj"/>
+    <mesh file="link7_5.obj"/>
+    <mesh file="link7_6.obj"/>
+    <mesh file="link7_7.obj"/>
+    <mesh file="hand_0.obj"/>
+    <mesh file="hand_1.obj"/>
+    <mesh file="hand_2.obj"/>
+    <mesh file="hand_3.obj"/>
+    <mesh file="hand_4.obj"/>
+    <mesh file="finger_0.obj"/>
+    <mesh file="finger_1.obj"/>
+    <mesh file="d405.stl"/>
+  </asset>
+
+  <worldbody>
+    <light name="top" pos="0 0 2" mode="trackcom"/>
+    <camera name="front_cam" pos="1.1 0 0.5" xyaxes="0 1 0 -0.4 0 0.7"/>
+    <body name="link0" childclass="panda">
+      <inertial mass="0.629769" pos="0 0 0"
+        fullinertia="0.00315 0.00388 0.004285 8.2904e-7 0.00015 8.2299e-6"/>
+      <geom mesh="link0_0" material="off_white" class="visual"/>
+      <geom mesh="link0_1" material="black" class="visual"/>
+      <geom mesh="link0_2" material="off_white" class="visual"/>
+      <geom mesh="link0_3" material="black" class="visual"/>
+      <geom mesh="link0_4" material="off_white" class="visual"/>
+      <geom mesh="link0_5" material="black" class="visual"/>
+      <geom mesh="link0_7" material="white" class="visual"/>
+      <geom mesh="link0_8" material="white" class="visual"/>
+      <geom mesh="link0_9" material="black" class="visual"/>
+      <geom mesh="link0_10" material="off_white" class="visual"/>
+      <geom mesh="link0_11" material="white" class="visual"/>
+      <geom mesh="link0_collision" class="collision"/>
+      <body name="link1" pos="0 0 0.333">
+        <inertial mass="4.970684" pos="0.003875 0.002081 -0.04762"
+          fullinertia="0.70337 0.70661 0.0091170 -0.00013900 0.0067720 0.019169"/>
+        <joint name="joint1" range="-2.8973 2.8973"/>
+        <geom material="white" mesh="link1" class="visual"/>
+        <geom mesh="link1_collision" class="collision"/>
+        <body name="link2" quat="1 -1 0 0">
+          <inertial mass="0.646926" pos="-0.003141 -0.02872 0.003495"
+            fullinertia="0.0079620 2.8110e-2 2.5995e-2 -3.925e-3 1.0254e-2 7.04e-4"/>
+          <joint name="joint2" range="-1.7628 1.7628"/>
+          <geom material="white" mesh="link2" class="visual"/>
+          <geom mesh="link2_collision" class="collision"/>
+          <body name="link3" pos="0 -0.316 0" quat="1 1 0 0">
+            <joint name="joint3" range="-2.8973 2.8973"/>
+            <inertial mass="3.228604" pos="2.7518e-2 3.9252e-2 -6.6502e-2"
+              fullinertia="3.7242e-2 3.6155e-2 1.083e-2 -4.761e-3 -1.1396e-2 -1.2805e-2"/>
+            <geom mesh="link3_0" material="white" class="visual"/>
+            <geom mesh="link3_1" material="white" class="visual"/>
+            <geom mesh="link3_2" material="white" class="visual"/>
+            <geom mesh="link3_3" material="black" class="visual"/>
+            <geom mesh="link3_collision" class="collision"/>
+            <body name="link4" pos="0.0825 0 0" quat="1 1 0 0">
+              <inertial mass="3.587895" pos="-5.317e-2 1.04419e-1 2.7454e-2"
+                fullinertia="2.5853e-2 1.9552e-2 2.8323e-2 7.796e-3 -1.332e-3 8.641e-3"/>
+              <joint name="joint4" range="-3.0718 -0.0698"/>
+              <geom mesh="link4_0" material="white" class="visual"/>
+              <geom mesh="link4_1" material="white" class="visual"/>
+              <geom mesh="link4_2" material="black" class="visual"/>
+              <geom mesh="link4_3" material="white" class="visual"/>
+              <geom mesh="link4_collision" class="collision"/>
+              <body name="link5" pos="-0.0825 0.384 0" quat="1 -1 0 0">
+                <inertial mass="1.225946" pos="-1.1953e-2 4.1065e-2 -3.8437e-2"
+                  fullinertia="3.5549e-2 2.9474e-2 8.627e-3 -2.117e-3 -4.037e-3 2.29e-4"/>
+                <joint name="joint5" range="-2.8973 2.8973"/>
+                <geom mesh="link5_0" material="black" class="visual"/>
+                <geom mesh="link5_1" material="white" class="visual"/>
+                <geom mesh="link5_2" material="white" class="visual"/>
+                <geom mesh="link5_collision_0" class="collision"/>
+                <geom mesh="link5_collision_1" class="collision"/>
+                <geom mesh="link5_collision_2" class="collision"/>
+                <body name="link6" quat="1 1 0 0">
+                  <inertial mass="1.666555" pos="6.0149e-2 -1.4117e-2 -1.0517e-2"
+                    fullinertia="1.964e-3 4.354e-3 5.433e-3 1.09e-4 -1.158e-3 3.41e-4"/>
+                  <joint name="joint6" range="-0.0175 3.7525"/>
+                  <geom mesh="link6_0" material="off_white" class="visual"/>
+                  <geom mesh="link6_1" material="white" class="visual"/>
+                  <geom mesh="link6_2" material="black" class="visual"/>
+                  <geom mesh="link6_3" material="white" class="visual"/>
+                  <geom mesh="link6_4" material="white" class="visual"/>
+                  <geom mesh="link6_5" material="white" class="visual"/>
+                  <geom mesh="link6_6" material="white" class="visual"/>
+                  <geom mesh="link6_7" material="light_blue" class="visual"/>
+                  <geom mesh="link6_8" material="light_blue" class="visual"/>
+                  <geom mesh="link6_9" material="black" class="visual"/>
+                  <geom mesh="link6_10" material="black" class="visual"/>
+                  <geom mesh="link6_11" material="white" class="visual"/>
+                  <geom mesh="link6_12" material="green" class="visual"/>
+                  <geom mesh="link6_13" material="white" class="visual"/>
+                  <geom mesh="link6_14" material="black" class="visual"/>
+                  <geom mesh="link6_15" material="black" class="visual"/>
+                  <geom mesh="link6_16" material="white" class="visual"/>
+                  <geom mesh="link6_collision" class="collision"/>
+                  <body name="link7" pos="0.088 0 0" quat="1 1 0 0">
+                    <inertial mass="7.35522e-01" pos="1.0517e-2 -4.252e-3 6.1597e-2"
+                      fullinertia="1.2516e-2 1.0027e-2 4.815e-3 -4.28e-4 -1.196e-3 -7.41e-4"/>
+                    <joint name="joint7" range="-2.8973 2.8973"/>
+                    <geom mesh="link7_0" material="white" class="visual"/>
+                    <geom mesh="link7_1" material="black" class="visual"/>
+                    <geom mesh="link7_2" material="black" class="visual"/>
+                    <geom mesh="link7_3" material="black" class="visual"/>
+                    <geom mesh="link7_4" material="black" class="visual"/>
+                    <geom mesh="link7_5" material="black" class="visual"/>
+                    <geom mesh="link7_6" material="black" class="visual"/>
+                    <geom mesh="link7_7" material="white" class="visual"/>
+                    <geom mesh="link7_collision" class="collision"/>
+                    <body name="hand" pos="0 0 0.107" quat="0.9238795 0 0 -0.3826834">
+                      <inertial mass="0.73" pos="-0.01 0 0.03" diaginertia="0.001 0.0025 0.0017"/>
+                      <geom mesh="hand_0" material="off_white" class="visual"/>
+                      <geom mesh="hand_1" material="black" class="visual"/>
+                      <geom mesh="hand_2" material="black" class="visual"/>
+                      <geom mesh="hand_3" material="white" class="visual"/>
+                      <geom mesh="hand_4" material="off_white" class="visual"/>
+                      <geom mesh="hand_collision" class="collision"/>
+                      <geom name="hand_capsule_collision" type="capsule" class="collision" size="0.04 0.06" quat="1 1 0 0" pos="0 0 0.03"/>
+                      <site name="gripper" pos="0 0 0.1"/>
+                      <site name="tip" pos="0 0 0.1034"/>
+                      <body name="camera_d405" pos="0.05 0.0 0.066021" quat="0.6644633 0.2418451 -0.2418451 -0.6644633">
+                        <geom class="visual" mesh="d405"/>
+                        <geom class="collision" type="box" size=".018 .018 .01" pos="0 0 -0.01" group="4"/>
+                        <body name="camera_frame" pos="0 0 0" quat="0 0 0 1">
+                          <camera name="camera_d405" pos="0 0 0" quat="0 1 0 0" resolution="1280 720"
+                            sensorsize="0.003896 0.00214" focal="0.00193 0.00193"/>
+                        </body>
+                      </body>
+                      <body name="left_finger" pos="0 0 0.0584">
+                        <inertial mass="0.015" pos="0 0 0" diaginertia="2.375e-6 2.375e-6 7.5e-7"/>
+                        <joint name="finger_joint1" class="finger" range="0 0.04"/>
+                        <geom mesh="finger_0" material="off_white" class="visual"/>
+                        <geom mesh="finger_1" material="black" class="visual"/>
+                        <geom name="finger_left_collision" mesh="finger_0" class="collision"/>
+                        <geom name="lf1_collision" class="fingertip_collision_1"/>
+                        <geom name="lf2_collision" class="fingertip_collision_2"/>
+                        <geom name="lf3_collision" class="fingertip_collision_3"/>
+                        <geom name="lfp_collision" class="fingertip_pad_collision_6"/>
+                      </body>
+                      <body name="right_finger" pos="0 0 0.0584" quat="0 0 0 1">
+                        <inertial mass="0.015" pos="0 0 0" diaginertia="2.375e-6 2.375e-6 7.5e-7"/>
+                        <joint name="finger_joint2" class="finger" range="0 0.04"/>
+                        <geom mesh="finger_0" material="off_white" class="visual"/>
+                        <geom mesh="finger_1" material="black" class="visual"/>
+                        <geom name="finger_right_collision" mesh="finger_0" class="collision"/>
+                        <geom name="rf1_collision" class="fingertip_collision_1"/>
+                        <geom name="rf2_collision" class="fingertip_collision_2"/>
+                        <geom name="rf3_collision" class="fingertip_collision_3"/>
+                        <geom name="rfp_collision" class="fingertip_pad_collision_6"/>
+                      </body>
+                    </body>
+                  </body>
+                </body>
+              </body>
+            </body>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+
+  <equality>
+    <joint joint1="finger_joint1" joint2="finger_joint2" solimp="0.95 0.99 0.001" solref="0.005 1"/>
+  </equality>
+
+  <contact>
+    <exclude body1="link0" body2="link1"/>
+  </contact>
+</mujoco>

tinysim_mujoco/manipulation/xmls/scene.xml

@@ -0,0 +1,23 @@
+<mujoco model="panda scene">
+  <include file="panda.xml"/>
+  <include file="table.xml"/>
+  <statistic center="0 0 0.1" extent="0.8"/>
+
+  <visual>
+    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
+    <rgba haze="0.15 0.25 0.35 1"/>
+    <global azimuth="120" elevation="-20"/>
+  </visual>
+
+  <asset>
+    <texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
+    <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3"
+      markrgb="0.8 0.8 0.8" width="300" height="300"/>
+    <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
+  </asset>
+
+  <worldbody>
+    <light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
+    <geom name="floor" size="0 0 0.05" type="plane" material="groundplane" contype="1" conaffinity="1"/>
+  </worldbody>
+</mujoco>

tinysim_mujoco/manipulation/xmls/table.xml

@@ -0,0 +1,37 @@
+<mujoco model="table">
+  <asset>
+    <material name="table_wood"
+              rgba="0.6 0.45 0.3 1"
+              specular="0.2"
+              shininess="0.2"/>
+  </asset>
+
+  <worldbody>
+    <body name="table" pos="0.55 0 0">
+
+      <geom name="table_top"
+            type="box"
+            size="0.25 0.35 0.02"
+            pos="0 0 0.50"
+            material="table_wood"
+            contype="1"
+            conaffinity="1"/>
+
+      <geom name="leg_fl" type="box" size="0.025 0.025 0.25" pos=" 0.22  0.32 0.25"/>
+      <geom name="leg_fr" type="box" size="0.025 0.025 0.25" pos=" 0.22 -0.32 0.25"/>
+      <geom name="leg_bl" type="box" size="0.025 0.025 0.25" pos="-0.22  0.32 0.25"/>
+      <geom name="leg_br" type="box" size="0.025 0.025 0.25" pos="-0.22 -0.32 0.25"/>
+    </body>
+
+    <body name="box1" pos="0.35 0 0.60">
+    <freejoint/>
+    <geom name="box1_geom"
+          type="box"
+          size="0.05 0.05 0.05"
+          rgba="0.8 0.2 0.2 1"
+          density="200"
+          contype="1" 
+          conaffinity="1"/>
+    </body>
+  </worldbody>
+</mujoco>

tinysim_mujoco/notebook_viewer.py

@@ -0,0 +1,33 @@
+import numpy as np
+import mujoco
+from IPython.display import display
+from PIL import Image as PILImage
+
+
+class NotebookViewer:
+    def __init__(self, model, data):
+        self.model = model
+        self.data = data
+        self.width = 640
+        self.height = 480
+
+        self.cam = mujoco.MjvCamera()
+        self.opt = mujoco.MjvOption()
+        self.renderer = mujoco.Renderer(model, self.height, self.width, max_geom=500)
+
+        self.cam.azimuth = 90
+        self.cam.elevation = -25
+        self.cam.distance = 4
+        self.cam.lookat[:] = np.array([0.0, 0.0, 0.0])
+        self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE
+        self.renderer.update_scene(self.data, camera=self.cam)
+        pil = PILImage.fromarray(self.renderer.render())
+        self.display = display(pil, display_id=True)
+
+    def render(self):
+        self.renderer.update_scene(self.data, camera=self.cam)
+        pil = PILImage.fromarray(self.renderer.render())
+        self.display.update(pil)
+
+    def close(self):
+        self.renderer.close()

tinysim_mujoco/unitree_a1/__init__.py

@@ -0,0 +1,282 @@
+import pathlib
+
+import numpy as np
+
+try:
+    import mujoco
+    from ..gl_viewer import GLViewer
+    from ..notebook_viewer import NotebookViewer
+except ImportError:
+    raise ImportError(
+        "Mujoco is not properly installed. Install using `pip install tinysim[mujoco]`"
+    )
+
+
+class UnitreeA1BaseEnv:
+
+    def __init__(self, headless=False, **kwargs):
+        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
+        mujoco.mj_forward(self.model, self.data)
+
+        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:
+            if kwargs.get("notebook", False):
+                self.viewer = NotebookViewer(self.model, self.data)
+            else:
+                self.viewer = GLViewer(self.model, self.data)
+        else:
+            self.viewer = None
+
+    def step(self, actions, n_frames=20):
+        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:
+            self.viewer.render()
+
+    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, **kwargs):
+        self.env = UnitreeA1BaseEnv(**kwargs)
+        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.