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

tinysim_mujoco/manipulation/push_env_cam.py

@@ -0,0 +1,274 @@
+import numpy as np
+from PIL import Image as PILImage
+
+try:
+    import mujoco
+
+    from . import ManipulationBaseEnv
+except ImportError:
+    raise ImportError(
+        "Mujoco is not properly installed. Install using `pip install tinysim[mujoco]`"
+    )
+
+_FLOAT_EPS = np.finfo(np.float64).eps
+_EPS4 = _FLOAT_EPS * 4.0
+
+
+class ManipulationEnvV1:
+
+    def __init__(self, **kwargs):
+        kwargs["use_d405_camera"] = True
+        self.env = ManipulationBaseEnv(**kwargs)
+        self._model = self.env.model
+        self._data = self.env.data
+        self.frame_skip = 25
+        self.dt = self._model.opt.timestep * self.frame_skip
+
+        self._body_name2id = {}
+        for body_id in range(self._model.nbody):
+            name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, body_id)
+            if name is not None:
+                self._body_name2id[name] = body_id
+
+        self.joint_names = []
+        for j in range(self._model.njnt):
+            name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_JOINT, j)
+            if name is not None:
+                self.joint_names.append(name)
+
+        distance_threshold = 0.05
+        self.neutral_joint_values = np.array(
+            [0.00, 0.41, 0.00, -1.85, 0.00, 2.26, 0.79]
+        )
+
+        self.current_step = 0
+        self.maximum_episode_steps = 128
+        self.distance_threshold = distance_threshold
+
+        self.goal = self.get_site_xpos(self._model, self._data, "goal_site").copy()
+        self.initial_object_pos = self.get_site_xpos(
+            self._model, self._data, "obj_site"
+        ).copy()
+        free_joint_index = self.joint_names.index("obj_joint")
+        self.arm_joint_names = self.joint_names[:free_joint_index][0:7]
+        # self.gripper_joint_names = self.joint_names[:free_joint_index][7:9]
+        self.set_joint_neutral()
+        self.reset_mocap_welds(self._model, self._data)
+        mujoco.mj_forward(self._model, self._data)
+
+    def reset_mocap_welds(self, model, data):
+        if model.nmocap > 0 and model.eq_data is not None:
+            for i in range(model.eq_data.shape[0]):
+                if model.eq_type[i] == mujoco.mjtEq.mjEQ_WELD:
+                    model.eq_data[i, :7] = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0])
+        mujoco.mj_forward(model, data)
+
+    def reset_mocap2body_xpos(self, model, data):
+        if model.eq_type is None or model.eq_obj1id is None or model.eq_obj2id is None:
+            return
+        for eq_type, obj1_id, obj2_id in zip(
+            model.eq_type, model.eq_obj1id, model.eq_obj2id
+        ):
+            if eq_type != mujoco.mjtEq.mjEQ_WELD:
+                continue
+
+            mocap_id = model.body_mocapid[obj1_id]
+            if mocap_id != -1:
+                # obj1 is the mocap, obj2 is the welded body
+                body_idx = obj2_id
+            else:
+                # obj2 is the mocap, obj1 is the welded body
+                mocap_id = model.body_mocapid[obj2_id]
+                body_idx = obj1_id
+
+            assert mocap_id != -1
+            data.mocap_pos[mocap_id][:] = data.xpos[body_idx]
+            data.mocap_quat[mocap_id][:] = data.xquat[body_idx]
+
+    def mocap_set_action(self, model, data, action):
+        if model.nmocap > 0:
+            action, _ = np.split(action, (model.nmocap * 7,))
+            action = action.reshape(model.nmocap, 7)
+            pos_delta = action[:, :3]
+            quat_delta = action[:, 3:]
+            self.reset_mocap2body_xpos(model, data)
+            data.mocap_pos[:] = data.mocap_pos + pos_delta
+            data.mocap_quat[:] = data.mocap_quat + quat_delta
+
+    def ctrl_set_action(self, model, data, action):
+        if model.nmocap > 0:
+            _, action = np.split(action, (model.nmocap * 7,))
+
+        if len(data.ctrl) > 0:
+            for i in range(action.shape[0]):
+                if model.actuator_biastype[i] == 0:
+                    data.ctrl[i] = action[i]
+                else:
+                    idx = model.jnt_qposadr[model.actuator_trnid[i, 0]]
+                    data.ctrl[i] = data.qpos[idx] + action[i]
+
+    def step(self, action):
+        self.current_step += 1
+        # just move the end effector in x, y, z
+        pos_ctrl = action.copy() * 0.05  # scale action
+        rot_ctrl = [
+            1.0,
+            0.0,
+            0.0,
+            0.0,
+        ]
+        action = np.concatenate([pos_ctrl, rot_ctrl])
+
+        self.ctrl_set_action(self._model, self._data, action)
+        self.mocap_set_action(self._model, self._data, action)
+
+        self.env.step(n_frames=self.frame_skip)
+        obs = self._get_obs()
+        terminated = self._is_success(obs["achieved_goal"], self.goal)
+        truncated = self.current_step >= self.maximum_episode_steps
+        reward = self.compute_reward(obs["achieved_goal"], self.goal)
+        return obs, reward, terminated, truncated, {}
+
+    def mat2euler(self, mat):
+        mat = np.asarray(mat, dtype=np.float64)
+        cy = np.sqrt(mat[..., 2, 2] * mat[..., 2, 2] + mat[..., 1, 2] * mat[..., 1, 2])
+        condition = cy > _EPS4
+        euler = np.empty(mat.shape[:-1], dtype=np.float64)
+        euler[..., 2] = np.where(
+            condition,
+            -np.arctan2(mat[..., 0, 1], mat[..., 0, 0]),
+            -np.arctan2(-mat[..., 1, 0], mat[..., 1, 1]),
+        )
+        euler[..., 1] = np.where(
+            condition,
+            -np.arctan2(-mat[..., 0, 2], cy),
+            -np.arctan2(-mat[..., 0, 2], cy),
+        )
+        euler[..., 0] = np.where(
+            condition, -np.arctan2(mat[..., 1, 2], mat[..., 2, 2]), 0.0
+        )
+        return euler
+
+    def get_site_xpos(self, model, data, name):
+        site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, name)
+        return data.site_xpos[site_id]
+
+    def get_site_xvelp(self, model, data, name):
+        site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, name)
+        jacp = self.get_site_jacp(model, data, site_id)
+        xvelp = jacp @ data.qvel
+        return xvelp
+
+    def get_site_jacp(self, model, data, site_id):
+        jacp = np.zeros((3, model.nv))
+        mujoco.mj_jacSite(model, data, jacp, None, site_id)
+        return jacp
+
+    def _get_obs(self) -> dict:
+        ee_position = self.get_site_xpos(self._model, self._data, "tip").copy()
+        ee_velocity = (
+            self.get_site_xvelp(self._model, self._data, "tip").copy() * self.dt
+        )
+        object_position = self.get_site_xpos(self._model, self._data, "obj_site").copy()
+        object_rotation = self.mat2euler(
+            self.get_site_xmat(self._model, self._data, "obj_site")
+        )
+
+        frame = self.env.d405_viewer.capture_frame()
+        img = np.array(PILImage.fromarray(frame).resize((64, 64))).transpose(2, 0, 1)
+
+        return {
+            "observation_cap": img,
+            "observation": np.concatenate(
+                [
+                    ee_position,
+                    ee_velocity,
+                    object_position,
+                    object_rotation,
+                ]
+            ).copy(),
+            "achieved_goal": object_position.copy(),
+            "desired_goal": self.goal.copy(),
+        }
+
+    def set_joint_qpos(self, model, data, name, value):
+        """Set the joint positions (qpos) of the model."""
+        joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name)
+        joint_addr = model.jnt_qposadr[joint_id]
+        # All joints are assumed to be of type mjJNT_HINGE in this environment
+        data.qpos[joint_addr] = value
+
+    def set_object_qpos(self, model, data):
+        joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "obj_joint")
+        joint_addr = model.jnt_qposadr[joint_id]
+        data.qpos[joint_addr : joint_addr + 3] = np.array(self.initial_object_pos)
+        # reset orientation
+        data.qpos[joint_addr + 3 : joint_addr + 7] = np.array([1.0, 0.0, 0.0, 0.0])
+
+    def _sample_goal(self):
+        goal = np.array([0.3, 0.0, 0.0])
+        return goal
+
+    def goal_distance(self, goal_a, goal_b):
+        assert goal_a.shape == goal_b.shape
+        return np.linalg.norm(goal_a - goal_b, axis=-1)
+
+    def compute_reward(self, achieved_goal, desired_goal, info=None):
+        d = self.goal_distance(achieved_goal, desired_goal)
+        # give reward for the end effector being close to the object
+        ee_position = self.get_site_xpos(self._model, self._data, "tip")
+        object_position = self.get_site_xpos(self._model, self._data, "obj_site")
+        ee_object_distance = np.linalg.norm(ee_position - object_position, axis=-1)
+
+        # the arm tries to cheat by moving the object with the side of the gripper
+        # negative ward if the arm contacts the ground with
+        effector_distance_reward = -0.1 * (ee_object_distance > 0.05).astype(np.float32)
+
+        # todo: check correctness
+        # give some reward if the cube is moving towards the goal
+        object_velocity = self.get_site_xvelp(self._model, self._data, "obj_site")
+        goal_direction = desired_goal[0] - object_position
+        goal_direction /= np.linalg.norm(goal_direction) + 1e-8
+        velocity_towards_goal = np.dot(object_velocity, goal_direction)
+        velocity_reward = max(0.0, velocity_towards_goal)
+
+        reward = -d + effector_distance_reward + velocity_reward
+        return reward
+
+        # sparse
+        # return -(d > self.distance_threshold).astype(np.float32)
+
+    def _is_success(self, achieved_goal, desired_goal):
+        d = self.goal_distance(achieved_goal, desired_goal)
+        return (d < self.distance_threshold).astype(np.float32)
+
+    def set_joint_neutral(self):
+        for name, value in zip(self.arm_joint_names, self.neutral_joint_values):
+            self.set_joint_qpos(self._model, self._data, name, value)
+
+    def _reset_sim(self) -> bool:
+        self.set_joint_neutral()
+        self.set_object_qpos(self._model, self._data)
+        mujoco.mj_forward(self._model, self._data)
+
+    def reset(self, **kwargs):
+        self.current_step = 0
+        self._reset_sim()
+        obs = self._get_obs()
+
+        # self.env.step(n_frames=self.frame_skip)
+        # frame = self.env.viewer.capture_frame()
+        # import matplotlib.pyplot as plt
+        # plt.imsave("debug_frame.png", frame)
+        # raise Exception("Debug")
+        return obs, {}
+
+    def get_site_xmat(self, model, data, name: str):
+        site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, name)
+        return data.site_xmat[site_id].reshape(3, 3)
+
+    def get_body_state(self, name):
+        body_id = self._body_name2id[name]
+        body_xpos = self._data.xpos[body_id]
+        body_xquat = self._data.xquat[body_id]
+        return np.concatenate([body_xpos, body_xquat])

