hansen1416 · October 9, 2023 02:29
diff --git a/mujoco_joints.py b/mujoco_joints.py
 import time
 import mujoco
 import mujoco.viewer
 import math
 from transforms3d import quaternions

 from lib.BaseRender import BaseRender


 class BaseRender:

    def __init__(self, xml) -> None:
        self.model = mujoco.MjModel.from_xml_string(xml)
        self.data = mujoco.MjData(self.model)

        # enable joint visualization option:
        scene_option = mujoco.MjvOption()
        scene_option.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True

    def run(self, init_callback=None, step_callback=None):
        mujoco.mj_kinematics(self.model, self.data)
        mujoco.mj_forward(self.model, self.data)

        with mujoco.viewer.launch_passive(self.model, self.data) as viewer:

            # set camera
            viewer.cam.lookat[0] = 0  # x position
            viewer.cam.lookat[1] = 0  # y position
            viewer.cam.lookat[2] = 0  # z position
            viewer.cam.distance = 3  # distance from the target
            viewer.cam.elevation = -30  # elevation angle
            viewer.cam.azimuth = 0  # azimuth angle

            # random initial rotational velocity:
            mujoco.mj_resetData(self.model, self.data)

            if callable(init_callback):
                init_callback(self.model, self.data)

            # Close the viewer automatically after 30 wall-seconds.
            start = time.time()

            while viewer.is_running() and time.time() - start < 6:

                step_start = time.time()

                mujoco.mj_step(self.model, self.data)

                if callable(step_callback):
                    step_callback(self.model, self.data)

                # Pick up changes to the physics state, apply perturbations, update options from GUI.
                viewer.sync()

                # Rudimentary time keeping, will drift relative to wall clock.
                time_until_next_step = self.model.opt.timestep - \
                    (time.time() - step_start)

                if time_until_next_step > 0:
                    time.sleep(time_until_next_step)
                    
 # `free` joint
 # This joint type is only allowed in bodies that are children of the world body.
 # No other joints can be defined in the body if a free joint is defined.
 free_joint = """
 <mujoco>
    <asset>
        <texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3"
        rgb2=".2 .3 .4" width="300" height="300"/>
        <material name="grid" texture="grid" texrepeat="8 8" reflectance=".2"/>
    </asset>

    <worldbody>
        <light pos="0 0 .6"/>
        <geom size=".3 .3 .01" type="plane" material="grid"/>
        <body name="bat" pos="0 0 0">
            <joint name="free" type="free"/>
            <geom name="bat" type="capsule" fromto="0 0 .04 0 -.3 .04"
            size=".04" rgba="1 1 1 1"/>
        </body>
    </worldbody>
 </mujoco>
 """

 # `ball` joint
 # The rotation is around the point defined by the pos attribute below.
 # If a body has a ball joint, it cannot have other rotational joints (ball or hinge).
 # Combining ball joints with slide joints in the same body is allowed.
 ball_joint = """
 <mujoco>
    <asset>
        <texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3"
        rgb2=".2 .3 .4" width="300" height="300"/>
        <material name="grid" texture="grid" texrepeat="8 8" reflectance=".2"/>
    </asset>

    <worldbody>
        <light pos="0 0 .6"/>
        <geom size=".3 .3 .01" type="plane" material="grid"/>
        <body name="bat" pos="0 0 0">
            <joint name="ball" type="ball" pos="0 0 0"/>
            <geom name="bat" type="capsule" fromto="0 0 .04 0 -.3 .04"
            size=".04" rgba="1 1 1 1"/>
        </body>
    </worldbody>
 </mujoco>
 """

 # `slide` joint
 # The slide type creates a sliding or prismatic joint with one translational degree of freedom.
 # Such joints are defined by a position and a sliding direction.
 # For simulation purposes only the direction is needed;
 # the joint position is used for rendering purposes.
 slide_joint = """
 <mujoco>
    <asset>
        <texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3"
        rgb2=".2 .3 .4" width="300" height="300"/>
        <material name="grid" texture="grid" texrepeat="8 8" reflectance=".2"/>
    </asset>

    <worldbody>
        <light pos="0 0 .6"/>
        <geom size=".3 .3 .01" type="plane" material="grid"/>
        <body name="bat" pos="0 0 0">
            <joint name="slide" type="slide" axis="0 1 0"/>
            <geom name="bat" type="capsule" fromto="0 0 .04 0 -.3 .04"
            size=".04" rgba="1 1 1 1"/>
        </body>
    </worldbody>
 </mujoco>
 """

 # 'hinge' joint
 # The hinge type creates a hinge joint with one rotational degree of freedom.
 # The rotation takes place around a specified axis through a specified position.
 # This is the most common type of joint and is therefore the default.
 hinge_joint = """
 <mujoco>
    <asset>
        <texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3"
        rgb2=".2 .3 .4" width="300" height="300"/>
        <material name="grid" texture="grid" texrepeat="8 8" reflectance=".2"/>
    </asset>

    <worldbody>
        <light pos="0 0 .6"/>
        <geom size=".3 .3 .01" type="plane" material="grid"/>
        <body name="bat" pos="0 0 0">
            <joint name="hinge" type="hinge" axis="0 0 1" pos="0 -0.2 0"/>
            <geom name="bat" type="capsule" fromto="0 0 .04 0 -.3 .04"
            size=".04" rgba="1 1 1 1"/>
        </body>
    </worldbody>
 </mujoco>
 """


 def free_init_callback(mode, data):
    # For a free joint, data.qpos has 6 elements: the first 3 are the Cartesian coordinates of the body,
    # and the last 3 are the Euler angles of the body’s orientation12.
    print(data.qpos)

    data.qpos[0] = 0.  # x
    data.qpos[1] = 0.  # y
    data.qpos[2] = 0.  # z, height

    data.qpos[3] = 0  # x, rotation
    data.qpos[4] = 0  # y, rotation
    data.qpos[5] = 0  # z, rotation

    # data.qvel has 6 elements: the first 3 are the linear velocities of the body,
    # and the last 3 are the angular velocities of the body.
    print(data.qvel)


 def free_step_callback(mode, data):

    data.qvel[5] = 1


 def ball_init_callback(mode, data):
    # The ball type creates a ball joint with three rotational degrees of freedom.
    # qpos of ball joint is a quaternion, w,x,y,z
    # q.vel are the angular velocities of the body.

    q = quaternions.axangle2quat(
        [0, 0, 1], math.radians(90), is_normalized=True)

    data.qpos[0] = q[0]  # w
    data.qpos[1] = q[1]  # x
    data.qpos[2] = q[2]  # y
    data.qpos[3] = q[3]  # z

    # data.qvel[1] = 3

    print(data.qpos)  # [1. 0. 0. 0.]
    print(data.qvel)  # [0. 0. 0.]


 def ball_step_callback(mode, data):

    data.qvel[0] = -1  # x rotation
    data.qvel[1] = 0  # y rotation
    data.qvel[2] = 0  # z rotation


 def slide_init_callback(mode, data):
    # The slide type creates a sliding or prismatic joint with one translational degree of freedom.
    data.qpos[0] = -1

    print(data.qpos)  # [0.]
    print(data.qvel)  # [0.]


 def slide_step_callback(mode, data):
    # along the axis of the slide joint
    data.qvel[0] = 1


 def hinge_init_callback(mode, data):
    # The slide type creates a sliding or prismatic joint with one translational degree of freedom.
    data.qpos[0] = math.pi/4

    print(data.qpos)  # [0.]
    print(data.qvel)  # [0.]


 def hinge_step_callback(mode, data):
    # along the axis of the slide joint
    data.qvel[0] = 1


 if __name__ == "__main__":
    rnder = BaseRender(hinge_joint)

    rnder.run(init_callback=hinge_init_callback,
              step_callback=hinge_step_callback)
	import time
	import mujoco
	import mujoco.viewer
	import math
	from transforms3d import quaternions

	from lib.BaseRender import BaseRender


	class BaseRender:

	def __init__(self, xml) -> None:
	self.model = mujoco.MjModel.from_xml_string(xml)
	self.data = mujoco.MjData(self.model)

	# enable joint visualization option:
	scene_option = mujoco.MjvOption()
	scene_option.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True

	def run(self, init_callback=None, step_callback=None):
	mujoco.mj_kinematics(self.model, self.data)
	mujoco.mj_forward(self.model, self.data)

	with mujoco.viewer.launch_passive(self.model, self.data) as viewer:

	# set camera
	viewer.cam.lookat[0] = 0 # x position
	viewer.cam.lookat[1] = 0 # y position
	viewer.cam.lookat[2] = 0 # z position
	viewer.cam.distance = 3 # distance from the target
	viewer.cam.elevation = -30 # elevation angle
	viewer.cam.azimuth = 0 # azimuth angle

	# random initial rotational velocity:
	mujoco.mj_resetData(self.model, self.data)

	if callable(init_callback):
	init_callback(self.model, self.data)

	# Close the viewer automatically after 30 wall-seconds.
	start = time.time()

	while viewer.is_running() and time.time() - start < 6:

	step_start = time.time()

	mujoco.mj_step(self.model, self.data)

	if callable(step_callback):
	step_callback(self.model, self.data)

	# Pick up changes to the physics state, apply perturbations, update options from GUI.
	viewer.sync()

	# Rudimentary time keeping, will drift relative to wall clock.
	time_until_next_step = self.model.opt.timestep - \
	(time.time() - step_start)

	if time_until_next_step > 0:
	time.sleep(time_until_next_step)

	# `free` joint
	# This joint type is only allowed in bodies that are children of the world body.
	# No other joints can be defined in the body if a free joint is defined.
	free_joint = """
	<mujoco>
	<asset>
	<texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3"
	rgb2=".2 .3 .4" width="300" height="300"/>
	<material name="grid" texture="grid" texrepeat="8 8" reflectance=".2"/>
	</asset>

	<worldbody>
	<light pos="0 0 .6"/>
	<geom size=".3 .3 .01" type="plane" material="grid"/>
	<body name="bat" pos="0 0 0">
	<joint name="free" type="free"/>
	<geom name="bat" type="capsule" fromto="0 0 .04 0 -.3 .04"
	size=".04" rgba="1 1 1 1"/>
	</body>
	</worldbody>
	</mujoco>
	"""

	# `ball` joint
	# The rotation is around the point defined by the pos attribute below.
	# If a body has a ball joint, it cannot have other rotational joints (ball or hinge).
	# Combining ball joints with slide joints in the same body is allowed.
	ball_joint = """
	<mujoco>
	<asset>
	<texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3"
	rgb2=".2 .3 .4" width="300" height="300"/>
	<material name="grid" texture="grid" texrepeat="8 8" reflectance=".2"/>
	</asset>

	<worldbody>
	<light pos="0 0 .6"/>
	<geom size=".3 .3 .01" type="plane" material="grid"/>
	<body name="bat" pos="0 0 0">
	<joint name="ball" type="ball" pos="0 0 0"/>
	<geom name="bat" type="capsule" fromto="0 0 .04 0 -.3 .04"
	size=".04" rgba="1 1 1 1"/>
	</body>
	</worldbody>
	</mujoco>
	"""

	# `slide` joint
	# The slide type creates a sliding or prismatic joint with one translational degree of freedom.
	# Such joints are defined by a position and a sliding direction.
	# For simulation purposes only the direction is needed;
	# the joint position is used for rendering purposes.
	slide_joint = """
	<mujoco>
	<asset>
	<texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3"
	rgb2=".2 .3 .4" width="300" height="300"/>
	<material name="grid" texture="grid" texrepeat="8 8" reflectance=".2"/>
	</asset>

	<worldbody>
	<light pos="0 0 .6"/>
	<geom size=".3 .3 .01" type="plane" material="grid"/>
	<body name="bat" pos="0 0 0">
	<joint name="slide" type="slide" axis="0 1 0"/>
	<geom name="bat" type="capsule" fromto="0 0 .04 0 -.3 .04"
	size=".04" rgba="1 1 1 1"/>
	</body>
	</worldbody>
	</mujoco>
	"""

	# 'hinge' joint
	# The hinge type creates a hinge joint with one rotational degree of freedom.
	# The rotation takes place around a specified axis through a specified position.
	# This is the most common type of joint and is therefore the default.
	hinge_joint = """
	<mujoco>
	<asset>
	<texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3"
	rgb2=".2 .3 .4" width="300" height="300"/>
	<material name="grid" texture="grid" texrepeat="8 8" reflectance=".2"/>
	</asset>

	<worldbody>
	<light pos="0 0 .6"/>
	<geom size=".3 .3 .01" type="plane" material="grid"/>
	<body name="bat" pos="0 0 0">
	<joint name="hinge" type="hinge" axis="0 0 1" pos="0 -0.2 0"/>
	<geom name="bat" type="capsule" fromto="0 0 .04 0 -.3 .04"
	size=".04" rgba="1 1 1 1"/>
	</body>
	</worldbody>
	</mujoco>
	"""


	def free_init_callback(mode, data):
	# For a free joint, data.qpos has 6 elements: the first 3 are the Cartesian coordinates of the body,
	# and the last 3 are the Euler angles of the body’s orientation12.
	print(data.qpos)

	data.qpos[0] = 0. # x
	data.qpos[1] = 0. # y
	data.qpos[2] = 0. # z, height

	data.qpos[3] = 0 # x, rotation
	data.qpos[4] = 0 # y, rotation
	data.qpos[5] = 0 # z, rotation

	# data.qvel has 6 elements: the first 3 are the linear velocities of the body,
	# and the last 3 are the angular velocities of the body.
	print(data.qvel)


	def free_step_callback(mode, data):

	data.qvel[5] = 1


	def ball_init_callback(mode, data):
	# The ball type creates a ball joint with three rotational degrees of freedom.
	# qpos of ball joint is a quaternion, w,x,y,z
	# q.vel are the angular velocities of the body.

	q = quaternions.axangle2quat(
	[0, 0, 1], math.radians(90), is_normalized=True)

	data.qpos[0] = q[0] # w
	data.qpos[1] = q[1] # x
	data.qpos[2] = q[2] # y
	data.qpos[3] = q[3] # z

	# data.qvel[1] = 3

	print(data.qpos) # [1. 0. 0. 0.]
	print(data.qvel) # [0. 0. 0.]


	def ball_step_callback(mode, data):

	data.qvel[0] = -1 # x rotation
	data.qvel[1] = 0 # y rotation
	data.qvel[2] = 0 # z rotation


	def slide_init_callback(mode, data):
	# The slide type creates a sliding or prismatic joint with one translational degree of freedom.
	data.qpos[0] = -1

	print(data.qpos) # [0.]
	print(data.qvel) # [0.]


	def slide_step_callback(mode, data):
	# along the axis of the slide joint
	data.qvel[0] = 1


	def hinge_init_callback(mode, data):
	# The slide type creates a sliding or prismatic joint with one translational degree of freedom.
	data.qpos[0] = math.pi/4

	print(data.qpos) # [0.]
	print(data.qvel) # [0.]


	def hinge_step_callback(mode, data):
	# along the axis of the slide joint
	data.qvel[0] = 1


	if __name__ == "__main__":
	rnder = BaseRender(hinge_joint)

	rnder.run(init_callback=hinge_init_callback,
	step_callback=hinge_step_callback)