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

tinysim/topdown_driving/widget.py

@@ -1,59 +1,25 @@
 import pathlib
-import anywidget
+
 import traitlets
-import asyncio
-import numpy as np
-from IPython.display import display
-from jupyter_ui_poll import ui_events
 
-from . import TopDownDrivingEnv, LOCAL_WALLS
+from .._widget_base import BaseWidget
+from . import LOCAL_WALLS, TopDownDrivingEnv
 
 
-class TopDownDrivingWidget(anywidget.AnyWidget):
+class TopDownDrivingWidget(BaseWidget):
     _esm = pathlib.Path(__file__).parent / "sim.js"
 
     sim_state = traitlets.Dict(default_value={}).tag(sync=True)
     wall_positions = traitlets.List(default_value=LOCAL_WALLS).tag(sync=True)
+    _viewport_size = traitlets.Tuple(default_value=(800, 600)).tag(sync=True)
 
-    _viewport_size = traitlets.Tuple(
-        traitlets.Int(), traitlets.Int(), default_value=(800, 600)
-    ).tag(sync=True)
-    _view_ready = traitlets.Bool(default_value=False).tag(sync=True)
-
-    def __init__(self, viewport_size=(800, 600), sim_env=None):
+    def __init__(self, viewport_size=(800, 600), sim_env=TopDownDrivingEnv()):
+        super().__init__(sim_env)
         self._viewport_size = viewport_size
-        super().__init__()
-        if sim_env is None:
-            sim_env = TopDownDrivingEnv()
-
-        self.sim_env = sim_env
-        self.copy_py_state(self.sim_env.reset())
-
-    def render(self):
-        display(self)
 
-        try:
-            with ui_events() as ui_poll:
-                while not self._view_ready:
-                    ui_poll(100)
-        except Exception:
-            pass
-
-    def copy_py_state(self, sim_state: dict):
+    def _update_props(self, sim_state):
         self.sim_state = {
             "x": sim_state["x"].tolist(),
             "y": sim_state["y"].tolist(),
             "angle": sim_state["angle"].tolist(),
         }
-
-    async def step(self, action: int, dt: float = 0.01) -> dict:
-        sim_state = self.sim_env.step(action)
-        self.copy_py_state(sim_state)
-
-        await asyncio.sleep(dt)
-        return sim_state
-
-    async def reset(self) -> dict:
-        sim_state = self.sim_env.reset()
-        self.copy_py_state(sim_state)
-        return sim_state