tinysim_mujoco/manipulation/xmls/panda.xml

@@ -52,7 +52,6 @@
     <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"/>
@@ -65,7 +64,6 @@
     <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"/>
@@ -128,6 +126,15 @@
   <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"/>
+
+    <!-- Added -->
+    <body mocap="true" name="panda_mocap" pos="0.6 0 0.4">
+        <geom conaffinity="0" contype="0" pos="0 0 0" rgba="0 0.5 0 0.7" size="0.005 0.005 0.005" type="box"></geom>
+        <geom conaffinity="0" contype="0" pos="0 0 0" rgba="0 0.5 0 0.1" size="0.5 0.005 0.005" type="box"></geom>
+        <geom conaffinity="0" contype="0" pos="0 0 0" rgba="0 0.5 0 0.1" size="0.005 0.5 0.001" type="box"></geom>
+        <geom conaffinity="0" contype="0" pos="0 0 0" rgba="0 0.5 0 0.1" size="0.005 0.005 0.5" type="box"></geom>
+    </body>
+
     <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"/>
@@ -146,17 +153,17 @@
       <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"/>
+        <joint name="joint1" range="-2.8973 2.8973" /> <!-- Here -->
         <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"/>
+          <joint name="joint2" range="-1.7628 1.7628" /> <!-- Here -->
           <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"/>
