Utility functions
Contents
Function summaries
Vector operations
Rotates the vectors with the rotation matrix. |
|
Rotates the vectors with the transpose of the rotation matrix. |
|
Adds two vectors with weights. |
|
Normalizes an array of vectors, such that each vector has unit length. |
MuJoCo access utilities
Convenience function to get geom ids.
Convenience function to get body ids.
Returns all geom ids belonging to a given body.
Returns the subtree of the body structure that has the provided body as its root.
Get sensor data from the sensor with the provided name.
Get the indices in the sensordata array corresponding to the given sensor.
Joint Manipulation utilities
Sets the joint position for the joint with name joint_name.
Get the indices in the qpos array corresponding to the given joint.
Get the indices in the qvel array corresponding to the given joint.
Sets the angle from default at which the joint will be locked.
Locks a joint to a fixed angle.
Unlocks a given joint.
MuJoCo coordinate frame utilities
Returns the position of geom in the world frame.
Returns the position of body in the world frame.
Returns the rotation matrix that rotates the geoms frame to the world frame.
Returns the rotation matrix that rotates the bodies frame to the world frame.
Converts position from the world coordinate frame to a geom specific frame.
Converts position from the world coordinate frame to a body specific frame.
Converts position from the geom specific coordinate frame to the world frame.
Converts position from the body specific coordinate frame to the world frame.
Converts position from the geom specific coordinate frame to the frame of a specific body.
Converts position from the body specific coordinate frame to the frame of a specific geom.
Converts position from one geoms coordinate frame to another.
Converts position from one bodies coordinate frame to another.
Converts a vectors direction from a geoms specific coordinate frame to the world frame.
Converts a vectors direction from a bodies specific coordinate frame to the world frame.
Converts a vectors direction from the world coordinate frame to a geoms specific frame.
Converts a vectors direction from the world coordinate frame to a bodies specific frame.
Converts a vectors direction from a geoms coordinate frame to a bodies frame.
Converts a vectors direction from a bodies coordinate frame to a geoms frame.
Converts a vectors direction from one geoms coordinate frame to another.
Converts a vectors direction from one bodies coordinate frame to another.
Plotting utilities
Plots an array of points.
Plots an array of points and vectors pointing from those points.
Assorted functions
Distribute a target mass over multiple bodies.
Data fields
- mimoEnv.utils.MUJOCO_JOINT_SIZES
Size of qpos entries for each joint type; free, ball, slide, hinge.
- mimoEnv.utils.MUJOCO_DOF_SIZES
Size of qvel entries for each joint type; free, ball, slide, hinge.
Detail documentation
- mimoEnv.utils.MUJOCO_JOINT_SIZES
Size of qpos entries for each joint type; free, ball, slide, hinge.
- mimoEnv.utils.MUJOCO_DOF_SIZES
Size of qvel entries for each joint type; free, ball, slide, hinge.
- mimoEnv.utils.rotate_vector(vector: numpy.ndarray, rot_matrix: numpy.ndarray)
Rotates the vectors with the rotation matrix.
The vector can be a 1d vector or a multidimensional array of vectors, as long as the final dimension has length 3. Convention for mujoco matrices: Use this function to convert from special frame to global and
rotate_vector_transpose()for the inverse rotation. The exception are the contact frames, which are transposed.- Parameters
vector (numpy.ndarray) – The vector(s). Must have shapes (3,) or (.., 3).
rot_matrix (numpy.ndarray) – The rotation matrix that will be applied to the vectors. Should be a (3,3) array.
- Returns
The rotated vector(s).
- Return type
numpy.ndarray
- mimoEnv.utils.rotate_vector_transpose(vector, rot_matrix)
Rotates the vectors with the transpose of the rotation matrix.
Works identical to rotate_vector, but transposes the rotation matrix first.
- Parameters
vector (numpy.ndarray) – The vector(s). Must have shapes (3,) or (.., 3).
rot_matrix (numpy.ndarray) – The rotation matrix that will be applied to the vectors. Should be a (3,3) array.
- Returns
The rotated vector(s).
- Return type
numpy.ndarray
- mimoEnv.utils.weighted_sum_vectors(vector1, vector2, weight1, weight2)
Adds two vectors with weights.
- mimoEnv.utils.normalize_vectors(vectors)
Normalizes an array of vectors, such that each vector has unit length.
- Parameters
vectors (numpy.ndarray) – The array of vectors. The last dimension should iterate over the elements of the vectors.
- Returns
The normalized vectors. Same shape as the input, but the length of each vector is reduced to 1.
- Return type
numpy.ndarray
- mimoEnv.utils.get_geom_id(mujoco_model, geom_id=None, geom_name=None)
Convenience function to get geom ids.
MuJoCo geoms can be referred to by either an id or a name. This function wraps this ambiguity and always returns the id of a geom when either is specified. If both an id and a name are specified the name is ignored!
- mimoEnv.utils.get_body_id(mujoco_model, body_id=None, body_name=None)
Convenience function to get body ids.
Works identical to
get_geom_id()
- mimoEnv.utils.get_geoms_for_body(sim_model, body_id)
Returns all geom ids belonging to a given body.
- mimoEnv.utils.get_child_bodies(sim_model, body_id)
Returns the subtree of the body structure that has the provided body as its root.
The body structure is defined in the MuJoCo XMLs. This function returns a list containing the ids of all descendant bodies of a given body, including the given body.
- mimoEnv.utils.get_data_for_sensor(mujoco_model, mujoco_data, sensor_name)
Get sensor data from the sensor with the provided name.
- Parameters
mujoco_model (mujoco.MjModel) – The MuJoCo model object.
mujoco_data (mujoco.MjData) – The MuJoCo data object.
sensor_name (str) – The name of the sensor.
- Returns
The output values of the sensor. The shape will depend on the sensor type.
- Return type
numpy.ndarray
- mimoEnv.utils.get_sensor_addr(mujoco_model, sensor_id)
Get the indices in the sensordata array corresponding to the given sensor.
- mimoEnv.utils.set_joint_qpos(mujoco_model, mujoco_data, joint_name, qpos)
Sets the joint position for the joint with name joint_name.
Directly sets the joint to the position provided by qpos. Note that the shape of qpos must match the joint! A free joint for example has length 7. The sizes for all types can be found in
MUJOCO_JOINT_SIZES.
- mimoEnv.utils.get_joint_qpos_addr(mujoco_model, joint_id)
Get the indices in the qpos array corresponding to the given joint.
- mimoEnv.utils.get_joint_qvel_addr(mujoco_model, joint_id)
Get the indices in the qvel array corresponding to the given joint.
- mimoEnv.utils.set_joint_locking_angle(mujoco_model, joint_name, angle, constraint_id=None)
Sets the angle from default at which the joint will be locked.
The angle is in radians, and can be positive or negative. This function does not lock or unlock a joint, merely changes the angle. This function requires that there be a constraint already existing the scene XML. This is the case for MIMo by default, with each joint having a constraint of the same name that is disabled at initialization.
- Parameters
mujoco_model (mujoco.MjModel) – The MuJoCo model object.
joint_name (str) – The name of the joint.
angle (float|ndarray) – The locking angle(s) in radians, as a delta from the model starting value.
constraint_id (int|ndarray) – If the ID(s) of the constraint is already known the id lookup can be bypassed by passing it here.
- mimoEnv.utils.lock_joint(mujoco_model, joint_name, joint_angle=None)
Locks a joint to a fixed angle.
This function utilizes MuJoCos equality constraints to achieve the locking effect, requiring that there be a constraint already existing the scene XML. This is the case for MIMo by default, with each joint having a constraint of the same name that is disabled at initialization. In effect this function enables the equality constraint with same name as the argument.
- Parameters
mujoco_model (mujoco.MjModel) – The MuJoCo model object.
joint_name (str) – The name of the joint.
joint_angle (float) – The locking angle in radians, as a delta from the model starting value. The angle that the joint will be locked to can be set separately using
set_joint_locking_angle(). By default, joints are locked into the value they have in the scene xml.
- mimoEnv.utils.unlock_joint(mujoco_model, joint_name)
Unlocks a given joint.
See
lock_joint().- Parameters
mujoco_model (mujoco.MjModel) – The MuJoCo model object.
joint_name (str) – The name of the joint.
- mimoEnv.utils.get_geom_position(mujoco_data, geom_id)
Returns the position of geom in the world frame.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
geom_id (int) – The id of the geom.
- Returns
The position of the geom in the world frame. Shape (3,).
- Return type
numpy.ndarray
- mimoEnv.utils.get_body_position(mujoco_data, body_id)
Returns the position of body in the world frame.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
body_id (int) – The id of the body.
- Returns
The position of the body in the world frame. Shape (3,).
- Return type
numpy.ndarray
- mimoEnv.utils.get_geom_rotation(mujoco_data, geom_id)
Returns the rotation matrix that rotates the geoms frame to the world frame.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
geom_id (int) – The id of the geom.
- Returns
A (3,3) array containing the rotation matrix.
- Return type
numpy.ndarray
- mimoEnv.utils.get_body_rotation(mujoco_data, body_id)
Returns the rotation matrix that rotates the bodies frame to the world frame.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
body_id (int) – The id of the body.
- Returns
A (3,3) array containing the rotation matrix.
- Return type
numpy.ndarray
- mimoEnv.utils.world_pos_to_geom(mujoco_data, position, geom_id)
Converts position from the world coordinate frame to a geom specific frame.
Position can be a vector or an array of vectors such that the last dimension has size 3.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
position (numpy.ndarray) – Array containing position(s). Its shape should be either (3,) or (.., 3).
geom_id (int) – The id of the geom.
- Returns
Array of the same shape as the input array with the converted coordinates.
- Return type
numpy.ndarray
- mimoEnv.utils.world_pos_to_body(mujoco_data, position, body_id)
Converts position from the world coordinate frame to a body specific frame.
Position can be a vector or an array of vectors such that the last dimension has size 3.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
position (numpy.ndarray) – Array containing position(s). Its shape should be either (3,) or (.., 3).
body_id (int) – The id of the geom.
- Returns
Array of the same shape as the input array with the converted coordinates.
- Return type
numpy.ndarray
- mimoEnv.utils.geom_pos_to_world(mujoco_data, position, geom_id)
Converts position from the geom specific coordinate frame to the world frame.
Position can be a vector or an array of vectors such that the last dimension has size 3.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
position (numpy.ndarray) – Array containing position(s). Its shape should be either (3,) or (.., 3).
geom_id (int) – The id of the geom.
- Returns
Array of the same shape as the input array with the converted coordinates.
- Return type
numpy.ndarray
- mimoEnv.utils.body_pos_to_world(mujoco_data, position, body_id)
Converts position from the body specific coordinate frame to the world frame.
Position can be a vector or an array of vectors such that the last dimension has size 3.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
position (numpy.ndarray) – Array containing position(s). Its shape should be either (3,) or (.., 3).
body_id (int) – The id of the body.
- Returns
Array of the same shape as the input array with the converted coordinates.
- Return type
numpy.ndarray
- mimoEnv.utils.geom_pos_to_body(mujoco_data, position, geom_id, body_id)
Converts position from the geom specific coordinate frame to the frame of a specific body.
Position can be a vector or an array of vectors such that the last dimension has size 3.
- Parameters
- Returns
Array of the same shape as the input array with the converted coordinates.
- Return type
numpy.ndarray
- mimoEnv.utils.body_pos_to_geom(mujoco_data, position, body_id, geom_id)
Converts position from the body specific coordinate frame to the frame of a specific geom.
Position can be a vector or an array of vectors such that the last dimension has size 3.
- Parameters
- Returns
Array of the same shape as the input array with the converted coordinates.
- Return type
numpy.ndarray
- mimoEnv.utils.geom_pos_to_geom(mujoco_data, position, geom_id_source, geom_id_target)
Converts position from one geoms coordinate frame to another.
Position can be a vector or an array of vectors such that the last dimension has size 3.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
position (numpy.ndarray) – Array containing position(s). Its shape should be either (3,) or (.., 3).
geom_id_source (int) – The id of the geom for the initial coordinate frame.
geom_id_target (int) – The id of the geom for the output coordinate frame.
- Returns
Array of the same shape as the input array with the converted coordinates.
- Return type
numpy.ndarray
- mimoEnv.utils.body_pos_to_body(mujoco_data, position, body_id_source, body_id_target)
Converts position from one bodies coordinate frame to another.
Position can be a vector or an array of vectors such that the last dimension has size 3.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
position (numpy.ndarray) – Array containing position(s). Its shape should be either (3,) or (.., 3).
body_id_source (int) – The id of the body for the initial coordinate frame.
body_id_target (int) – The id of the body for the output coordinate frame.
- Returns
Array of the same shape as the input array with the converted coordinates.
- Return type
numpy.ndarray
- mimoEnv.utils.geom_rot_to_world(mujoco_data, vector, geom_id)
Converts a vectors direction from a geoms specific coordinate frame to the world frame.
Unlike the functions for the positions, this only converts the direction.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
vector (numpy.ndarray) – A vector or array of vectors. Shape must be either (3,) or (.., 3).
geom_id (int) – The id of the geom.
- Returns
Array of the same shape as the input array with the rotated vector.
- Return type
numpy.ndarray
- mimoEnv.utils.body_rot_to_world(mujoco_data, vector, body_id)
Converts a vectors direction from a bodies specific coordinate frame to the world frame.
Unlike the functions for the positions, this only converts the direction.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
vector (numpy.ndarray) – A vector or array of vectors. Shape must be either (3,) or (.., 3).
body_id (int) – The id of the body.
- Returns
Array of the same shape as the input array with the rotated vector.
- Return type
numpy.ndarray
- mimoEnv.utils.world_rot_to_geom(mujoco_data, vector, geom_id)
Converts a vectors direction from the world coordinate frame to a geoms specific frame.
Unlike the functions for the positions, this only converts the direction.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
vector (numpy.ndarray) – A vector or array of vectors. Shape must be either (3,) or (.., 3).
geom_id (int) – The id of the geom.
- Returns
Array of the same shape as the input array with the rotated vector.
- Return type
numpy.ndarray
- mimoEnv.utils.world_rot_to_body(mujoco_data, vector, body_id)
Converts a vectors direction from the world coordinate frame to a bodies specific frame.
Unlike the functions for the positions, this only converts the direction.
- Parameters
mujoco_data (mujoco.MjData) – The MuJoCo data object.
vector (numpy.ndarray) – A vector or array of vectors. Shape must be either (3,) or (.., 3).
body_id (int) – The id of the body.
- Returns
Array of the same shape as the input array with the rotated vector.
- Return type
numpy.ndarray
- mimoEnv.utils.geom_rot_to_body(mujoco_data, vector, geom_id, body_id)
Converts a vectors direction from a geoms coordinate frame to a bodies frame.
Unlike the functions for the positions, this only converts the direction.
- Parameters
- Returns
Array of the same shape as the input array with the rotated vector.
- Return type
numpy.ndarray
- mimoEnv.utils.body_rot_to_geom(mujoco_data, vector, body_id, geom_id)
Converts a vectors direction from a bodies coordinate frame to a geoms frame.
Unlike the functions for the positions, this only converts the direction.
- Parameters
- Returns
Array of the same shape as the input array with the rotated vector.
- Return type
numpy.ndarray
- mimoEnv.utils.geom_rot_to_geom(mujoco_data, vector, geom_id_source, geom_id_target)
Converts a vectors direction from one geoms coordinate frame to another.
Unlike the functions for the positions, this only converts the direction.
- Parameters
- Returns
Array of the same shape as the input array with the rotated vector.
- Return type
numpy.ndarray
- mimoEnv.utils.body_rot_to_body(mujoco_data, vector, body_id_source, body_id_target)
Converts a vectors direction from one bodies coordinate frame to another.
Unlike the functions for the positions, this only converts the direction.
- Parameters
- Returns
Array of the same shape as the input array with the rotated vector.
- Return type
numpy.ndarray
- mimoEnv.utils.plot_points(points, limit=1.0, title='', show=True)
Plots an array of points.
- Parameters
points (numpy.ndarray) – An array containing points. Shape should be (n, 3) for n points.
limit (float) – The limit that is applied to the axis. Default 1.
title (str) – The title for the plot. Empty by default.
show (bool) – If
Truethe plot is rendered to a window, ifFalsethe figure and axis objects are returned instead.
- Returns
A tuple (fig, ax) containing the pyplot figure and axis objects if show is
False,Noneotherwise.- Return type
Tuple[plt.Figure, plt.Axes]|None
- mimoEnv.utils.plot_forces(points, vectors, limit=1.0, title='', show=True)
Plots an array of points and vectors pointing from those points.
The arrays points and vectors must have the same shape. For each point there is a vector, plotted as arrows, the direction and size of which is determined by the vectors argument, starting from that point.
- Parameters
points (numpy.ndarray) – An array containing the points. Shape should be (n, 3) for n points.
vectors (numpy.ndarray) – An array of vectors with one for each point. Shape should be (n, 3) for n points.
limit (float) – The limit that is applied to the axis. Default 1.
title (str) – The title for the plot. Empty by default.
show (bool) – If
Truethe plot is rendered to a window, ifFalsethe figure and axis objects are returned instead.
- Returns
A tuple (fig, ax) containing the pyplot figure and axis objects if show is
False,Noneotherwise.- Return type
Tuple[plt.Figure, plt.Axes]|None
- mimoEnv.utils.determine_geom_masses(mujoco_model, mujoco_data, body_ids, target_mass, print_out=False)
Distribute a target mass over multiple bodies.
Given a list of bodies and a desired target mass, calculate the mass of component geoms assuming identical density such that the total mass of the bodies matches the target mass. This function takes account of overlap between geoms within each body, but not between bodies.
- Parameters
mujoco_model (mujoco.MjModel) – The MuJoCo model object.
mujoco_data (mujoco.MjData) – The MuJoCo data object.
body_ids (List[int]) – A list of bodies by ID over which the mass will be distributed.
target_mass (float) – The target mass.
print_out (bool) – If
True, target masses and body names are printed to console.
- Returns
A dictionary with body names as keys and a list of geom masses as values.
- Return type