{tinysim-0.0.4.dist-info → tinysim-0.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinysim
-Version: 0.0.4
+Version: 0.0.5
 Summary: small modular simulation environments
 Author-email: Matthew Taylor <matthew.taylor.andre@gmail.com>
 Project-URL: Homepage, https://github.com/MatthewAndreTaylor/TinySim
@@ -30,7 +30,7 @@ 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)
+[![Jupyterlite](https://jupyterlite.rtfd.io/en/latest/_static/badge.svg)](https://matthewandretaylor.github.io/TinySim)
 
 The goal of this project is to create more minimal simulation environments for robotics and machine learning.
 
@@ -73,8 +73,7 @@ To use the [mujoco](https://github.com/google-deepmind/mujoco) simulation enviro
 ```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" /> |
-
+| [UnitreeA1WalkEnv](https://github.com/MatthewAndreTaylor/TinySim/tree/main/tinysim_mujoco/unitree_a1) | [ManipulationEnvV0](https://github.com/MatthewAndreTaylor/TinySim/tree/main/tinysim_mujoco/manipulation) |
+|:-------------------------:|:-------------------------:|
+| <img width="312" alt="mujoco_sim" src="https://github.com/user-attachments/assets/1109488e-0563-4694-8926-682edd0c8278" /> | <img width="312" alt="mujoco_sim" src="https://github.com/user-attachments/assets/07b406b5-15fc-4cfd-83c4-cf6b855d773a" /> |
 

{tinysim-0.0.4.dist-info → tinysim-0.0.5.dist-info}/RECORD

@@ -1,35 +1,34 @@
-tinysim/__init__.py,sha256=ptkxqwWun-Tw-L2-zUTm2fqunNgdLBkdU3DHF1wAMBI,197
-tinysim/_tk_base.py,sha256=kzF-41H56A5FWNxtMhlvEbl3iH8Wp04zTqAxTcjRPDM,1279
-tinysim/flappy/__init__.py,sha256=lUHGteFeYRb2vQ3-x9WxnKWM7sfsDwIHoVps732JIDY,3482
+tinysim/__init__.py,sha256=971PtvBBFPJ9zBfRRLdfuGRE2oEX739127n_Y6MNJIg,263
+tinysim/_tk_base.py,sha256=vUJEx0NXfo8TLZycQNOhZgbx_2ybrBjvaJC_KQ6jvBo,1790
+tinysim/_widget_base.py,sha256=_L6xSOH4LOpXIvCFCczDkcorS8xH2E6UjdTH3L6B6tM,1084
+tinysim/tinyspace.py,sha256=x4MKrou4-9aKmNQ0V7RSxYhHk3akzruF56Sn14wnO7s,1950
+tinysim/flappy/__init__.py,sha256=_0OBPUyymQCVkwCRiG7Km_oIwlrhWP808r4e1COSZ3A,3551
 tinysim/flappy/sim.js,sha256=g8ycKdJGopu0Z5CeBgtF8Rr2Pe64dTFcSiYh7hzYPzU,3190
-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/flappy/tk.py,sha256=HR3AlHzGLxLWw2qdtY8pbyFxEELb3sItGWNeOqqjHDA,2111
+tinysim/flappy/widget.py,sha256=BSrEvAmUKbaqF9zZq7rZ4WxgIDjt1NTmUIb0gg0gt4M,748
+tinysim/frogger/__init__.py,sha256=YHumyU2tR8oRZoJACLPQ3Ncyw7b3MbR_W2LJSmtdv9g,5712
 tinysim/frogger/sim.js,sha256=b1HfEtFZFbbsrZNtltuWO4h-QM4Fw3XjMRydwm5l2NU,3081
-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/frogger/tk.py,sha256=gXxp_9sbcFM__-JdX6Jhymf9ZzoIt7N1GmK9DPTW724,2518
+tinysim/frogger/widget.py,sha256=mYlc7zToV7C1Gy_OdvFjCk2xygAdO2ggGupNxXruckk,952
+tinysim/mountain_car/__init__.py,sha256=zR-WaZGF3PPz4ajBlhOqWDxaKSa68bYuKJNaR_PwKjI,2207
 tinysim/mountain_car/sim.js,sha256=z7kxl7LbkRglfzXxmrBbwz2vMYdQV84QP3WhE0hoPVQ,4581
 tinysim/mountain_car/styles.css,sha256=a8Bc64RdRNB1uWsBdpMEQA9STZO8mVAB2vvIPhhYuik,63
-tinysim/mountain_car/tk.py,sha256=EbnLirA0Q5egJaDUzu_WkGjbofbclwYUcHIvvifNt0Y,4256
-tinysim/mountain_car/widget.py,sha256=bUIXsofy5moBXMDBdIOVvRi59NFbLRRQ1WZcwm_mYG8,1726
-tinysim/simple_amr/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tinysim/simple_amr/example_maps.py,sha256=38duuHbd6HvhQBIIcg9PWokNMjkasPSsKPjHnBhv9qQ,3307
-tinysim/simple_amr/sim.js,sha256=2YW2YWEEe-QWgJGLAdpt1aZi__sLL5ML84fv1s90poY,11575
-tinysim/simple_amr/styles.css,sha256=t85kXf0xyrq5ldL-j_pV0FEekWoorK63T_ZorWEgBOw,1087
-tinysim/simple_amr/widget.py,sha256=dmNcftB7-Z9Kf3hQ0MBaJvo97eyjCA6R74oEkHASMas,2557
-tinysim/topdown_driving/__init__.py,sha256=HiRpxQNeu4IIOPFraGHLTO2F0VNISAk6bUiK1rsbZTE,6228
+tinysim/mountain_car/tk.py,sha256=WMlktnnSEWXJ7ppWb-WKg2wm_jtWN-XwIrFVGk64GIo,3753
+tinysim/mountain_car/widget.py,sha256=3EasBKyA0_MsjqlG97OLBNRlGyK8qzT2zQijWnVf9SM,860
+tinysim/topdown_driving/__init__.py,sha256=ZA4DcQHRs9zeraTYE7Id8ityGZksYtGPN8GiSq2RsJI,5965
 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=Eh7tvtmMbgmzFEvQZ10D6npfF2s1k7QD5AHLYX5Jaqw,1718
+tinysim/topdown_driving/tk.py,sha256=WfiFq1zub7x9MGi5aQ4JVQB8qVGKwBT47e3n-_Mju48,4272
+tinysim/topdown_driving/track_0.json,sha256=tZVZtftQCdsB11Gwqj6OECOyOa52RBFTE3vk_e4n8_k,9515
+tinysim/topdown_driving/widget.py,sha256=KVn8BCHicVql07BqdcO8x3i1H67woDiTbu60ypJPZc4,802
 tinysim_mujoco/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tinysim_mujoco/gl_viewer.py,sha256=s0x4HuYwspsIOl7CBywz7cs0g37p22e7fWGtcLXuq6g,2754
-tinysim_mujoco/notebook_viewer.py,sha256=4hbQKcJLC8jAOd5xQyz5yjia5N0_CaKe0s4LgP0VnmM,1035
-tinysim_mujoco/manipulation/__init__.py,sha256=7JB29iraYN64MXAuWNoB6vY38Qr_KvZB66BxPmhZw9U,1129
-tinysim_mujoco/manipulation/xmls/panda.xml,sha256=y5tZFUX6zbj8nb32jux4WuT4nMQ4u3i-ziOEnfoFtFA,14803
-tinysim_mujoco/manipulation/xmls/scene.xml,sha256=cM_FJCnt-aK4b4L2kTVNSd7P-rL6l4aQodOcd-5zB-s,921
-tinysim_mujoco/manipulation/xmls/table.xml,sha256=TFG7lbTr-93_V94sDpJMveuGgtNVPrip6g1YFeE8eXs,1068
+tinysim_mujoco/gl_viewer.py,sha256=0cPmjk3GDhtca9cMiV3px_w8McjXF577NNNGJIdAWxQ,3401
+tinysim_mujoco/notebook_viewer.py,sha256=TwLa4VCKovE3y8NkcdSOkUhoQ3JHvq7Pqni5VpANhQ4,996
+tinysim_mujoco/manipulation/__init__.py,sha256=8sEszaf47LjRaVlSO6k81cwK-UT9mRRHB2mXZyfCl-s,1675
+tinysim_mujoco/manipulation/push_env.py,sha256=IEQi4dR8DVwWxwvmzZhzYLBzVTkZ_62tKsy2yL2PDYw,10456
+tinysim_mujoco/manipulation/push_env_cam.py,sha256=B-GSB9WHew2iq4iUDWc-X78OqwHBl1jtpTGvSgVaOyY,10707
+tinysim_mujoco/manipulation/xmls/panda.xml,sha256=5uTQ473BIGDIO1aASEeLIxw2aaMBjistu_PYMjaN6BU,17868
+tinysim_mujoco/manipulation/xmls/scene.xml,sha256=_6nYdaFPi_Of-GTJPBUg2nCLKab_uz5SM-eFWfvz6Gc,918
+tinysim_mujoco/manipulation/xmls/table.xml,sha256=Gm0-ejM7mmrJQXGkZjEI2JpyOElQsJNbh3srGSWTuUI,1637
 tinysim_mujoco/manipulation/xmls/assets/d405.stl,sha256=Zhfz-yK2BKZsJ-FtI5fzBfzcqKeBJIq9urv2nJ_mo8U,50084
 tinysim_mujoco/manipulation/xmls/assets/finger_0.obj,sha256=8UwcM_XWXKrdyKWg0-Hwg-GjPsqqe-HpHLgdTyHE4vM,23013
 tinysim_mujoco/manipulation/xmls/assets/finger_1.obj,sha256=tLZPSqpntdM74TQDqHT-r0-ov3swyLJr15JCgpx-QIg,16772
@@ -98,7 +97,7 @@ tinysim_mujoco/manipulation/xmls/assets/link7_4.obj,sha256=SZSny1BhpxKpgq5AYX1Hy
 tinysim_mujoco/manipulation/xmls/assets/link7_5.obj,sha256=MsT-PFfY9bpGdtVULacqRtZKBSKmZ15Lf6J_tChFlkE,61442
 tinysim_mujoco/manipulation/xmls/assets/link7_6.obj,sha256=1avsmFwLFqymiEQN7-usXrwQcYDJAhoiZBgs1RXnqNw,27403
 tinysim_mujoco/manipulation/xmls/assets/link7_7.obj,sha256=s56R21mHTno3skIqRb9lnEoS5Qw8G6N74E6lulVmdgo,221627
-tinysim_mujoco/unitree_a1/__init__.py,sha256=VH71dmkOeBWBdNVWyYJ5MMfBx-wpPfZ-GuFJY9kbYfA,10632
+tinysim_mujoco/unitree_a1/__init__.py,sha256=tcWaG1yL8tnF-e6EdQcKNVhQYUsiFsoqmkUD8qZ7D1I,10542
 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
@@ -109,11 +108,12 @@ tinysim_mujoco/unitree_a1/unitree_a1/assets/hip.obj,sha256=fVW_Gc8Csumvrb-z-4Ztd
 tinysim_mujoco/unitree_a1/unitree_a1/assets/thigh.obj,sha256=CdXaRgx4N_02RSH8R6h5k3jyvbd8df04YM8qlP-En1c,979501
 tinysim_mujoco/unitree_a1/unitree_a1/assets/thigh_mirror.obj,sha256=2LL2hKx8de-Ivj7lM5jkqCEOHTm_AMj78bVyqh99orw,1010010
 tinysim_mujoco/unitree_a1/unitree_a1/assets/trunk.obj,sha256=8DX8cdkBSlOq4ay-KVhXghmTTjqJwR0qqj9O_5DAt3c,2014440
-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.4.dist-info/METADATA,sha256=iCf6ezMgRK8qjZt1Q_Sk0dqEWYyMlgao-XtDni0MqYw,3941
-tinysim-0.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tinysim-0.0.4.dist-info/top_level.txt,sha256=v-9NmTLfZukVjPwJwEdtCU7yY4TdLsBf1H8s25tR3V4,36
-tinysim-0.0.4.dist-info/RECORD,,
+tinysim_warp/__init__.py,sha256=i8lHo6ZtR-XGD-isd9zifnp3F1DEs24HOcOKeF9mCjg,3525
+tinysim_warp/cart_pole/__init__.py,sha256=ED-xfKpTlrRftjfBHDyIQ8eDuE3HxuniMAeRWNVAJXg,3636
+tinysim_warp/quadruped/__init__.py,sha256=3fJBrG-mBdim2lT_ON5LlFEYPCfFCs608zNvbB7Q6gs,2247
+tinysim_warp/simple_quadruped/__init__.py,sha256=QoIXplgKUoARiZn4N9Xqu6o2OZn0XLOSk61Ki-5xiz8,3085
+tinysim_warp/simple_quadruped/simple_quadruped.urdf,sha256=60eXZG8nH-ABbtJzYuzdGqdLq-gFKSaDM7brNYt5PN4,8177
+tinysim-0.0.5.dist-info/METADATA,sha256=-EqR_V8RIjVImxDFldhYHM3legWynSl2cM9XIMTv4ko,4177
+tinysim-0.0.5.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+tinysim-0.0.5.dist-info/top_level.txt,sha256=v-9NmTLfZukVjPwJwEdtCU7yY4TdLsBf1H8s25tR3V4,36
+tinysim-0.0.5.dist-info/RECORD,,

{tinysim-0.0.4.dist-info → tinysim-0.0.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

tinysim_mujoco/gl_viewer.py

@@ -1,10 +1,35 @@
-import numpy as np
-import mujoco
 import glfw
+import mujoco
+import numpy as np
+
+
+class OffScreenRenderer:
+    def __init__(self, model, data, cam_id):
+        self.model = model
+        self.data = data
+        self.cam_id = cam_id
+        self.renderer = mujoco.Renderer(self.model)
+        self.renderer.update_scene(self.data, camera=self.cam_id)
+
+    def render(self):
+        self.renderer.update_scene(self.data, camera=self.cam_id)
+
+    def capture_frame(self):
+        return self.renderer.render()
+
+
+def default_camera():
+    cam = mujoco.MjvCamera()
+    cam.azimuth = 90.0
+    cam.elevation = -25.0
+    cam.distance = 4.0
+    cam.lookat[:] = np.array([0.0, 0.0, 0.0])
+    cam.type = mujoco.mjtCamera.mjCAMERA_FREE
+    return cam
 
 
 class GLViewer:
-    def __init__(self, model, data, width=1200, height=900):
+    def __init__(self, model, data, cam_id=-1, width=1200, height=900):
         self.width = width
         self.height = height
         self.model = model
@@ -15,6 +40,13 @@ class GLViewer:
         if not glfw.init():
             raise RuntimeError("Failed to initialize GLFW")
 
+        if cam_id != -1:
+            self.cam = mujoco.MjvCamera()
+            self.cam.fixedcamid = cam_id
+            self.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED
+        else:
+            self.cam = default_camera()
+
         self._create_window()
 
     def _create_window(self):
@@ -27,19 +59,12 @@ class GLViewer:
         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
@@ -85,4 +110,4 @@ class GLViewer:
 
         # OpenGL framebuffer origin is bottom-left, flip vertically
         rgb_buffer = np.flip(rgb_buffer, axis=0)
-        return rgb_buffer
+        return rgb_buffer

tinysim_mujoco/manipulation/__init__.py

@@ -1,10 +1,9 @@
 import pathlib
 
-import numpy as np
-
 try:
     import mujoco
-    from ..gl_viewer import GLViewer
+
+    from ..gl_viewer import GLViewer, OffScreenRenderer
     from ..notebook_viewer import NotebookViewer
 except ImportError:
     raise ImportError(
@@ -13,8 +12,9 @@ except ImportError:
 
 
 class ManipulationBaseEnv:
-
-    def __init__(self, headless=False, **kwargs):
+    def __init__(
+        self, headless=False, use_d405_camera=False, show_cam_renders=False, **kwargs
+    ):
         model_path = str(pathlib.Path(__file__).parent / "xmls/scene.xml")
         self.model = mujoco.MjModel.from_xml_path(model_path)
         self.data = mujoco.MjData(self.model)
@@ -28,11 +28,25 @@ class ManipulationBaseEnv:
         else:
             self.viewer = None
 
+        self.d405_viewer = None
+
+        if use_d405_camera:
+            self.cam_id = mujoco.mj_name2id(
+                self.model, mujoco.mjtObj.mjOBJ_CAMERA, "camera_d405"
+            )
+            if show_cam_renders:
+                self.d405_viewer = GLViewer(self.model, self.data, cam_id=self.cam_id)
+            else:
+                self.d405_viewer = OffScreenRenderer(self.model, self.data, self.cam_id)
+
     def step(self, action=None, n_frames=20):
         self.render()
         mujoco.mj_step(self.model, self.data, nstep=n_frames)
         mujoco.mj_rnePostConstraint(self.model, self.data)
 
+        if self.d405_viewer:
+            self.d405_viewer.render()
+
     def render(self):
         if self.viewer:
             self.viewer.render()

tinysim_mujoco/manipulation/push_env.py

@@ -0,0 +1,268 @@
+import numpy as np
+
+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 ManipulationEnvV0:
+
+    def __init__(self, **kwargs):
+        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")
+        )
+
+        return {
+            "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])