+            <joint name="joint3" range="-2.8973 2.8973" /> <!-- Here -->
             <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"/>
@@ -167,7 +174,7 @@
             <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"/>
+              <joint name="joint4" range="-3.0718 -0.0698" /> <!-- Here -->
               <geom mesh="link4_0" material="white" class="visual"/>
               <geom mesh="link4_1" material="white" class="visual"/>
               <geom mesh="link4_2" material="black" class="visual"/>
@@ -176,7 +183,7 @@
               <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"/>
+                <joint name="joint5" range="-2.8973 2.8973" /> <!-- Here -->
                 <geom mesh="link5_0" material="black" class="visual"/>
                 <geom mesh="link5_1" material="white" class="visual"/>
                 <geom mesh="link5_2" material="white" class="visual"/>
@@ -186,7 +193,8 @@
                 <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"/>
+                  <joint name="joint6" range="-0.0175 3.7525" /> <!-- Here -->
+
                   <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"/>
@@ -208,7 +216,7 @@
                   <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"/>
+                    <joint name="joint7" range="-2.8973 2.8973" /> <!-- Here -->
                     <geom mesh="link7_0" material="white" class="visual"/>
                     <geom mesh="link7_1" material="black" class="visual"/>
                     <geom mesh="link7_2" material="black" class="visual"/>
@@ -226,9 +234,13 @@
                       <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"/>
+
+                      <!-- These are the red circles between the fingers -->
+                      <site name="gripper" pos="0 0 0.1" rgba="1 0 0 1"/>
+                      <site name="tip" pos="0 0 0.1034" rgba="1 0 0 1" /> 
+
                       <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"/>
@@ -237,9 +249,9 @@
                             sensorsize="0.003896 0.00214" focal="0.00193 0.00193"/>
                         </body>
                       </body>
-                      <body name="left_finger" pos="0 0 0.0584">
+                      <body name="left_finger" pos="0 0.0 0.0584"> <!-- Changed -->
                         <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"/>
+                        <!-- <joint name="finger_joint1" class="finger"/> -->
                         <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"/>
@@ -248,9 +260,9 @@
                         <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">
+                      <body name="right_finger" pos="0 0.0 0.0584" quat="0 0 0 1"> <!-- Changed -->
                         <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"/>
+                        <!-- <joint name="finger_joint2" class="finger"/> -->
                         <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"/>
@@ -270,10 +282,41 @@
     </body>
   </worldbody>
 
-  <equality>
+  <!-- <equality>
     <joint joint1="finger_joint1" joint2="finger_joint2" solimp="0.95 0.99 0.001" solref="0.005 1"/>
+  </equality> -->
+  <equality>
+    <weld body1="panda_mocap" body2="hand" solimp="0.9 0.95 0.001" solref="0.02 1"></weld>
   </equality>
 
+  <!-- Actuators for the panda robot -->
+
+  <!-- <actuator>
+    <general class="panda" name="actuator1" joint="joint1" kp="870" forcerange="-87 87" ctrlrange="-2.9671 2.9671"/>
+    <general class="panda" name="actuator2" joint="joint2" kp="870" forcerange="-87 87" ctrlrange="-1.8326 1.8326"/>
+    <general class="panda" name="actuator3" joint="joint3" gainprm="3500" biasprm="0 -3500 -350"/>
+    <general class="panda" name="actuator4" joint="joint4" gainprm="3500" biasprm="0 -3500 -350" ctrlrange="-3.0718 -0.0698"/>
+    <general class="panda" name="actuator5" joint="joint5" gainprm="2000" biasprm="0 -2000 -200" forcerange="-12 12"/>
+    <general class="panda" name="actuator6" joint="joint6" gainprm="2000" biasprm="0 -2000 -200" forcerange="-12 12" ctrlrange="-0.0175 3.7525"/>
+    <general class="panda" name="actuator7" joint="joint7" gainprm="2000" biasprm="0 -2000 -200" forcerange="-12 12"/>
+  </actuator> -->
+
+  <actuator>
+    <position class="panda" name="actuator1" joint="joint1" kv="870" kp="870" forcerange="-87 87" ctrlrange="-2.9671 2.9671"/>
+    <position class="panda" name="actuator2" joint="joint2" kv="870" kp="870" forcerange="-87 87" ctrlrange="-1.8326 1.8326"/>
+    <position class="panda" name="actuator3" joint="joint3" kv="870" kp="870" forcerange="-87 87" ctrlrange="-2.9671 2.9671"/>
+    <position class="panda" name="actuator4" joint="joint4" kv="870" kp="870" forcerange="-87 87" ctrlrange="-3.1416 0.0"/>
+    <position class="panda" name="actuator5" joint="joint5" kv="120" kp="120" forcerange="-12 12" ctrlrange="-2.9671 2.9671"/>
+    <position class="panda" name="actuator6" joint="joint6" kv="120" kp="120" forcerange="-12 12" ctrlrange="-0.0175 3.7525"/>
+    <position class="panda" name="actuator7" joint="joint7" kv="120" kp="120" forcerange="-12 12" ctrlrange="-2.9671 2.9671"/>
+  </actuator>
+
+
+  <!-- <actuator>
+    <general name="r_gripper_finger_joint" joint="finger_joint1" class="panda" gainprm="500 0 0" biasprm="0 -500 -10" forcerange="-120 120" ctrlrange="0 0.04"/>
+    <general name="l_gripper_finger_joint" joint="finger_joint2" class="panda" gainprm="500 0 0" biasprm="0 -500 -10" forcerange="-120 120" ctrlrange="0 0.04"/>
+  </actuator> -->
+
   <contact>
     <exclude body1="link0" body2="link1"/>
   </contact>

tinysim_mujoco/manipulation/xmls/scene.xml

@@ -2,20 +2,17 @@
   <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"/>

tinysim_mujoco/manipulation/xmls/table.xml

@@ -5,10 +5,8 @@
               specular="0.2"
               shininess="0.2"/>
   </asset>
-
   <worldbody>
-    <body name="table" pos="0.55 0 0">
-
+    <!-- <body name="table" pos="0.55 0 0">
       <geom name="table_top"
             type="box"
             size="0.25 0.35 0.02"
@@ -16,22 +14,28 @@
             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">
+      <joint type="free" damping="0.01" name="obj_joint"/>
+      <site name="obj_site" pos="0 0 0" size="0.01" type="sphere"/>
+      <geom name="box1_geom"
+            type="box"
+            size="0.02 0.02 0.02"
+            rgba="0.8 0.2 0.2 1"
+            friction="1" condim="6"/>
+    </body> -->
+    <body name="obj" pos="0.6 0.0 0.05" quat="1 0 0 0">
+      <joint type="free" damping="0.01" name="obj_joint"/>
+      <geom name="obj_geom" mass="0.1" type="box" size="0.03 0.03 0.03" rgba="0.5 0.7 0.5 1" friction="1" condim="6"/>
+      <site name="obj_site" pos="0 0 0" size="0.01" type="sphere"/>
     </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 name="goal_site" pos="0.3 0.0 0.0" quat="1 0 0 0">
+      <site name="goal_site" pos="0 0 0" size="0.05" type="sphere" rgba="0 1 0 0.5"/>
     </body>
   </worldbody>
 </mujoco>

tinysim_mujoco/notebook_viewer.py

@@ -1,19 +1,16 @@
-import numpy as np
 import mujoco
+import numpy as np
 from IPython.display import display
 from PIL import Image as PILImage
 
 
 class NotebookViewer:
-    def __init__(self, model, data):
+    def __init__(self, model, data, width=640, height=480):
         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.renderer = mujoco.Renderer(model, height, width, max_geom=500)
 
         self.cam.azimuth = 90
         self.cam.elevation = -25

tinysim_mujoco/unitree_a1/__init__.py

@@ -4,6 +4,7 @@ import numpy as np
 
 try:
     import mujoco
+
     from ..gl_viewer import GLViewer
     from ..notebook_viewer import NotebookViewer
 except ImportError:
@@ -50,9 +51,6 @@ class UnitreeA1BaseEnv:
         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]
@@ -77,6 +75,7 @@ class UnitreeA1BaseEnv:
 class UnitreeA1WalkEnv:
     def __init__(self, reward_weights, cost_weights, **kwargs):
         self.env = UnitreeA1BaseEnv(**kwargs)
+        self.viewer = self.env.viewer
         self.frame_skip = 20
         self.obs_dim = 45
         self._previous_observation = np.zeros(self.obs_dim, dtype=np.float32)
@@ -165,14 +164,14 @@ class UnitreeA1WalkEnv:
         action_diff_penalty = np.sum(np.abs(action - self.prev_action))
         action_sym = self.action_sym()
 
-        Positive_rewards = (
+        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 = (
+        neg_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"]
@@ -180,11 +179,9 @@ class UnitreeA1WalkEnv:
             + 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"]
+            + self.default_joint_pos_cost * self.cost_weights["default_joint_pos"]
         )
-        reward = Positive_rewards - Negative_rewards
-        return reward
+        return rewards - neg_rewards
 
     def _tracking_velocity_penalty(self, lin_vel, ang_vel):
         lin_vel_error = np.sum(np.abs(self.target_lin_vel[0] - lin_vel[0]))
@@ -235,7 +232,7 @@ class UnitreeA1WalkEnv:
         return np.square(roll) + np.square(pitch)
 
     @property
-    def default_joint_position_cost(self):
+    def default_joint_pos_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