diff --git a/dm_control/locomotion/walkers/assets/__init__.py b/dm_control/locomotion/walkers/assets/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dm_control/locomotion/walkers/assets/dog_v2/__init__.py b/dm_control/locomotion/walkers/assets/dog_v2/__init__.py new file mode 100644 index 00000000..422d4bb5 --- /dev/null +++ b/dm_control/locomotion/walkers/assets/dog_v2/__init__.py @@ -0,0 +1,21 @@ +import build_torso +import build_neck +import add_markers +import build_tail +import create_skin +import build_front_legs +import build_back_legs +import add_muscles +import add_wrapping_geoms +import add_torque_actuators + +build_torso = build_torso.create_torso +build_neck = build_neck.create_neck +add_markers = add_markers.add_markers +build_tail = build_tail.create_tail +create_skin = create_skin.create +build_back_legs = build_back_legs.create_back_legs +build_front_legs = build_front_legs.create_front_legs +add_muscles = add_muscles.add_muscles +add_wrapping_geoms = add_wrapping_geoms.add_wrapping_geoms +add_torque_actuators = add_torque_actuators.add_motors diff --git a/dm_control/locomotion/walkers/assets/dog_v2/add_markers.py b/dm_control/locomotion/walkers/assets/dog_v2/add_markers.py new file mode 100644 index 00000000..16ad0427 --- /dev/null +++ b/dm_control/locomotion/walkers/assets/dog_v2/add_markers.py @@ -0,0 +1,75 @@ +# Copyright 2023 The dm_control Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +"""Add markers to the dog model, used for motion capture tracking.""" + +MARKERS_PER_BODY = { + "torso": ( + (-0.25, 0, 0.13), + (-0.1, 0, 0.13), + (0.05, 0, 0.15), + (-0.25, -0.06, -0.01), + (-0.1, -0.08, -0.02), + (0.05, -0.09, -0.02), + (-0.25, 0.06, -0.01), + (-0.1, 0.08, -0.02), + (0.05, 0.09, -0.02), + ), + "C_6": ((0.05, -0.05, 0), (0.05, 0.05, 0),), + "upper_arm": ((0.04, 0, 0),), + "lower_arm": ((0.025, 0, 0), (0.025, 0, -0.1),), + "hand": ((0, -0.02, 0), (0, 0.02, 0), (0.02, -0.02, -0.08), (0.02, 0.02, -0.08),), + "finger": ((0.065, 0, 0),), + "pelvis": ((0.04, -0.07, 0.02), (0.04, 0.07, 0.02),), + "upper_leg": ((-0.03, 0.04, 0),), + "lower_leg": ((0.04, 0, 0), (-0.04, 0.035, 0),), + "foot": ((0, -0.02, 0), (0, 0.02, 0), (0, -0.02, -0.1), (0, 0.03, -0.1),), + "toe": ((0.055, 0, 0),), + "Ca_10": ((0, 0, 0.01),), + "skull": ((0.01, 0, 0.04), (0, -0.05, 0), (0, 0.05, 0), (-0.1, 0, 0.02),), +} + + +def add_markers(model): + """Add markers to the given model. + + Args: + model: The model to add markers to. + """ + bodies = model.find_all("body") + + total_markers = 0 + for body in bodies: + for name, markers_pos in MARKERS_PER_BODY.items(): + if name in body.name and "anchor" not in body.name: + marker_idx = 0 + for pos in markers_pos: + pos = list(pos) + if "_R" in body.name: + pos[1] *= -1 + body.add( + "site", + name="marker_" + body.name + "_" + str(marker_idx), + pos=pos, + dclass="marker", + ) + + marker_idx += 1 + total_markers += 1 + + for i in range(total_markers): + marker_body = model.worldbody.add( + "body", name="marker_" + str(i), mocap=True) + marker_body.add("site", name="marker_" + str(i), dclass="mocap_marker") diff --git a/dm_control/locomotion/walkers/assets/dog_v2/add_muscles.py b/dm_control/locomotion/walkers/assets/dog_v2/add_muscles.py new file mode 100644 index 00000000..30747a42 --- /dev/null +++ b/dm_control/locomotion/walkers/assets/dog_v2/add_muscles.py @@ -0,0 +1,537 @@ +# Copyright 2023 The dm_control Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +"""Add muscle actuators to the dog model.""" + +from os import listdir +from os.path import isfile, join + +import numpy as np +import trimesh +from scipy.spatial import KDTree +from tqdm import tqdm +from utils import array_to_string, slices2paths + +from dm_control import mjcf as mjcf_module +from dm_control.locomotion.walkers.assets.dog_v2 import muscles as muscles_constants +from dm_control.mujoco.wrapper import mjbindings +from dm_control.utils.transformations import quat_to_mat + +MUSCLE_LEGS = ( + muscles_constants.EXTENSORS_BACK + + muscles_constants.EXTENSORS_FRONT + + muscles_constants.FLEXORS_BACK + + muscles_constants.FLEXORS_FRONT +) + +WRAP_GEOMS_LEGS = ( + "shoulder_L_wrapping", + "elbow_L_wrapping", + "wrist_L_wrapping", + "finger_L_wrapping", + "hip_L_wrapping", + "knee_L_wrapping", + "toe_L_wrapping", + "shoulder_R_wrapping", + "elbow_R_wrapping", + "wrist_R_wrapping", + "finger_R_wrapping", + "hip_R_wrapping", + "knee_R_wrapping", + "toe_R_wrapping", +) + + +def get_closest_geom(kd_tree, point, mjcf, mtu): + """Finds the closest geometric body to a given point using a KD-tree. + + Args: + kd_tree : A dictionary containing KD-trees for different bones. + point : The coordinates of the target point. + mjcf: The MJCF (MuJoCo XML) object. + mtu: The muscle-tendon unit (MTU) identifier. + + Returns: + closest_geom_body: The closest geometric body (as a Mujoco body object) + to the target point. + """ + dist = np.inf + closest_geom_body = 0 + for bone, tree in kd_tree.items(): + dist_new, i = tree.query(np.array([point]), k=1) + geom = mjcf.find("geom", bone) + body = geom.parent + + if mtu in MUSCLE_LEGS: + if body.name not in [ + "upper_leg_L", + "upper_leg_R", + "lower_leg_L", + "lower_leg_R", + "foot_L", + "foot_R", + "toe_L", + "toe_R", + "scapula_L", + "scapula_R", + "upper_arm_L", + "upper_arm_R", + "lower_arm_L", + "lower_arm_R", + "hand_L", + "hand_R", + "finger_L", + "finger_R", + "pelvis", + ]: + continue + else: + if body.name in [ + "upper_leg_L", + "upper_leg_R", + "lower_leg_L", + "lower_leg_R", + "foot_L", + "foot_R", + "toe_L", + "toe_R", + "lower_arm_L", + "lower_arm_R", + "hand_L", + "hand_R", + "finger_L", + "finger_R", + ]: + continue + + if dist_new < dist: + dist = dist_new + closest_geom_body = body + + return closest_geom_body + + +def calculate_transformation(element): + """Calculates the transformation matrix from the root to the given element. + + This function iterates through the parent elements of the given element and + calculates the transformation matrices from the root until reaching the + given element. It considers the position and orientation (quaternion) + attributes of each element to construct the transformation matrices. + The final transformation matrix represents the cumulative transformation + from the root to the given element. + + Args: + element: The MuJoCo element for which the transformation matrix + is calculated. + + Returns: + T: The transformation matrix from the root to the given element. + """ + # Calculate all transformation matrices from root until this element + all_transformations = [] + + while type(element) == mjcf_module.element._ElementImpl: + if element.pos is not None: + pos = np.array(element.pos, dtype=float) + else: + pos = np.zeros(3) + + if element.quat is not None: + rot = quat_to_mat(element.get("quat")) + else: + rot = np.eye(3) + + all_transformations.append( + np.vstack( + (np.hstack((rot, pos.reshape([-1, 1]))), np.array([0, 0, 0, 1]))) + ) + + element = element.parent + + # Apply all transformations + transform = np.eye(4) + for transformation in reversed(all_transformations): + transform = np.matmul(transform, transformation) + + return transform + + +def add_muscles( + model, scale_multiplier, muscle_dynamics, asset_dir, lengthrange_from_joints +): + """Add muscles to a Mujoco (MJCF) model. + + This function adds muscles to a Mujoco model, + creating MTUs (muscle-tendon units) and defining + their properties. + + Args: + model (MJCFModel): The input Mujoco model to which muscles will be added. + scale_multiplier (float): A scaling factor for muscle forces (0 to disable anatomical scaling). + muscle_dynamics (str): Muscle dynamics model, either 'Millard', 'Sigmoid', or 'General'. + asset_dir (str): The directory path containing muscle and bone assets. + lengthrange_from_joints (bool): If True, compute length ranges from joint limits. + + Returns: + None + + Note: + - The function modifies the input `model` in place by adding muscles and related properties. + - The specific muscles added depend on the provided `muscle_dynamics` and other parameters. + + Raises: + NameError: If an unsupported `muscle_dynamics` value is provided. + """ + physics = mjcf_module.Physics.from_mjcf_model(model) + muscle_meshes_path = join(asset_dir, "muscles/") + bones_path = asset_dir + + bones = [ + f for f in listdir(bones_path) if isfile(join(bones_path, f)) and f[0] != "." + ] + + bones_kd_trees = {} + bones_meshes = {} + for bone in bones: + if ( + "BONE" in bone + and "simpl" not in bone + and "Lingual" not in bone + and "cartilage" not in bone + ): + bone_mesh = trimesh.load_mesh(bones_path + "/" + bone, process=False) + bones_kd_trees[bone[4:-4]] = KDTree(bone_mesh.vertices) + bones_meshes[bone[4:-4]] = bone_mesh + + muscle = model.default.default["muscle"] + muscle.tendon.rgba = [0.5, 0, 0, 1] + + length_range = model.compiler.lengthrange + length_range.inttotal = 5 + length_range.accel = 50 + length_range.interval = 2 + length_range.timestep = 0.001 + length_range.timeconst = 1 + length_range.tolrange = 0.05 + + muscles = ( + muscles_constants.EXTENSORS_FRONT + + muscles_constants.FLEXORS_FRONT + + muscles_constants.EXTENSORS_BACK + + muscles_constants.FLEXORS_BACK + + muscles_constants.TORSO + + muscles_constants.NECK + + muscles_constants.TAIL + ) + + used_muscles = [] + volumes = [] + spatials = [] + # create MTUs + for mtu in muscles: + # Load stl file of site + m = trimesh.load_mesh(muscle_meshes_path + mtu + ".stl") + + volumes.append(m.volume) + # check along which axis to slice + if mtu not in muscles_constants.LATERAL: + plane_normal = [0, 0, -1] + start = np.argmax(m.vertices[:, 2]) + end = np.argmin(m.vertices[:, 2]) + else: + plane_normal = [1, 0, 0] + start = np.argmin(m.vertices[:, 0]) + end = np.argmax(m.vertices[:, 0]) + + muscle_length = np.linalg.norm(m.vertices[start] - m.vertices[end]) + + if muscle_length < 0.1: + print("muscle too short", mtu) + continue + + heights = np.linspace(start=-0.5, stop=muscle_length, + num=1000, dtype=float) + + slices = m.section_multiplane( + plane_origin=m.vertices[start], plane_normal=plane_normal, heights=heights + ) + + paths = slices2paths(mtu, slices, muscle_length) + spatial = model.tendon.add( + "spatial", name=f"{mtu}_tendon", dclass="muscle") + spatials.append(spatial) + used_muscles.append(mtu) + + counter = 0 # used for site naming + prev_body = None + prev_point = None + + for idx in range(len(paths)): + point = paths[idx][0] + closest_geom_body = get_closest_geom(bones_kd_trees, point, model, mtu) + + # Add site to tendon + min_dist = 0.05 + max_dist = 0.11 + if mtu in muscles_constants.TAIL: + max_dist = 0.05 + elif mtu in muscles_constants.NECK: + min_dist = 0.11 + max_dist = 0.15 + + d = 0 + if prev_point is not None: + d = np.linalg.norm(prev_point - point) + + # we add a site to the tendon path if: + # 1. it's the last or first site + # 2. the geometry we are attaching it to is different from the previous + # one, and we have passed the min distance between sites + # 3. we passed the max dist between sites + if ( + (closest_geom_body != prev_body and d > min_dist) + or idx == len(paths) - 1 + or idx == 0 + or d > max_dist + ): + site_matrix = np.eye(4) + site_matrix[:3, 3] = point + body_matrix = calculate_transformation(closest_geom_body) + site_matrix = np.matmul(np.linalg.inv(body_matrix), site_matrix) + + site_name = mtu + "_" + str(counter) + # Create the site + closest_geom_body.add( + "site", + name=site_name, + pos=array_to_string(site_matrix[:3, 3]), + group="4", + dclass="connector", + ) + + counter += 1 + spatial.add("site", site=site_name) + + prev_point = point + prev_body = closest_geom_body + + # check for wrapping geoms + if idx != len(paths) - 1: + geoms = physics.named.data.geom_xpos.axes.row.names + dist = np.inf + closest_g = None + closest_geom_pos = None + for g_name in geoms: + if "wrapping" in g_name: + pos = physics.named.data.geom_xpos[g_name] + new_dist = np.linalg.norm(pos - point) + if ( + ( + mtu in muscles_constants.LATERAL + or mtu in muscles_constants.NECK + ) + and g_name not in WRAP_GEOMS_LEGS + and pos[2] < point[2] + and new_dist < dist + ): + closest_g = g_name + dist = new_dist + closest_geom_pos = pos + elif ( + ( + mtu in muscles_constants.FLEXORS_BACK + or mtu in muscles_constants.FLEXORS_FRONT + or mtu in muscles_constants.EXTENSORS_FRONT + or mtu in muscles_constants.EXTENSORS_BACK + ) + and pos[2] < point[2] + and dist > 0.01 + and new_dist < dist + and g_name in WRAP_GEOMS_LEGS + ): + # in the legs we are interested only in wrapping + # geometries that are lower than the site we + # are adding + closest_g = g_name + dist = new_dist + closest_geom_pos = pos + + if dist < 0.1: + if ( + mtu in muscles_constants.FLEXORS_BACK + or mtu in muscles_constants.FLEXORS_FRONT + ) and closest_geom_pos[2] < point[2]: + spatial.add( + "geom", geom=closest_g, sidesite=closest_g + "_backward" + ) + elif ( + mtu in muscles_constants.EXTENSORS_BACK + or mtu in muscles_constants.EXTENSORS_FRONT + ) and closest_geom_pos[2] < point[2]: + spatial.add( + "geom", geom=closest_g, sidesite=closest_g + "_forward" + ) + elif mtu in muscles_constants.NECK: + pos1 = physics.named.data.site_xpos[closest_g + "_forward"] + d1 = np.linalg.norm(pos1 - point) + pos2 = physics.named.data.site_xpos[closest_g + "_backward"] + d2 = np.linalg.norm(pos2 - point) + if d1 < d2: + spatial.add( + "geom", + geom=closest_g, + sidesite=closest_g + "_forward", + ) + else: + spatial.add( + "geom", + geom=closest_g, + sidesite=closest_g + "_backward", + ) + elif mtu in muscles_constants.TORSO and dist < 0.05: + spatial.add( + "geom", geom=closest_g, sidesite=closest_g + "_forward" + ) + + print("Added tendon {}".format(spatial.name)) + + physics = mjcf_module.Physics.from_mjcf_model(model) + # Compute length ranges: + # This computation of lenght ranges differs from the computation + # provided in MuJoCo. MuJoCo contracts the muscles + # individually, and records it's minimum and maximum length. However this + # method relies on the assumption that the muscles dynamics and parameters + # are perfectly set up. This is not the case, so the method below is used. + # This method randomly samples the joint limits and records the minimum and + # maximum length of the muscles. Sampling is necessary because some muscles + # cross multiple joints (Biarticular muscles). + if lengthrange_from_joints: + vector_min = np.ones(physics.model.ntendon) * np.inf + vector_max = np.ones(physics.model.ntendon) * -np.inf + + for _ in tqdm(range(500000)): + hinge = mjbindings.enums.mjtJoint.mjJNT_HINGE + slide = mjbindings.enums.mjtJoint.mjJNT_SLIDE + + for joint_id in range(physics.model.njnt): + joint_name = physics.model.id2name(joint_id, "joint") + joint_type = physics.model.jnt_type[joint_id] + is_limited = physics.model.jnt_limited[joint_id] + range_min, range_max = physics.model.jnt_range[joint_id] + + if is_limited and (joint_type == hinge or joint_type == slide): + physics.named.data.qpos[joint_name] = np.random.uniform( + range_min, range_max + ) + else: + continue + + physics.forward() + vector_min = np.minimum(vector_min, physics.data.ten_length[:]) + vector_max = np.maximum(vector_max, physics.data.ten_length[:]) + physics.reset() + + for tend_idx, spatial in enumerate(spatials): + if muscle_dynamics in ["Millard", "Sigmoid"]: + muscle = model.actuator.add( + "general", + tendon=spatial, + name=spatial.name, + dclass=spatial.dclass, + dyntype="muscle", + gaintype="muscle", + biastype="muscle", + ) + prms = [0] * 10 + prms[0] = 0.01 + prms[1] = 0.04 + if muscle_dynamics == "Millard": + muscle.dynprm = prms + else: + prms[2] = 2.0 + muscle.dynprm = prms + + if spatial.name[: -len("_tendon")] in muscles_constants.NECK: + scale = 1500 + else: + scale = 8000 + + # range(2), force, scale, lmin, lmax, vmax, fpmax, fvmax + gainprm = [0.75, 1.05, -1, scale, 0.5, 1.6, 1.5, 1.3, 1.2, 0] + muscle.gainprm = gainprm + muscle.biasprm = gainprm + muscle.ctrllimited = True + muscle.ctrlrange = [0.0, 1.0] + if lengthrange_from_joints: + muscle.lengthrange = [vector_min[tend_idx], vector_max[tend_idx]] + elif muscle_dynamics == "General": + muscle = model.actuator.add( + "general", tendon=spatial, name=spatial.name, dclass=spatial.dclass + ) + muscle.ctrllimited = True + muscle.ctrlrange = [0.0, 1.0] + muscle.gaintype = "affine" + muscle.dyntype = "muscle" + muscle.dynprm = [0.01, 0.04, 0.001] + muscle.gainprm = [-200, -50, -10] + if lengthrange_from_joints: + muscle.lengthrange = [vector_min[tend_idx], vector_max[tend_idx]] + else: + raise NameError(muscle_dynamics) + + everything_ok = False + while not everything_ok: + try: + physics = mjcf_module.Physics.from_mjcf_model(model) + everything_ok = True + + # assign length ranges and forces + if muscle_dynamics in ["Millard", "Sigmoid"]: + lr = physics.named.model.actuator_lengthrange + forces = None + + # compute anatomical forces if we provide a scale_multiplier > 0 + if scale_multiplier > 0: + volumes = np.array(volumes) + lm = [] + for mtu in used_muscles: + mtu_name = mtu + "_tendon" + L0 = (lr[mtu_name, 1] - lr[mtu_name, 0]) / (1.05 - 0.75) + LM = physics.named.model.actuator_length0[mtu_name] - ( + lr[mtu_name, 0] - 0.75 * L0 + ) + lm.append(LM) + lm = np.array(lm) + forces = (volumes / lm) * scale_multiplier + + actuators = model.find_all("actuator") + # assign forces and length ranges + for idx, muscle in enumerate(actuators): + muscle.lengthrange = lr[idx] + if forces is not None: + muscle.gainprm[2] = forces[idx] + + except Exception as inst: + print(inst) + s = str(inst) + s = s.split() + i = s.index("actuator") + muscle_index = s[i + 1] + if ":" in muscle_index: + muscle_index = muscle_index[:-1] + + muscle_index = int(muscle_index) + del model.actuator.general[muscle_index] diff --git a/dm_control/locomotion/walkers/assets/dog_v2/add_torque_actuators.py b/dm_control/locomotion/walkers/assets/dog_v2/add_torque_actuators.py index 71c66fb4..13f98783 100644 --- a/dm_control/locomotion/walkers/assets/dog_v2/add_torque_actuators.py +++ b/dm_control/locomotion/walkers/assets/dog_v2/add_torque_actuators.py @@ -33,18 +33,18 @@ def add_motors(physics, model, lumbar_joints, cervical_joints, caudal_joints): """ # Fixed Tendons: spinal_joints = collections.OrderedDict() - spinal_joints['lumbar_'] = lumbar_joints - spinal_joints['cervical_'] = cervical_joints - spinal_joints['caudal_'] = caudal_joints + spinal_joints["lumbar_"] = lumbar_joints + spinal_joints["cervical_"] = cervical_joints + spinal_joints["caudal_"] = caudal_joints tendons = [] for region in spinal_joints.keys(): - for direction in ['extend', 'bend', 'twist']: + for direction in ["extend", "bend", "twist"]: joints = [ joint for joint in spinal_joints[region] if direction in joint.name ] if joints: tendon = model.tendon.add( - 'fixed', name=region + direction, dclass=joints[0].dclass + "fixed", name=region + direction, dclass=joints[0].dclass ) tendons.append(tendon) joint_inertia = physics.bind(joints).M0 @@ -52,34 +52,34 @@ def add_motors(physics, model, lumbar_joints, cervical_joints, caudal_joints): coefs /= coefs.sum() coefs *= len(joints) for i, joint in enumerate(joints): - tendon.add('joint', joint=joint, coef=coefs[i]) + tendon.add("joint", joint=joint, coef=coefs[i]) # Actuators: all_spinal_joints = [] for region in spinal_joints.values(): all_spinal_joints.extend(region) - root_joint = model.find('joint', 'root') + + root_joint = model.find("joint", "root") actuated_joints = [ joint - for joint in model.find_all('joint') + for joint in model.find_all("joint") if joint not in all_spinal_joints and joint is not root_joint ] for tendon in tendons: gain = 0.0 for joint in tendon.joint: # joint.joint.user = physics.bind(joint.joint).damping - def_joint = model.default.find('default', joint.joint.dclass) + def_joint = model.default.find("default", joint.joint.dclass) j_gain = def_joint.general.gainprm or def_joint.parent.general.gainprm gain += j_gain[0] * joint.coef gain /= len(tendon.joint) model.actuator.add( - 'general', tendon=tendon, name=tendon.name, dclass=tendon.dclass + "general", tendon=tendon, name=tendon.name, dclass=tendon.dclass ) for joint in actuated_joints: - model.actuator.add( - 'general', joint=joint, name=joint.name, dclass=joint.dclass - ) + model.actuator.add("general", joint=joint, + name=joint.name, dclass=joint.dclass) return actuated_joints diff --git a/dm_control/locomotion/walkers/assets/dog_v2/add_wrapping_geoms.py b/dm_control/locomotion/walkers/assets/dog_v2/add_wrapping_geoms.py new file mode 100644 index 00000000..deb7356b --- /dev/null +++ b/dm_control/locomotion/walkers/assets/dog_v2/add_wrapping_geoms.py @@ -0,0 +1,378 @@ +# Copyright 2023 The dm_control Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + + +def add_wrapping_geoms(model): + """ Adds wrapping geometries and sites to the specified model. + + Args: + model: The model to which the wrapping geometries and sites will be added. + """ + pelvis = model.find("body", "pelvis") + pelvis.add( + "geom", + name="hip_R_wrapping", + dclass="wrapping", + size=[0.05, 0.04], + pos=[0, 0.05, -0.02], + ) + pelvis.add( + "site", + name="hip_R_wrapping_forward", + dclass="wrapping", + pos=[0.08, -0.05, -0.00], + ) + pelvis.add( + "site", + name="hip_R_wrapping_backward", + dclass="wrapping", + pos=[-0.04, -0.05, -0.04], + ) + + pelvis.add( + "geom", + name="hip_L_wrapping", + dclass="wrapping", + size=[0.05, 0.04], + pos=[0, -0.05, -0.02], + ) + pelvis.add( + "site", + name="hip_L_wrapping_forward", + dclass="wrapping", + pos=[0.08, 0.05, -0.00], + ) + pelvis.add( + "site", + name="hip_L_wrapping_backward", + dclass="wrapping", + pos=[-0.04, 0.05, -0.04], + ) + + C_2 = model.find("body", "C_2") + C_2.add( + "geom", + name="C_2_wrapping", + dclass="wrapping", + pos=[0.02, 0, 0.03], + size=[0.015, 0.04], + ) + C_2.add("site", name="C_2_wrapping_forward", + pos=[0.05, 0, 0.02], dclass="wrapping") + C_2.add( + "site", name="C_2_wrapping_backward", pos=[-0.005, 0, 0.04], dclass="wrapping" + ) + C_7 = model.find("body", "C_7") + C_7.add("geom", name="C_7_wrapping", dclass="wrapping", size=[0.03, 0.07]) + + Ca_1 = model.find("body", "Ca_1") + Ca_1.add("geom", name="Ca_1_wrapping", size=[0.008, 0.04], dclass="wrapping") + Ca_1.add("site", name="Ca_1_wrapping_forward", + pos=[0, 0, -0.05], dclass="wrapping") + Ca_1.add("site", name="Ca_1_wrapping_backward", + pos=[0, 0, 0.05], dclass="wrapping") + + torso = model.find("body", "torso") + torso.add( + "site", name="C_7_wrapping_forward", pos=[0.19, 0, 0.055], dclass="wrapping" + ) + torso.add( + "site", name="C_7_wrapping_backward", pos=[0.09, 0, 0.12], dclass="wrapping" + ) + + for i in range(1, 8, 1): + body = model.find("body", f"L_{i}") + body.add("geom", name=f"L_{i}_wrapping", size=[ + 0.009, 0.04], dclass="wrapping") + body.add( + "site", + name=f"L_{i}_wrapping_forward", + pos=[0, 0, 0.1], + dclass="wrapping", + ) + body.add( + "site", + name=f"L_{i}_wrapping_backward", + pos=[0, 0, -0.1], + dclass="wrapping", + ) + + # Left Side + upper_leg_L = model.find("body", "upper_leg_L") + upper_leg_L.add( + "site", + name="knee_L_wrapping_forward", + pos=[0.07, 0.0311, -0.175], + dclass="wrapping", + ) + upper_leg_L.add( + "site", + name="knee_L_wrapping_backward", + pos=[-0.09, 0.0311, 0], + dclass="wrapping", + ) + + lower_leg_L = model.find("body", "lower_leg_L") + lower_leg_L.add( + "geom", name="knee_L_wrapping", dclass="wrapping", size=[0.02, 0.05] + ) + lower_leg_L.add( + "site", + name="ankle_L_wrapping_forward", + pos=[-0.08, 0, -0.13824], + dclass="wrapping", + ) + lower_leg_L.add( + "site", + name="ankle_L_wrapping_backward", + pos=[-0.25, 0, -0.13824], + dclass="wrapping", + ) + + foot_L = model.find("body", "foot_L") + foot_L.add("geom", name="ankle_L_wrapping", + dclass="wrapping", size=[0.02, 0.05]) + foot_L.add( + "site", + name="toe_L_wrapping_forward", + pos=[0.07, 0, -0.11993], + dclass="wrapping", + ) + foot_L.add( + "site", + name="toe_L_wrapping_backward", + pos=[-0.07, 0, -0.11993], + dclass="wrapping", + ) + + toe_L = model.find("body", "toe_L") + toe_L.add("geom", name="toe_L_wrapping", + dclass="wrapping", size=[0.01, 0.05]) + + scapula_L = model.find("body", "scapula_L") + scapula_L.add( + "geom", + name="shoulder_L_wrapping", + pos=[0.075, 0.033, -0.13], + size=[0.015, 0.05], + dclass="wrapping", + ) + scapula_L.add( + "site", + name="shoulder_L_wrapping_forward", + pos=[0.1, 0.033, -0.13], + dclass="wrapping", + ) + scapula_L.add( + "site", + name="shoulder_L_wrapping_backward", + pos=[0.02, 0.033, -0.13], + dclass="wrapping", + ) + upper_arm_L = model.find("body", "upper_arm_L") + upper_arm_L.add( + "geom", + name="elbow_L_wrapping", + pos=[-0.05, 0.015, -0.135], + size=[0.015, 0.05], + dclass="wrapping", + ) + upper_arm_L.add( + "site", + name="elbow_L_wrapping_forward", + pos=[0.03, 0.015, -0.15], + dclass="wrapping", + ) + upper_arm_L.add( + "site", + name="elbow_L_wrapping_backward", + pos=[-0.1, 0.015, -0.15], + dclass="wrapping", + ) + lower_arm_L = model.find("body", "lower_arm_L") + lower_arm_L.add( + "site", + name="wrist_L_wrapping_forward", + pos=[0.015, -0.015, -0.18], + dclass="wrapping", + ) + lower_arm_L.add( + "site", + name="wrist_L_wrapping_backward", + pos=[-0.05, -0.015, -0.19], + dclass="wrapping", + ) + hand_anchor_L = model.find("body", "hand_anchor_L") + hand_anchor_L.add( + "geom", + name="wrist_L_wrapping", + size=[0.011, 0.05], + pos=[0, 0, 0.012], + dclass="wrapping", + ) + hand_L = model.find("body", "hand_L") + hand_L.add( + "geom", + name="finger_L_wrapping", + size=[0.015, 0.05], + pos=[0.01, 0, -0.05], + dclass="wrapping", + ) + hand_L.add( + "site", + name="finger_L_wrapping_forward", + pos=[0.05, 0, -0.06], + dclass="wrapping", + ) + hand_L.add( + "site", + name="finger_L_wrapping_backward", + pos=[-0.05, 0, -0.06], + dclass="wrapping", + ) + + # Right Side + upper_leg_R = model.find("body", "upper_leg_R") + upper_leg_R.add( + "site", + name="knee_R_wrapping_forward", + pos=[0.07, -0.0311, -0.175], + dclass="wrapping", + ) + upper_leg_R.add( + "site", + name="knee_R_wrapping_backward", + pos=[-0.09, -0.0311, 0], + dclass="wrapping", + ) + + lower_leg_R = model.find("body", "lower_leg_R") + lower_leg_R.add( + "geom", name="knee_R_wrapping", dclass="wrapping", size=[0.02, 0.05] + ) + lower_leg_R.add( + "site", + name="ankle_R_wrapping_forward", + pos=[-0.08, 0, -0.13824], + dclass="wrapping", + ) + lower_leg_R.add( + "site", + name="ankle_R_wrapping_backward", + pos=[-0.25, 0, -0.13824], + dclass="wrapping", + ) + + foot_R = model.find("body", "foot_R") + foot_R.add("geom", name="ankle_R_wrapping", + dclass="wrapping", size=[0.02, 0.05]) + foot_R.add( + "site", + name="toe_R_wrapping_forward", + pos=[0.07, 0, -0.11993], + dclass="wrapping", + ) + foot_R.add( + "site", + name="toe_R_wrapping_backward", + pos=[-0.07, 0, -0.11993], + dclass="wrapping", + ) + + toe_R = model.find("body", "toe_R") + toe_R.add("geom", name="toe_R_wrapping", + dclass="wrapping", size=[0.01, 0.05]) + + scapula_R = model.find("body", "scapula_R") + scapula_R.add( + "geom", + name="shoulder_R_wrapping", + pos=[0.075, -0.033, -0.13], + size=[0.015, 0.05], + dclass="wrapping", + ) + scapula_R.add( + "site", + name="shoulder_R_wrapping_forward", + pos=[0.1, -0.033, -0.13], + dclass="wrapping", + ) + scapula_R.add( + "site", + name="shoulder_R_wrapping_backward", + pos=[0.02, -0.033, -0.13], + dclass="wrapping", + ) + upper_arm_R = model.find("body", "upper_arm_R") + upper_arm_R.add( + "geom", + name="elbow_R_wrapping", + pos=[-0.05, -0.015, -0.135], + size=[0.015, 0.05], + dclass="wrapping", + ) + upper_arm_R.add( + "site", + name="elbow_R_wrapping_forward", + pos=[0.03, -0.015, -0.15], + dclass="wrapping", + ) + upper_arm_R.add( + "site", + name="elbow_R_wrapping_backward", + pos=[-0.1, -0.015, -0.15], + dclass="wrapping", + ) + lower_arm_R = model.find("body", "lower_arm_R") + lower_arm_R.add( + "site", + name="wrist_R_wrapping_forward", + pos=[0.015, 0.015, -0.18], + dclass="wrapping", + ) + lower_arm_R.add( + "site", + name="wrist_R_wrapping_backward", + pos=[-0.05, 0.015, -0.19], + dclass="wrapping", + ) + hand_anchor_R = model.find("body", "hand_anchor_R") + hand_anchor_R.add( + "geom", + name="wrist_R_wrapping", + size=[0.011, 0.05], + pos=[0, 0, 0.012], + dclass="wrapping", + ) + hand_R = model.find("body", "hand_R") + hand_R.add( + "geom", + name="finger_R_wrapping", + size=[0.015, 0.05], + pos=[0.01, 0, -0.05], + dclass="wrapping", + ) + hand_R.add( + "site", + name="finger_R_wrapping_forward", + pos=[0.05, 0, -0.06], + dclass="wrapping", + ) + hand_R.add( + "site", + name="finger_R_wrapping_backward", + pos=[-0.05, 0, -0.06], + dclass="wrapping", + ) diff --git a/dm_control/locomotion/walkers/assets/dog_v2/build_back_legs.py b/dm_control/locomotion/walkers/assets/dog_v2/build_back_legs.py index bbc36e2c..058ef2ba 100644 --- a/dm_control/locomotion/walkers/assets/dog_v2/build_back_legs.py +++ b/dm_control/locomotion/walkers/assets/dog_v2/build_back_legs.py @@ -15,9 +15,10 @@ """Make back legs for the dog model.""" -from dm_control import mjcf import numpy as np +from dm_control import mjcf + def create_back_legs( model, @@ -49,67 +50,67 @@ def create_back_legs( # Hip joint sites: scale = np.asarray([bone_size[bone] for bone in pelvic_bones]).mean() hip_pos = np.array((-0.23, -0.6, -0.16)) * scale - for side in ['_L', '_R']: + for side in ["_L", "_R"]: pelvis.add( - 'site', name='hip' + side, size=[0.011], pos=hip_pos * side_sign[side] + "site", name="hip" + side, size=[0.011], pos=hip_pos * side_sign[side] ) # Upper legs: upper_leg = {} - femurs = [b for b in bones if 'Fem' in b] + femurs = [b for b in bones if "Fem" in b] use_tendons = False if not use_tendons: - femurs += [b for b in bones if 'Patella' in b] - for side in ['_L', '_R']: + femurs += [b for b in bones if "Patella" in b] + for side in ["_L", "_R"]: body_pos = hip_pos * side_sign[side] - leg = pelvis.add('body', name='upper_leg' + side, pos=body_pos) + leg = pelvis.add("body", name="upper_leg" + side, pos=body_pos) upper_leg[side] = leg for bone in [b for b in femurs if side in b]: leg.add( - 'geom', + "geom", name=bone, mesh=bone, - pos=-bone_position['Pelvis'] - body_pos, - dclass='bone', + pos=-bone_position["Pelvis"] - body_pos, + dclass="bone", ) # Hip joints - for dof in ['_supinate', '_abduct', '_extend']: + for dof in ["_supinate", "_abduct", "_extend"]: axis = primary_axis[dof].copy() - if dof != '_extend': - axis *= 1.0 if side != '_R' else -1.0 - leg.add('joint', name='hip' + side + dof, dclass='hip' + dof, axis=axis) + if dof != "_extend": + axis *= 1.0 if side != "_R" else -1.0 + leg.add("joint", name="hip" + side + dof, dclass="hip" + dof, axis=axis) # Knee sites - scale = bone_size['Femoris_L'] + scale = bone_size["Femoris_L"] knee_pos = np.array((-0.2, -0.27, -1.45)) * scale leg.add( - 'site', - type='cylinder', - name='knee' + side, + "site", + type="cylinder", + name="knee" + side, size=[0.003, 0.02], zaxis=(0, 1, 0), pos=knee_pos * side_sign[side], ) pos = np.array((-0.01, -0.02, -0.08)) * side_sign[side] - euler = [-10 * (1.0 if side == '_R' else -1.0), 20, 0] + euler = [-10 * (1.0 if side == "_R" else -1.0), 20, 0] leg.add( - 'geom', - name=leg.name + '0_collision', + "geom", + name=leg.name + "0_collision", pos=pos, size=[0.04, 0.08], euler=euler, - dclass='collision_primitive', + dclass="collision_primitive", ) pos = np.array((-0.03, 0, -0.05)) - euler = [-10 * (1.0 if side == '_R' else -1.0), 5, 0] + euler = [-10 * (1.0 if side == "_R" else -1.0), 5, 0] leg.add( - 'geom', - name=leg.name + '1_collision', + "geom", + name=leg.name + "1_collision", pos=pos, size=[0.04, 0.04], euler=euler, - dclass='collision_primitive', + dclass="collision_primitive", ) # Patella @@ -119,29 +120,29 @@ def create_back_legs( # Lower legs: lower_leg = {} - lower_leg_bones = [b for b in bones if 'Tibia_' in b or 'Fibula' in b] - for side in ['_L', '_R']: + lower_leg_bones = [b for b in bones if "Tibia_" in b or "Fibula" in b] + for side in ["_L", "_R"]: body_pos = knee_pos * side_sign[side] - leg = upper_leg[side].add('body', name='lower_leg' + side, pos=body_pos) + leg = upper_leg[side].add("body", name="lower_leg" + side, pos=body_pos) lower_leg[side] = leg for bone in [b for b in lower_leg_bones if side in b]: leg.add( - 'geom', + "geom", name=bone, mesh=bone, - pos=-bone_position['Pelvis'] - upper_leg[side].pos - body_pos, - dclass='bone', + pos=-bone_position["Pelvis"] - upper_leg[side].pos - body_pos, + dclass="bone", ) # Knee joints - leg.add('joint', name='knee' + side, dclass='knee', axis=(0, -1, 0)) + leg.add("joint", name="knee" + side, dclass="knee", axis=(0, -1, 0)) # Ankle sites - scale = bone_size['Tibia_L'] + scale = bone_size["Tibia_L"] ankle_pos = np.array((-1.27, 0.04, -0.98)) * scale leg.add( - 'site', - type='cylinder', - name='ankle' + side, + "site", + type="cylinder", + name="ankle" + side, size=[0.003, 0.013], zaxis=(0, 1, 0), pos=ankle_pos * side_sign[side], @@ -149,133 +150,132 @@ def create_back_legs( # Feet: foot = {} - foot_bones = [b for b in bones if 'tars' in b.lower() or 'tuber' in b] - for side in ['_L', '_R']: + foot_bones = [b for b in bones if "tars" in b.lower() or "tuber" in b] + for side in ["_L", "_R"]: body_pos = ankle_pos * side_sign[side] - leg = lower_leg[side].add('body', name='foot' + side, pos=body_pos) + leg = lower_leg[side].add("body", name="foot" + side, pos=body_pos) foot[side] = leg for bone in [b for b in foot_bones if side in b]: leg.add( - 'geom', + "geom", name=bone, mesh=bone, - pos=-bone_position['Pelvis'] + pos=-bone_position["Pelvis"] - upper_leg[side].pos - lower_leg[side].pos - body_pos, - dclass='bone', + dclass="bone", ) # Ankle joints - leg.add('joint', name='ankle' + side, dclass='ankle', axis=(0, 1, 0)) + leg.add("joint", name="ankle" + side, dclass="ankle", axis=(0, 1, 0)) pos = np.array((-0.01, -0.005, -0.05)) * side_sign[side] leg.add( - 'geom', - name=leg.name + '_collision', + "geom", + name=leg.name + "_collision", size=[0.015, 0.07], pos=pos, - dclass='collision_primitive', + dclass="collision_primitive", ) # Toe sites - scale = bone_size['Metatarsi_R_2'] + scale = bone_size["Metatarsi_R_2"] toe_pos = np.array((-0.37, -0.2, -2.95)) * scale leg.add( - 'site', - type='cylinder', - name='toe' + side, + "site", + type="cylinder", + name="toe" + side, size=[0.003, 0.025], zaxis=(0, 1, 0), pos=toe_pos * side_sign[side], ) # Toes: - toe_bones = [b for b in bones if 'Phalange' in b] + toe_bones = [b for b in bones if "Phalange" in b] toe_geoms = [] sole_sites = [] nails = [] - for side in ['_L', '_R']: + for side in ["_L", "_R"]: body_pos = toe_pos * side_sign[side] foot_anchor = foot[side].add( - 'body', name='foot_anchor' + side, pos=body_pos - ) + "body", name="foot_anchor" + side, pos=body_pos) foot_anchor.add( - 'geom', + "geom", name=foot_anchor.name, - dclass='foot_primitive', - type='box', + dclass="foot_primitive", + type="box", size=(0.005, 0.005, 0.005), contype=0, conaffinity=0, ) - foot_anchor.add('site', name=foot_anchor.name, dclass='sensor') - leg = foot_anchor.add('body', name='toe' + side) + foot_anchor.add("site", name=foot_anchor.name, dclass="sensor") + leg = foot_anchor.add("body", name="toe" + side) for bone in [b for b in toe_bones if side in b]: geom = leg.add( - 'geom', + "geom", name=bone, mesh=bone, - pos=-bone_position['Pelvis'] + pos=-bone_position["Pelvis"] - upper_leg[side].pos - lower_leg[side].pos - foot[side].pos - body_pos, - dclass='bone', + dclass="bone", ) - if 'B_3' in bone: + if "B_3" in bone: nails.append(bone) - geom.dclass = 'visible_bone' + geom.dclass = "visible_bone" else: toe_geoms.append(geom) # Toe joints - leg.add('joint', name='toe' + side, dclass='toe', axis=(0, 1, 0)) + leg.add("joint", name="toe" + side, dclass="toe", axis=(0, 1, 0)) # Collision geoms leg.add( - 'geom', - name=leg.name + '0_collision', + "geom", + name=leg.name + "0_collision", size=[0.018, 0.012], pos=[0.015, 0, -0.02], euler=(90, 0, 0), - dclass='foot_primitive', + dclass="foot_primitive", ) leg.add( - 'geom', - name=leg.name + '1_collision', + "geom", + name=leg.name + "1_collision", size=[0.01, 0.015], pos=[0.035, 0, -0.028], euler=(90, 0, 0), - dclass='foot_primitive', + dclass="foot_primitive", ) leg.add( - 'geom', - name=leg.name + '2_collision', + "geom", + name=leg.name + "2_collision", size=[0.008, 0.01], pos=[0.045, 0, -0.03], euler=(90, 0, 0), - dclass='foot_primitive', + dclass="foot_primitive", ) sole = leg.add( - 'site', - name='sole' + side, + "site", + name="sole" + side, size=(0.025, 0.03, 0.008), pos=(0.026, 0, -0.033), - type='box', - dclass='sensor', + type="box", + dclass="sensor", ) sole_sites.append(sole) physics = mjcf.Physics.from_mjcf_model(model) - for side in ['_L', '_R']: + for side in ["_L", "_R"]: # lower leg: leg = lower_leg[side] leg.add( - 'geom', - name=leg.name + '_collision', + "geom", + name=leg.name + "_collision", pos=physics.bind(leg).ipos * 1.3, size=[0.02, 0.1], quat=physics.bind(leg).iquat, - dclass='collision_primitive', + dclass="collision_primitive", ) return nails, sole_sites diff --git a/dm_control/locomotion/walkers/assets/dog_v2/build_dog.py b/dm_control/locomotion/walkers/assets/dog_v2/build_dog.py index ffcaea18..02a96641 100644 --- a/dm_control/locomotion/walkers/assets/dog_v2/build_dog.py +++ b/dm_control/locomotion/walkers/assets/dog_v2/build_dog.py @@ -17,44 +17,68 @@ import os -from absl import app -from absl import flags -from dm_control import mjcf -from dm_control.locomotion.walkers.assets.dog_v2 import add_torque_actuators -from dm_control.locomotion.walkers.assets.dog_v2 import build_back_legs -from dm_control.locomotion.walkers.assets.dog_v2 import build_front_legs -from dm_control.locomotion.walkers.assets.dog_v2 import build_neck -from dm_control.locomotion.walkers.assets.dog_v2 import build_tail -from dm_control.locomotion.walkers.assets.dog_v2 import build_torso -from dm_control.locomotion.walkers.assets.dog_v2 import create_skin -from lxml import etree import numpy as np +from absl import app, flags +from lxml import etree +from dm_control import mjcf +from dm_control.locomotion.walkers.assets import dog_v2 from dm_control.utils import io as resources -flags.DEFINE_boolean('make_skin', True, 'Whether to make a new dog_skin.skn') +flags.DEFINE_boolean("make_skin", True, "Whether to make a new dog_skin.skn") +flags.DEFINE_boolean( + "lengthrange_from_joints", + False, + "Whether to compute the length ranges of spatial tendons from joints", +) +flags.DEFINE_boolean("use_muscles", False, + "Whether to use muscles or torque actuators") +flags.DEFINE_boolean( + "add_markers", False, "Whether to add markers; used for motion capture tracking" +) +flags.DEFINE_integer( + "muscle_strength_scale", + -1, + "muscle force multiplier, if negative we let MuJoCo computer the strength", +) +flags.DEFINE_string( + "muscle_dynamics", "Millard", "muscle dynamics one of [Millard, Sigmoid, General]" +) +flags.DEFINE_boolean( + "make_muscles_skin", False, "Whether to make new skins for muscles" +) flags.DEFINE_float( - 'lumbar_dofs_per_vertebra', + "lumbar_dofs_per_vertebra", 1.5, - 'Number of degrees of freedom per vertebra in lumbar spine.', + "Number of degrees of freedom per vertebra in lumbar spine.", ) flags.DEFINE_float( - 'cervical_dofs_per_vertebra', + "cervical_dofs_per_vertebra", 1.5, - 'Number of degrees of freedom vertebra in cervical spine.', + "Number of degrees of freedom vertebra in cervical spine.", ) flags.DEFINE_float( - 'caudal_dofs_per_vertebra', + "caudal_dofs_per_vertebra", 1, - 'Number of degrees of freedom vertebra in caudal spine.', + "Number of degrees of freedom vertebra in caudal spine.", +) +flags.DEFINE_boolean( + "composer", + False, + "Make slight adjustments to the model to make it compatible for the composer.", +) +flags.DEFINE_string( + "output_dir", + os.path.dirname(__file__), + "Output directory for the created dog model" ) FLAGS = flags.FLAGS -BASE_MODEL = 'dog_base.xml' -ASSET_RELPATH = '../../../../suite/dog_assets' -ASSET_DIR = os.path.dirname(__file__) + '/' + ASSET_RELPATH -print(ASSET_DIR) +BASE_MODEL = "dog_base.xml" +ASSET_RELPATH = "../../../../suite/dog_assets" +CURRENT_DIR = os.path.dirname(__file__) +ASSET_DIR = os.path.join(CURRENT_DIR, ASSET_RELPATH) def exclude_contacts(model): @@ -66,73 +90,71 @@ def exclude_contacts(model): physics = mjcf.Physics.from_mjcf_model(model) excluded_pairs = [] for c in physics.data.contact: - body1 = physics.model.id2name(physics.model.geom_bodyid[c.geom1], 'body') - body2 = physics.model.id2name(physics.model.geom_bodyid[c.geom2], 'body') - pair = body1 + ':' + body2 + body1 = physics.model.id2name(physics.model.geom_bodyid[c.geom1], "body") + body2 = physics.model.id2name(physics.model.geom_bodyid[c.geom2], "body") + pair = body1 + ":" + body2 if pair not in excluded_pairs: excluded_pairs.append(pair) - model.contact.add('exclude', name=pair, body1=body1, body2=body2) + model.contact.add("exclude", name=pair, body1=body1, body2=body2) # manual exclusions model.contact.add( - 'exclude', - name='C_1:jaw', - body1=model.find('body', 'C_1'), - body2=model.find('body', 'jaw'), + "exclude", + name="C_1:jaw", + body1=model.find("body", "C_1"), + body2=model.find("body", "jaw"), ) model.contact.add( - 'exclude', - name='torso:lower_arm_L', - body1=model.find('body', 'torso'), - body2='lower_arm_L', + "exclude", + name="torso:lower_arm_L", + body1=model.find("body", "torso"), + body2="lower_arm_L", ) model.contact.add( - 'exclude', - name='torso:lower_arm_R', - body1=model.find('body', 'torso'), - body2='lower_arm_R', + "exclude", + name="torso:lower_arm_R", + body1=model.find("body", "torso"), + body2="lower_arm_R", ) + model.contact.add("exclude", name="C_4:scapula_R", + body1="C_4", body2="scapula_R") + model.contact.add("exclude", name="C_4:scapula_L", + body1="C_4", body2="scapula_L") model.contact.add( - 'exclude', name='C_4:scapula_R', body1='C_4', body2='scapula_R' + "exclude", name="C_5:upper_arm_R", body1="C_5", body2="upper_arm_R" ) model.contact.add( - 'exclude', name='C_4:scapula_L', body1='C_4', body2='scapula_L' + "exclude", name="C_5:upper_arm_L", body1="C_5", body2="upper_arm_L" ) model.contact.add( - 'exclude', name='C_5:upper_arm_R', body1='C_5', body2='upper_arm_R' + "exclude", name="C_6:upper_arm_R", body1="C_6", body2="upper_arm_R" ) model.contact.add( - 'exclude', name='C_5:upper_arm_L', body1='C_5', body2='upper_arm_L' + "exclude", name="C_6:upper_arm_L", body1="C_6", body2="upper_arm_L" ) model.contact.add( - 'exclude', name='C_6:upper_arm_R', body1='C_6', body2='upper_arm_R' + "exclude", name="C_7:upper_arm_R", body1="C_7", body2="upper_arm_R" ) model.contact.add( - 'exclude', name='C_6:upper_arm_L', body1='C_6', body2='upper_arm_L' + "exclude", name="C_7:upper_arm_L", body1="C_7", body2="upper_arm_L" ) model.contact.add( - 'exclude', name='C_7:upper_arm_R', body1='C_7', body2='upper_arm_R' + "exclude", + name="upper_leg_L:upper_leg_R", + body1="upper_leg_L", + body2="upper_leg_R", ) - model.contact.add( - 'exclude', name='C_7:upper_arm_L', body1='C_7', body2='upper_arm_L' - ) - model.contact.add( - 'exclude', - name='upper_leg_L:upper_leg_R', - body1='upper_leg_L', - body2='upper_leg_R', - ) - for side in ['_L', '_R']: + for side in ["_L", "_R"]: model.contact.add( - 'exclude', - name='lower_leg' + side + ':pelvis', - body1='lower_leg' + side, - body2='pelvis', + "exclude", + name="lower_leg" + side + ":pelvis", + body1="lower_leg" + side, + body2="pelvis", ) model.contact.add( - 'exclude', - name='upper_leg' + side + ':foot' + side, - body1='upper_leg' + side, - body2='foot' + side, + "exclude", + name="upper_leg" + side + ":foot" + side, + body1="upper_leg" + side, + body2="foot" + side, ) @@ -145,56 +167,85 @@ def main(argv): cervical_dofs_per_vertebra = FLAGS.cervical_dofs_per_vertebra caudal_dofs_per_vertebra = FLAGS.caudal_dofs_per_vertebra make_skin = FLAGS.make_skin + use_muscles = FLAGS.use_muscles + muscle_strength_scale = FLAGS.muscle_strength_scale + make_muscles_skin = FLAGS.make_muscles_skin + muscle_dynamics = FLAGS.muscle_dynamics + add_markers = FLAGS.add_markers + lengthrange_from_joints = FLAGS.lengthrange_from_joints + composer = FLAGS.composer + output_dir = FLAGS.output_dir else: - lumbar_dofs_per_vert = FLAGS['lumbar_dofs_per_vertebra'].default - cervical_dofs_per_vertebra = FLAGS['cervical_dofs_per_vertebra'].default - caudal_dofs_per_vertebra = FLAGS['caudal_dofs_per_vertebra'].default - make_skin = FLAGS['make_skin'].default - - print('Load base model.') - with open(BASE_MODEL, 'r') as f: + lumbar_dofs_per_vert = FLAGS["lumbar_dofs_per_vertebra"].default + cervical_dofs_per_vertebra = FLAGS["cervical_dofs_per_vertebra"].default + caudal_dofs_per_vertebra = FLAGS["caudal_dofs_per_vertebra"].default + make_skin = FLAGS["make_skin"].default + use_muscles = FLAGS["use_muscles"].default + muscle_strength_scale = FLAGS["muscle_strength_scale"].default + make_muscles_skin = FLAGS["make_muscles_skin"].default + muscle_dynamics = FLAGS["muscle_dynamics"].default + add_markers = FLAGS["add_markers"].default + lengthrange_from_joints = FLAGS["lengthrange_from_joints"].default + composer = FLAGS["composer"].default + output_dir = FLAGS["output_dir"].default + + print("Load base model.") + with open(os.path.join(CURRENT_DIR, BASE_MODEL), "r") as f: model = mjcf.from_file(f) + if not composer: + floor = model.worldbody.add( + "body", name="floor", pos=(0, 0, 0) + ) + floor.add( + "geom", + name="floor", + type="plane", + conaffinity=1, + size=[10, 10, 0.1], + material="grid" + ) # Helper constants: side_sign = { - '_L': np.array((1.0, -1.0, 1.0)), - '_R': np.array((1.0, 1.0, 1.0)), + "_L": np.array((1.0, -1.0, 1.0)), + "_R": np.array((1.0, 1.0, 1.0)), } primary_axis = { - '_abduct': np.array((-1.0, 0.0, 0.0)), - '_extend': np.array((0.0, 1.0, 0.0)), - '_supinate': np.array((0.0, 0.0, -1.0)), + "_abduct": np.array((-1.0, 0.0, 0.0)), + "_extend": np.array((0.0, 1.0, 0.0)), + "_supinate": np.array((0.0, 0.0, -1.0)), } # Add meshes: - print('Loading all meshes, getting positions and sizes.') + print("Loading all meshes, getting positions and sizes.") meshdir = ASSET_DIR model.compiler.meshdir = meshdir texturedir = ASSET_DIR model.compiler.texturedir = texturedir bones = [] + skin_msh = None for dirpath, _, filenames in resources.WalkResources(meshdir): - prefix = 'extras/' if 'extras' in dirpath else '' + prefix = "extras/" if "extras" in dirpath else "" for filename in filenames: - if 'dog_skin.msh' in filename: + if "dog_skin.msh" in filename: skin_msh = model.asset.add( - 'mesh', name='skin_msh', file=filename, scale=(1.25, 1.25, 1.25) + "mesh", name="dog_skin", file=filename, scale=(1.25, 1.25, 1.25) ) name = filename[4:-4] - name = name.replace('*', ':') - if filename.startswith('BONE'): - if 'Lingual' not in name: + name = name.replace("*", ":") + if filename.startswith("BONE"): + if "Lingual" not in name: bones.append(name) - model.asset.add('mesh', name=name, file=prefix + filename) + model.asset.add("mesh", name=name, file=prefix + filename) # Put all bones in worldbody, get positions, remove bones: bone_geoms = [] for bone in bones: geom = model.worldbody.add( - 'geom', + "geom", name=bone, mesh=bone, - type='mesh', + type="mesh", contype=0, conaffinity=0, rgba=[1, 0.5, 0.5, 0.4], @@ -204,36 +255,37 @@ def main(argv): bone_position = {} bone_size = {} for bone in bones: - geom = model.find('geom', bone) + geom = model.find("geom", bone) bone_position[bone] = np.array(physics.bind(geom).xpos) bone_size[bone] = np.array(physics.bind(geom).rbound) geom.remove() # Torso - print('Torso, lumbar spine, pelvis.') - pelvic_bones, lumbar_joints = build_torso.create_torso( + print("Torso, lumbar spine, pelvis.") + pelvic_bones, lumbar_joints = dog_v2.build_torso( model, bones, bone_position, lumbar_dofs_per_vert, side_sign, parent=model.worldbody, + composer=composer, ) - print('Neck, skull, jaw.') + print("Neck, skull, jaw.") # Cervical spine (neck) bodies: - cervical_joints = build_neck.create_neck( + cervical_joints = dog_v2.build_neck( model, bone_position, cervical_dofs_per_vertebra, bones, side_sign, bone_size, - parent=model.find('body', 'torso'), + parent=model.find("body", "torso"), ) - print('Back legs.') - nails, sole_sites = build_back_legs.create_back_legs( + print("Back legs.") + nails, sole_sites = dog_v2.build_back_legs( model, primary_axis, bone_position, @@ -241,53 +293,51 @@ def main(argv): side_sign, bone_size, pelvic_bones, - parent=model.find('body', 'pelvis'), + parent=model.find("body", "pelvis"), ) - print('Shoulders, front legs.') - palm_sites = build_front_legs.create_front_legs( + print("Shoulders, front legs.") + palm_sites = dog_v2.build_front_legs( nails, model, primary_axis, bones, side_sign, - parent=model.find('body', 'torso'), + parent=model.find("body", "torso"), ) - print('Tail.') - caudal_joints = build_tail.create_tail( + print("Tail.") + caudal_joints = dog_v2.build_tail( caudal_dofs_per_vertebra, bone_size, model, bone_position, - parent=model.find('body', 'pelvis'), + parent=model.find("body", "pelvis"), ) - print('Collision geoms, fixed tendons.') + print("Collision geoms, fixed tendons.") physics = mjcf.Physics.from_mjcf_model(model) - print('Unify ribcage and jaw meshes.') - for body in model.find_all('body'): + print("Unify ribcage and jaw meshes.") + for body in model.find_all("body"): body_meshes = [ geom for geom in body.all_children() - if geom.tag == 'geom' - and hasattr(geom, 'mesh') - and geom.mesh is not None + if geom.tag == "geom" and hasattr(geom, "mesh") and geom.mesh is not None ] if len(body_meshes) > 10: mergables = [ - ('torso', 'Ribcage'), - ('jaw', 'Jaw'), - ('skull', 'MergedSkull'), + ("torso", "Ribcage"), + ("jaw", "Jaw"), + ("skull", "MergedSkull"), ] for bodyname, meshname in mergables: if body.name == bodyname: - print('==== Merging ', bodyname) + print("==== Merging ", bodyname) for mesh in body_meshes: print(mesh.name) body.add( - 'inertial', + "inertial", mass=physics.bind(body).mass, pos=physics.bind(body).ipos, quat=physics.bind(body).iquat, @@ -295,41 +345,102 @@ def main(argv): ) for mesh in body_meshes: - if 'eye' not in mesh.name: - model.find('mesh', mesh.name).remove() + if "eye" not in mesh.name: + model.find("mesh", mesh.name).remove() mesh.remove() body.add( - 'geom', + "geom", name=meshname, mesh=meshname, - dclass='bone', + dclass="bone", pos=-bone_position[meshname], ) - print('Add Actuators') - actuated_joints = add_torque_actuators.add_motors( - physics, model, lumbar_joints, cervical_joints, caudal_joints - ) - - print('Excluding contacts.') + print("Excluding contacts.") exclude_contacts(model) - if make_skin: - create_skin.create(model, skin_msh) + if make_skin and skin_msh is not None: + dog_v2.create_skin(model=model, mesh_file=skin_msh, + asset_dir=ASSET_DIR, mesh_name="dog_skin", composer=composer) - # Add skin from .skn - print('Adding Skin.') - skin_texture = model.asset.add( - 'texture', name='skin', file='skin_texture.png', type='2d' - ) - skin_material = model.asset.add('material', name='skin', texture=skin_texture) - model.asset.add( - 'skin', name='skin', file='dog_skin.skn', material=skin_material - ) - skin_msh.remove() + # Add skin from .skn + print("Adding Skin.") + skin_texture = model.asset.add( + "texture", name="skin", file="skin_texture.png", type="2d" + ) + skin_material = model.asset.add( + "material", name="skin", texture=skin_texture) + model.asset.add( + "skin", + name="skin", + file="./skins/dog_skin.skn", + material=skin_material, + group=0, + ) + skin_msh.remove() + + if make_muscles_skin: + muscle_material = model.asset.add( + "material", name="muscle", rgba=[0.63, 0.17, 0.17, 1] + ) + for dirpath, _, filenames in resources.WalkResources(ASSET_DIR + "/muscles/"): + for filename in filenames: + if filename[-4:] not in [".obj", ".stl"]: + continue + print("filename", filename) + muscle_msh = model.asset.add( + "mesh", name=filename[:-4], file=ASSET_DIR + "/muscles/" + filename + ) + dog_v2.create_skin( + model=model, + asset_dir=ASSET_DIR, + mesh_file=muscle_msh, + mesh_name=filename[:-4], + tex_coords=False, + transform=False, + composer=composer, + ) + + # Add skin from .skn + print("Adding Skin --> " + filename[:-4]) + file_path = "./skins/" + filename[:-4] + ".skn" + model.asset.add( + "skin", + name=filename[:-4], + file=file_path, + material=muscle_material, + group=1, + ) + muscle_msh.remove() + + model.option.timestep = 0.005 + model.option.integrator = "implicitfast" + model.option.noslip_iterations = 3 + model.option.cone = "elliptic" + + if add_markers: + print("Add Markers") + dog_v2.add_markers(model) + + print("Add Actuators") + if use_muscles: + # Add wrapping geometries + dog_v2.add_wrapping_geoms(model) + # Add muscles + dog_v2.add_muscles( + model, + muscle_strength_scale, + muscle_dynamics, + ASSET_DIR, + lengthrange_from_joints, + ) + else: + actuated_joints = dog_v2.add_torque_actuators( + physics, model, lumbar_joints, cervical_joints, caudal_joints + ) - print('Removing non-essential sites.') - all_sites = model.find_all('site') + print("Removing non-essential sites.") + all_sites = model.find_all("site") for site in all_sites: if site.dclass is None: site.remove() @@ -337,90 +448,106 @@ def main(argv): physics = mjcf.Physics.from_mjcf_model(model) # sensors model.sensor.add( - 'accelerometer', name='accelerometer', site=model.find('site', 'head') + "accelerometer", name="accelerometer", site=model.find("site", "head") ) + model.sensor.add("velocimeter", name="velocimeter", + site=model.find("site", "head")) + model.sensor.add("gyro", name="gyro", site=model.find("site", "head")) model.sensor.add( - 'velocimeter', name='velocimeter', site=model.find('site', 'head') + "subtreelinvel", name="torso_linvel", body=model.find("body", "torso") ) - model.sensor.add('gyro', name='gyro', site=model.find('site', 'head')) model.sensor.add( - 'subtreelinvel', name='torso_linvel', body=model.find('body', 'torso') - ) - model.sensor.add( - 'subtreeangmom', name='torso_angmom', body=model.find('body', 'torso') + "subtreeangmom", name="torso_angmom", body=model.find("body", "torso") ) for site in palm_sites + sole_sites: - model.sensor.add('touch', name=site.name, site=site) - anchors = [site for site in model.find_all('site') if 'anchor' in site.name] + model.sensor.add("touch", name=site.name, site=site) + anchors = [site for site in model.find_all("site") if "anchor" in site.name] for site in anchors: - model.sensor.add('force', name=site.name.replace('_anchor', ''), site=site) + model.sensor.add("force", name=site.name.replace("_anchor", ""), site=site) - # Print stuff - joint_acts = [model.find('actuator', j.name) for j in actuated_joints] - print( - '{:20} {:>10} {:>10} {:>10} {:>10} {:>10}'.format( - 'name', 'mass', 'damping', 'stiffness', 'ratio', 'armature' - ) - ) - for i, j in enumerate(actuated_joints): - dmp = physics.bind(j).damping[0] - mass_eff = physics.bind(j).M0[0] - dmp = physics.bind(j).damping[0] - stf = physics.bind(joint_acts[i]).gainprm[0] - arma = physics.bind(j).armature[0] + # Print torque actuators stuff + if not use_muscles: + joint_acts = [model.find("actuator", j.name) for j in actuated_joints] print( - '{:20} {:10.4} {:10} {:10.4} {:10.4} {:10}'.format( - j.name, - mass_eff, - dmp, - stf, - dmp / (2 * np.sqrt(mass_eff * stf)), - arma, + "{:20} {:>10} {:>10} {:>10} {:>10} {:>10}".format( + "name", "mass", "damping", "stiffness", "ratio", "armature" ) ) + for i, j in enumerate(actuated_joints): + mass_eff = physics.bind(j).M0[0] + dmp = physics.bind(j).damping[0] + stf = physics.bind(joint_acts[i]).gainprm[0] + arma = physics.bind(j).armature[0] + print( + "{:20} {:10.4} {:10} {:10.4} {:10.4} {:10}".format( + j.name, + mass_eff, + dmp, + stf, + dmp / (2 * np.sqrt(mass_eff * stf)), + arma, + ) + ) + + if composer: + torso = model.find("body", "torso") + torso.pos = np.zeros(3) - print('Finalising and saving model.') - xml_string = model.to_xml_string('float', precision=4, zero_threshold=1e-7) + print("Finalising and saving model.") + xml_string = model.to_xml_string("float", precision=4, zero_threshold=1e-7) root = etree.XML(xml_string, etree.XMLParser(remove_blank_text=True)) - print('Remove hashes from filenames') - assets = list(root.find('asset').iter()) + print("Remove hashes from filenames") + assets = list(root.find("asset").iter()) for asset in assets: - asset_filename = asset.get('file') + asset_filename = asset.get("file") if asset_filename is not None: name = asset_filename[:-4] extension = asset_filename[-4:] - asset.set('file', name[:-41] + extension) + if asset_filename[-3:] == "skn": + asset.set("file", "skins/" + name[:-41] + extension) + else: + asset.set("file", name[:-41] + extension) - print('Add , for locally-loadable model') + print("Add , for locally-loadable model") compiler = etree.Element( - 'compiler', meshdir=ASSET_RELPATH, texturedir=ASSET_RELPATH + "compiler", meshdir=ASSET_RELPATH, texturedir=ASSET_RELPATH ) root.insert(0, compiler) print('Remove class="/"') - default_elem = root.find('default') + default_elem = root.find("default") root.insert(6, default_elem[0]) root.remove(default_elem) xml_string = etree.tostring(root, pretty_print=True) - xml_string = xml_string.replace(b' class="/"', b'') + xml_string = xml_string.replace(b' class="/"', b"") - print('Insert spaces between top level elements') + print("Insert spaces between top level elements") lines = xml_string.splitlines() newlines = [] for line in lines: newlines.append(line) - if line.startswith(b' <'): - if line.startswith(b' '): - newlines.append(b'') - newlines.append(b'') - xml_string = b'\n'.join(newlines) + if line.startswith(b" <"): + if line.startswith(b" "): + newlines.append(b"") + newlines.append(b"") + xml_string = b"\n".join(newlines) + name_prefix = "composer_" if composer else "" # Save to file. - f = open('dog.xml', 'wb') + if not use_muscles: + name = name_prefix + "dog.xml" + elif muscle_strength_scale < 0: + name = name_prefix + "dog_muscles_{}.xml".format(muscle_dynamics) + else: + name = name_prefix + "dog_muscles_{}_{}.xml".format( + muscle_strength_scale, muscle_dynamics + ) + + f = open(os.path.join(output_dir, name), "wb") f.write(xml_string) f.close() -if __name__ == '__main__': +if __name__ == "__main__": app.run(main) diff --git a/dm_control/locomotion/walkers/assets/dog_v2/build_front_legs.py b/dm_control/locomotion/walkers/assets/dog_v2/build_front_legs.py index e059da07..88ec1e83 100644 --- a/dm_control/locomotion/walkers/assets/dog_v2/build_front_legs.py +++ b/dm_control/locomotion/walkers/assets/dog_v2/build_front_legs.py @@ -15,9 +15,10 @@ """Make front legs for the dog model.""" -from dm_control import mjcf import numpy as np +from dm_control import mjcf + def create_front_legs(nails, model, primary_axis, bones, side_sign, parent): """Add front legs in the model. @@ -34,45 +35,45 @@ def create_front_legs(nails, model, primary_axis, bones, side_sign, parent): Returns: A list of palm sites. """ - def_scapula_supinate = model.default.find('default', 'scapula_supinate') - def_scapula_abduct = model.default.find('default', 'scapula_abduct') - def_scapula_extend = model.default.find('default', 'scapula_extend') + def_scapula_supinate = model.default.find("default", "scapula_supinate") + def_scapula_abduct = model.default.find("default", "scapula_abduct") + def_scapula_extend = model.default.find("default", "scapula_extend") scapula_defaults = { - '_abduct': def_scapula_abduct, - '_extend': def_scapula_extend, - '_supinate': def_scapula_supinate, + "_abduct": def_scapula_abduct, + "_extend": def_scapula_extend, + "_supinate": def_scapula_supinate, } torso = parent # Shoulders scapula = {} - scapulae = [b for b in bones if 'Scapula' in b] + scapulae = [b for b in bones if "Scapula" in b] scapula_pos = np.array((0.08, -0.02, 0.14)) - for side in ['_L', '_R']: + for side in ["_L", "_R"]: body_pos = scapula_pos * side_sign[side] - arm = torso.add('body', name='scapula' + side, pos=body_pos) + arm = torso.add("body", name="scapula" + side, pos=body_pos) scapula[side] = arm for bone in [b for b in scapulae if side in b]: arm.add( - 'geom', name=bone, mesh=bone, pos=-torso.pos - body_pos, dclass='bone' + "geom", name=bone, mesh=bone, pos=-torso.pos - body_pos, dclass="bone" ) # Shoulder joints - for dof in ['_supinate', '_abduct', '_extend']: + for dof in ["_supinate", "_abduct", "_extend"]: joint_axis = primary_axis[dof].copy() joint_pos = scapula_defaults[dof].joint.pos.copy() - if dof != '_extend': - joint_axis *= 1.0 if side == '_R' else -1.0 + if dof != "_extend": + joint_axis *= 1.0 if side == "_R" else -1.0 joint_pos *= side_sign[side] else: joint_axis += ( - 0.3 * (1 if side == '_R' else -1) * primary_axis['_abduct'] + 0.3 * (1 if side == "_R" else -1) * primary_axis["_abduct"] ) arm.add( - 'joint', - name='scapula' + side + dof, - dclass='scapula' + dof, + "joint", + name="scapula" + side + dof, + dclass="scapula" + dof, axis=joint_axis, pos=joint_pos, ) @@ -80,8 +81,8 @@ def create_front_legs(nails, model, primary_axis, bones, side_sign, parent): # Shoulder sites shoulder_pos = np.array((0.075, -0.033, -0.13)) arm.add( - 'site', - name='shoulder' + side, + "site", + name="shoulder" + side, size=[0.01], pos=shoulder_pos * side_sign[side], ) @@ -89,77 +90,75 @@ def create_front_legs(nails, model, primary_axis, bones, side_sign, parent): # Upper Arms: upper_arm = {} parent_pos = {} - humeri = ['humerus_R', 'humerus_L'] - for side in ['_L', '_R']: + humeri = ["humerus_R", "humerus_L"] + for side in ["_L", "_R"]: body_pos = shoulder_pos * side_sign[side] parent = scapula[side] parent_pos[side] = torso.pos + parent.pos - arm = parent.add('body', name='upper_arm' + side, pos=body_pos) + arm = parent.add("body", name="upper_arm" + side, pos=body_pos) upper_arm[side] = arm for bone in [b for b in humeri if side in b]: arm.add( - 'geom', + "geom", name=bone, mesh=bone, pos=-parent_pos[side] - body_pos, - dclass='bone', + dclass="bone", ) parent_pos[side] += body_pos # Shoulder joints - for dof in ['_supinate', '_extend']: + for dof in ["_supinate", "_extend"]: joint_axis = primary_axis[dof].copy() - if dof == '_supinate': + if dof == "_supinate": joint_axis[0] = 1 - joint_axis *= 1.0 if side == '_R' else -1.0 + joint_axis *= 1.0 if side == "_R" else -1.0 arm.add( - 'joint', - name='shoulder' + side + dof, - dclass='shoulder' + dof, + "joint", + name="shoulder" + side + dof, + dclass="shoulder" + dof, axis=joint_axis, ) # Elbow sites elbow_pos = np.array((-0.05, -0.015, -0.145)) arm.add( - 'site', - type='cylinder', - name='elbow' + side, + "site", + type="cylinder", + name="elbow" + side, size=[0.003, 0.02], - zaxis=(0, 1, -(1.0 if side == '_R' else -1.0) * 0.2), + zaxis=(0, 1, -(1.0 if side == "_R" else -1.0) * 0.2), pos=elbow_pos * side_sign[side], ) # Lower arms: lower_arm = {} lower_arm_bones = [ - b - for b in bones - if 'ulna' in b.lower() or 'Radius' in b or 'accessory' in b + b for b in bones if "ulna" in b.lower() or "Radius" in b or "accessory" in b ] - for side in ['_L', '_R']: + for side in ["_L", "_R"]: body_pos = elbow_pos * side_sign[side] - arm = upper_arm[side].add('body', name='lower_arm' + side, pos=body_pos) + arm = upper_arm[side].add("body", name="lower_arm" + side, pos=body_pos) lower_arm[side] = arm for bone in [b for b in lower_arm_bones if side in b]: arm.add( - 'geom', + "geom", name=bone, mesh=bone, pos=-parent_pos[side] - body_pos, - dclass='bone', + dclass="bone", ) # Elbow joints - elbow_axis = upper_arm[side].find_all('site')[0].zaxis - arm.add('joint', name='elbow' + side, dclass='elbow', axis=elbow_axis) + elbow_axis = upper_arm[side].find_all("site")[0].zaxis + arm.add("joint", name="elbow" + side, dclass="elbow", axis=elbow_axis) parent_pos[side] += body_pos # Wrist sites wrist_pos = np.array((0.003, 0.015, -0.19)) arm.add( - 'site', - type='cylinder', - name='wrist' + side, + "site", + type="cylinder", + name="wrist" + side, size=[0.004, 0.017], zaxis=(0, 1, 0), pos=wrist_pos * side_sign[side], @@ -170,44 +169,44 @@ def create_front_legs(nails, model, primary_axis, bones, side_sign, parent): hand_bones = [ b for b in bones - if ('carpal' in b.lower() and 'acces' not in b and 'ulna' not in b) - or ('distalis_digiti_I_' in b) + if ("carpal" in b.lower() and "acces" not in b and "ulna" not in b) + or ("distalis_digiti_I_" in b) ] - for side in ['_L', '_R']: + for side in ["_L", "_R"]: body_pos = wrist_pos * side_sign[side] hand_anchor = lower_arm[side].add( - 'body', name='hand_anchor' + side, pos=body_pos + "body", name="hand_anchor" + side, pos=body_pos ) hand_anchor.add( - 'geom', + "geom", name=hand_anchor.name, - dclass='foot_primitive', - type='box', + dclass="foot_primitive", + type="box", size=(0.005, 0.005, 0.005), contype=0, conaffinity=0, ) - hand_anchor.add('site', name=hand_anchor.name, dclass='sensor') - hand = hand_anchor.add('body', name='hand' + side) + hand_anchor.add("site", name=hand_anchor.name, dclass="sensor") + hand = hand_anchor.add("body", name="hand" + side) hands[side] = hand for bone in [b for b in hand_bones if side in b]: hand.add( - 'geom', + "geom", name=bone, mesh=bone, pos=-parent_pos[side] - body_pos, - dclass='bone', + dclass="bone", ) # Wrist joints - hand.add('joint', name='wrist' + side, dclass='wrist', axis=(0, -1, 0)) + hand.add("joint", name="wrist" + side, dclass="wrist", axis=(0, -1, 0)) hand.add( - 'geom', - name=hand.name + '_collision', + "geom", + name=hand.name + "_collision", size=[0.03, 0.016, 0.012], pos=[0.01, 0, -0.04], euler=(0, 65, 0), - dclass='collision_primitive', - type='box', + dclass="collision_primitive", + type="box", ) parent_pos[side] += body_pos @@ -215,122 +214,122 @@ def create_front_legs(nails, model, primary_axis, bones, side_sign, parent): # Finger sites finger_pos = np.array((0.02, 0, -0.06)) hand.add( - 'site', - type='cylinder', - name='finger' + side, + "site", + type="cylinder", + name="finger" + side, size=[0.003, 0.025], - zaxis=((1.0 if side == '_R' else -1.0) * 0.2, 1, 0), + zaxis=((1.0 if side == "_R" else -1.0) * 0.2, 1, 0), pos=finger_pos * side_sign[side], ) # Fingers: finger_bones = [ - b for b in bones if 'Phalanx' in b and 'distalis_digiti_I_' not in b + b for b in bones if "Phalanx" in b and "distalis_digiti_I_" not in b ] palm_sites = [] - for side in ['_L', '_R']: + for side in ["_L", "_R"]: body_pos = finger_pos * side_sign[side] - hand = hands[side].add('body', name='finger' + side, pos=body_pos) + hand = hands[side].add("body", name="finger" + side, pos=body_pos) for bone in [b for b in finger_bones if side in b]: geom = hand.add( - 'geom', + "geom", name=bone, mesh=bone, pos=-parent_pos[side] - body_pos, - dclass='bone', + dclass="bone", ) - if 'distalis' in bone: + if "distalis" in bone: nails.append(bone) - geom.dclass = 'visible_bone' + geom.dclass = "visible_bone" # Finger joints - finger_axis = upper_arm[side].find_all('site')[0].zaxis - hand.add('joint', name='finger' + side, dclass='finger', axis=finger_axis) + finger_axis = upper_arm[side].find_all("site")[0].zaxis + hand.add("joint", name="finger" + side, dclass="finger", axis=finger_axis) hand.add( - 'geom', - name=hand.name + '0_collision', + "geom", + name=hand.name + "0_collision", size=[0.018, 0.012], pos=[0.012, 0, -0.012], euler=(90, 0, 0), - dclass='foot_primitive', + dclass="foot_primitive", ) hand.add( - 'geom', - name=hand.name + '1_collision', + "geom", + name=hand.name + "1_collision", size=[0.01, 0.015], pos=[0.032, 0, -0.02], euler=(90, 0, 0), - dclass='foot_primitive', + dclass="foot_primitive", ) hand.add( - 'geom', - name=hand.name + '2_collision', + "geom", + name=hand.name + "2_collision", size=[0.008, 0.01], pos=[0.042, 0, -0.022], euler=(90, 0, 0), - dclass='foot_primitive', + dclass="foot_primitive", ) palm = hand.add( - 'site', - name='palm' + side, + "site", + name="palm" + side, size=(0.028, 0.03, 0.007), pos=(0.02, 0, -0.024), - type='box', - dclass='sensor', + type="box", + dclass="sensor", ) palm_sites.append(palm) physics = mjcf.Physics.from_mjcf_model(model) - for side in ['_L', '_R']: + for side in ["_L", "_R"]: # scapula: scap = scapula[side] - geom = scap.get_children('geom')[0] + geom = scap.get_children("geom")[0] bound_geom = physics.bind(geom) scap.add( - 'geom', - name=geom.name + '_collision', + "geom", + name=geom.name + "_collision", pos=bound_geom.pos, size=bound_geom.size * 0.8, quat=bound_geom.quat, - type='box', - dclass='collision_primitive', + type="box", + dclass="collision_primitive", ) # upper arm: arm = upper_arm[side] - geom = arm.get_children('geom')[0] + geom = arm.get_children("geom")[0] bound_geom = physics.bind(geom) arm.add( - 'geom', - name=geom.name + '_collision', + "geom", + name=geom.name + "_collision", pos=bound_geom.pos, size=[0.02, 0.08], quat=bound_geom.quat, - dclass='collision_primitive', + dclass="collision_primitive", ) - all_geoms = model.find_all('geom') + all_geoms = model.find_all("geom") for geom in all_geoms: - if 'Ulna' in geom.name: + if "Ulna" in geom.name: bound_geom = physics.bind(geom) geom.parent.add( - 'geom', - name=geom.name + '_collision', + "geom", + name=geom.name + "_collision", pos=bound_geom.pos, size=[0.015, 0.06], quat=bound_geom.quat, - dclass='collision_primitive', + dclass="collision_primitive", ) - if 'Radius' in geom.name: + if "Radius" in geom.name: bound_geom = physics.bind(geom) pos = bound_geom.pos + np.array((-0.005, 0.0, -0.01)) geom.parent.add( - 'geom', - name=geom.name + '_collision', + "geom", + name=geom.name + "_collision", pos=pos, size=[0.017, 0.09], quat=bound_geom.quat, - dclass='collision_primitive', + dclass="collision_primitive", ) return palm_sites diff --git a/dm_control/locomotion/walkers/assets/dog_v2/build_neck.py b/dm_control/locomotion/walkers/assets/dog_v2/build_neck.py index 8435796a..e3d4e2d6 100644 --- a/dm_control/locomotion/walkers/assets/dog_v2/build_neck.py +++ b/dm_control/locomotion/walkers/assets/dog_v2/build_neck.py @@ -17,9 +17,10 @@ import collections -from dm_control import mjcf import numpy as np +from dm_control import mjcf + def create_neck( model, @@ -47,17 +48,17 @@ def create_neck( A list of cervical joints. """ # Cervical Spine - def_cervical = model.default.find('default', 'cervical') - def_cervical_extend = model.default.find('default', 'cervical_extend') - def_cervical_bend = model.default.find('default', 'cervical_bend') - def_cervical_twist = model.default.find('default', 'cervical_twist') + def_cervical = model.default.find("default", "cervical") + def_cervical_extend = model.default.find("default", "cervical_extend") + def_cervical_bend = model.default.find("default", "cervical_bend") + def_cervical_twist = model.default.find("default", "cervical_twist") cervical_defaults = { - 'extend': def_cervical_extend, - 'bend': def_cervical_bend, - 'twist': def_cervical_twist, + "extend": def_cervical_extend, + "bend": def_cervical_bend, + "twist": def_cervical_twist, } - cervical_bones = ['C_' + str(i) for i in range(7, 0, -1)] + cervical_bones = ["C_" + str(i) for i in range(7, 0, -1)] parent_pos = parent.pos cervical_bodies = [] cervical_geoms = [] @@ -65,16 +66,17 @@ def create_neck( for i, bone in enumerate(cervical_bones): bone_pos = bone_position[bone] rel_pos = bone_pos - parent_pos - child = parent.add('body', name=bone, pos=rel_pos) + child = parent.add("body", name=bone, pos=rel_pos) cervical_bodies.append(child) - dclass = 'bone' if i > 3 else 'light_bone' - geom = child.add('geom', name=bone, mesh=bone, pos=-bone_pos, dclass=dclass) + dclass = "bone" if i > 3 else "light_bone" + geom = child.add("geom", name=bone, mesh=bone, + pos=-bone_pos, dclass=dclass) child.add( - 'geom', - name=bone + '_collision', - type='sphere', + "geom", + name=bone + "_collision", + type="sphere", size=[radius], - dclass='nonself_collision_primitive', + dclass="nonself_collision_primitive", ) radius -= 0.006 cervical_geoms.append(geom) @@ -86,35 +88,60 @@ def create_neck( # Cervical (neck) spine joints: cervical_axis = collections.OrderedDict() - cervical_axis['extend'] = np.array((0.0, 1.0, 0.0)) - cervical_axis['bend'] = np.array((0.0, 0.0, 1.0)) - cervical_axis['twist'] = np.array((1.0, 0.0, 0)) + cervical_axis["extend"] = np.array((0.0, 1.0, 0.0)) + cervical_axis["bend"] = np.array((0.0, 0.0, 1.0)) + cervical_axis["twist"] = np.array((1.0, 0.0, 0)) + + neck_ranges = { + "C_7_extend": [-14.2, 35.4], + "C_7_bend": [-23.5, 23.5], + "C_7_twist": [-42.4, 42.4], + "C_6_extend": [-14.2, 35.4], + "C_6_bend": [-23.5, 23.5], + "C_6_twist": [-42.4, 42.4], + "C_5_extend": [-11.9, 11.2], + "C_5_bend": [-27.9, 27.9], + "C_5_twist": [-33, 33], + "C_4_extend": [-11.9, 11.2], + "C_4_bend": [-27.9, 27.9], + "C_4_twist": [-33, 33], + "C_3_extend": [-13.4, 18.3], + "C_3_bend": [-32, 32], + "C_3_twist": [-10.1, 10.1], + "C_2_extend": [-13.4, 18.3], + "C_2_bend": [-32, 32], + "C_2_twist": [-10.1, 10.1], + "C_1_extend": [-16.4, 12.9], + "C_1_bend": [-26, 26], + "C_1_twist": [-10.1, 10.1], + } num_dofs = 0 cervical_joints = [] cervical_joint_names = [] - torso = model.find('body', 'torso') - parent = torso.find('geom', 'T_1') + torso = model.find("body", "torso") + parent = torso.find("geom", "T_1") for i, vertebra in enumerate(cervical_bodies): while num_dofs < (i + 1) * cervical_dofs_per_vertebra: dof = num_dofs % 3 dof_name = list(cervical_axis.keys())[dof] - cervical_joint_names.append(vertebra.name + '_' + dof_name) + cervical_joint_names.append(vertebra.name + "_" + dof_name) rel_pos = physics.bind(vertebra).xpos - physics.bind(parent).xpos twist_dir = rel_pos / np.linalg.norm(rel_pos) - bend_dir = np.cross(twist_dir, cervical_axis['extend']) - cervical_axis['bend'] = bend_dir - cervical_axis['twist'] = twist_dir - joint_frame = np.vstack((twist_dir, cervical_axis['extend'], bend_dir)) + bend_dir = np.cross(twist_dir, cervical_axis["extend"]) + cervical_axis["bend"] = bend_dir + cervical_axis["twist"] = twist_dir + joint_frame = np.vstack((twist_dir, cervical_axis["extend"], bend_dir)) joint_pos = ( def_cervical.joint.pos - * physics.bind(vertebra.find('geom', vertebra.name)).size.mean() + * physics.bind(vertebra.find("geom", vertebra.name)).size.mean() ) joint = vertebra.add( - 'joint', + "joint", name=cervical_joint_names[-1], - dclass='cervical_' + dof_name, + dclass="cervical_" + dof_name, + range=neck_ranges[cervical_joint_names[-1]], axis=cervical_axis[dof_name], pos=joint_pos.dot(joint_frame), ) @@ -124,15 +151,14 @@ def create_neck( # Lumbar spine joints: lumbar_axis = collections.OrderedDict() - lumbar_axis['extend'] = np.array((0.0, 1.0, 0.0)) - lumbar_axis['bend'] = np.array((0.0, 0.0, 1.0)) - lumbar_axis['twist'] = np.array((1.0, 0.0, 0)) + lumbar_axis["extend"] = np.array((0.0, 1.0, 0.0)) + lumbar_axis["bend"] = np.array((0.0, 0.0, 1.0)) + lumbar_axis["twist"] = np.array((1.0, 0.0, 0)) # Scale joint defaults relative to 3 cervical_dofs_per_vertebra for dof in lumbar_axis.keys(): - axis_scale = 7.0 / [dof in joint for joint in cervical_joint_names].count( - True - ) + axis_scale = 7.0 / \ + [dof in joint for joint in cervical_joint_names].count(True) cervical_defaults[dof].joint.range *= axis_scale # Reload @@ -140,45 +166,45 @@ def create_neck( # Skull c_1 = cervical_bodies[-1] - upper_teeth = [m for m in bones if 'Top' in m] - skull_bones = upper_teeth + ['Skull', 'Ethmoid', 'Vomer', 'eye_L', 'eye_R'] + upper_teeth = [m for m in bones if "Top" in m] + skull_bones = upper_teeth + ["Skull", "Ethmoid", "Vomer", "eye_L", "eye_R"] skull = c_1.add( - 'body', name='skull', pos=bone_position['Skull'] - physics.bind(c_1).xpos + "body", name="skull", pos=bone_position["Skull"] - physics.bind(c_1).xpos ) skull_geoms = [] for bone in skull_bones: geom = skull.add( - 'geom', + "geom", name=bone, mesh=bone, - pos=-bone_position['Skull'], - dclass='light_bone', + pos=-bone_position["Skull"], + dclass="light_bone", ) - if 'eye' in bone: + if "eye" in bone: geom.rgba = [1, 1, 1, 1] - geom.dclass = 'visible_bone' + geom.dclass = "visible_bone" skull_geoms.append(geom) if bone in upper_teeth: - geom.dclass = 'visible_bone' + geom.dclass = "visible_bone" - for side in ['_L', '_R']: + for side in ["_L", "_R"]: pos = np.array((0.023, -0.027, 0.01)) * side_sign[side] skull.add( - 'geom', - name='iris' + side, - type='ellipsoid', - dclass='visible_bone', + "geom", + name="iris" + side, + type="ellipsoid", + dclass="visible_bone", rgba=(0.45, 0.45, 0.225, 0.4), size=(0.003, 0.007, 0.007), pos=pos, - euler=[0, 0, -20 * (1.0 if side == '_R' else -1.0)], + euler=[0, 0, -20 * (1.0 if side == "_R" else -1.0)], ) pos = np.array((0.0215, -0.0275, 0.01)) * side_sign[side] skull.add( - 'geom', - name='pupil' + side, - type='sphere', - dclass='visible_bone', + "geom", + name="pupil" + side, + type="sphere", + dclass="visible_bone", rgba=(0, 0, 0, 1), size=(0.003, 0, 0), pos=pos, @@ -186,83 +212,82 @@ def create_neck( # collision geoms skull.add( - 'geom', - name='skull0' + '_collision', - type='ellipsoid', - dclass='collision_primitive', + "geom", + name="skull0" + "_collision", + type="ellipsoid", + dclass="collision_primitive", size=(0.06, 0.06, 0.04), pos=(-0.02, 0, 0.01), euler=[0, 10, 0], ) skull.add( - 'geom', - name='skull1' + '_collision', - type='capsule', - dclass='collision_primitive', + "geom", + name="skull1" + "_collision", + type="capsule", + dclass="collision_primitive", size=(0.015, 0.04, 0.015), pos=(0.06, 0, -0.01), euler=[0, 110, 0], ) skull.add( - 'geom', - name='skull2' + '_collision', - type='box', - dclass='collision_primitive', + "geom", + name="skull2" + "_collision", + type="box", + dclass="collision_primitive", size=(0.03, 0.028, 0.008), pos=(0.02, 0, -0.03), ) skull.add( - 'geom', - name='skull3' + '_collision', - type='box', - dclass='collision_primitive', + "geom", + name="skull3" + "_collision", + type="box", + dclass="collision_primitive", size=(0.02, 0.018, 0.006), pos=(0.07, 0, -0.03), ) skull.add( - 'geom', - name='skull4' + '_collision', - type='box', - dclass='collision_primitive', + "geom", + name="skull4" + "_collision", + type="box", + dclass="collision_primitive", size=(0.005, 0.015, 0.004), pos=(0.095, 0, -0.03), ) skull.add( - 'joint', - name='atlas', - dclass='atlas', - pos=np.array((-0.5, 0, 0)) * bone_size['Skull'], + "joint", + name="atlas", + dclass="atlas", + pos=np.array((-0.5, 0, 0)) * bone_size["Skull"], ) + skull.add("site", name="head", size=( + 0.01, 0.01, 0.01), type="box", dclass="sensor") skull.add( - 'site', name='head', size=(0.01, 0.01, 0.01), type='box', dclass='sensor' - ) - skull.add( - 'site', - name='upper_bite', + "site", + name="upper_bite", size=(0.005,), - dclass='sensor', + dclass="sensor", pos=(0.065, 0, -0.07), ) # Jaw - lower_teeth = [m for m in bones if 'Bottom' in m] - jaw_bones = lower_teeth + ['Mandible'] + lower_teeth = [m for m in bones if "Bottom" in m] + jaw_bones = lower_teeth + ["Mandible"] jaw = skull.add( - 'body', name='jaw', pos=bone_position['Mandible'] - bone_position['Skull'] + "body", name="jaw", pos=bone_position["Mandible"] - bone_position["Skull"] ) jaw_geoms = [] for bone in jaw_bones: geom = jaw.add( - 'geom', + "geom", name=bone, mesh=bone, - pos=-bone_position['Mandible'], - dclass='light_bone', + pos=-bone_position["Mandible"], + dclass="light_bone", ) jaw_geoms.append(geom) if bone in lower_teeth: - geom.dclass = 'visible_bone' + geom.dclass = "visible_bone" # Jaw collision geoms: jaw_col_pos = [ (-0.03, 0, 0.01), @@ -279,49 +304,49 @@ def create_neck( jaw_col_angle = [55, 30, 25, 15] for i in range(4): jaw.add( - 'geom', - name='jaw' + str(i) + '_collision', - type='box', - dclass='collision_primitive', + "geom", + name="jaw" + str(i) + "_collision", + type="box", + dclass="collision_primitive", size=jaw_col_size[i], pos=jaw_col_pos[i], euler=[0, jaw_col_angle[i], 0], ) jaw.add( - 'joint', - name='mandible', - dclass='mandible', + "joint", + name="mandible", + dclass="mandible", axis=[0, 1, 0], pos=np.array((-0.043, 0, 0.05)), ) jaw.add( - 'site', - name='lower_bite', + "site", + name="lower_bite", size=(0.005,), - dclass='sensor', + dclass="sensor", pos=(0.063, 0, 0.005), ) - print('Make collision ellipsoids for teeth.') + print("Make collision ellipsoids for teeth.") visible_bones = upper_teeth + lower_teeth for bone in visible_bones: - bone_geom = torso.find('geom', bone) - bone_geom.type = 'ellipsoid' + bone_geom = torso.find("geom", bone) + bone_geom.type = "ellipsoid" physics = mjcf.Physics.from_mjcf_model(model) for bone in visible_bones: - bone_geom = torso.find('geom', bone) + bone_geom = torso.find("geom", bone) pos = physics.bind(bone_geom).pos quat = physics.bind(bone_geom).quat size = physics.bind(bone_geom).size bone_geom.parent.add( - 'geom', - name=bone + '_collision', - dclass='tooth_primitive', + "geom", + name=bone + "_collision", + dclass="tooth_primitive", pos=pos, size=size * 1.2, quat=quat, - type='ellipsoid', + type="ellipsoid", ) bone_geom.type = None diff --git a/dm_control/locomotion/walkers/assets/dog_v2/build_tail.py b/dm_control/locomotion/walkers/assets/dog_v2/build_tail.py index 40740bda..3218b685 100644 --- a/dm_control/locomotion/walkers/assets/dog_v2/build_tail.py +++ b/dm_control/locomotion/walkers/assets/dog_v2/build_tail.py @@ -17,13 +17,12 @@ import collections -from dm_control import mjcf import numpy as np +from dm_control import mjcf + -def create_tail( - caudal_dofs_per_vertebra, bone_size, model, bone_position, parent -): +def create_tail(caudal_dofs_per_vertebra, bone_size, model, bone_position, parent): """Add tail in the dog model. Args: @@ -38,8 +37,8 @@ def create_tail( A list of caudal joints. """ # Caudal spine (tail) bodies: - caudal_bones = ['Ca_' + str(i + 1) for i in range(21)] - parent_pos = bone_position['Pelvis'] + caudal_bones = ["Ca_" + str(i + 1) for i in range(21)] + parent_pos = bone_position["Pelvis"] caudal_bodies = [] caudal_geoms = [] for bone in caudal_bones: @@ -47,9 +46,10 @@ def create_tail( rel_pos = bone_pos - parent_pos xyaxes = np.hstack((-rel_pos, (0, 1, 0))) xyaxes[1] = 0 - child = parent.add('body', name=bone, pos=rel_pos) + child = parent.add("body", name=bone, pos=rel_pos) caudal_bodies.append(child) - geom = child.add('geom', name=bone, mesh=bone, pos=-bone_pos, dclass='bone') + geom = child.add("geom", name=bone, mesh=bone, + pos=-bone_pos, dclass="bone") caudal_geoms.append(geom) parent = child parent_pos = bone_pos @@ -59,29 +59,29 @@ def create_tail( # Caudal spine joints: caudal_axis = collections.OrderedDict() - caudal_axis['extend'] = np.array((0.0, 1.0, 0.0)) + caudal_axis["extend"] = np.array((0.0, 1.0, 0.0)) scale = np.asarray([bone_size[bone] for bone in caudal_bones]).mean() joint_pos = np.array((0.3, 0, 0.26)) * scale num_dofs = 0 caudal_joints = [] caudal_joint_names = [] - parent = model.find('geom', 'Sacrum') + parent = model.find("geom", "Sacrum") for i, vertebra in enumerate(caudal_bodies): while num_dofs < (i + 1) * caudal_dofs_per_vertebra: dof = num_dofs % 2 dof_name = list(caudal_axis.keys())[dof] - caudal_joint_names.append(vertebra.name + '_' + dof_name) + caudal_joint_names.append(vertebra.name + "_" + dof_name) rel_pos = physics.bind(parent).xpos - physics.bind(vertebra).xpos twist_dir = rel_pos / np.linalg.norm(rel_pos) - bend_dir = np.cross(caudal_axis['extend'], twist_dir) - caudal_axis['bend'] = bend_dir + bend_dir = np.cross(caudal_axis["extend"], twist_dir) + caudal_axis["bend"] = bend_dir joint_pos = twist_dir * physics.bind(caudal_geoms[i]).size[2] joint = vertebra.add( - 'joint', + "joint", name=caudal_joint_names[-1], - dclass='caudal_' + dof_name, + dclass="caudal_" + dof_name, axis=caudal_axis[dof_name], pos=joint_pos, ) @@ -89,22 +89,22 @@ def create_tail( num_dofs += 1 parent = vertebra - parent.add('site', name='tail_tip', dclass='sensor', size=(0.005,)) + parent.add("site", name="tail_tip", dclass="sensor", size=(0.005,)) physics = mjcf.Physics.from_mjcf_model(model) - all_geoms = model.find_all('geom') + all_geoms = model.find_all("geom") for geom in all_geoms: - if 'Ca_' in geom.name: + if "Ca_" in geom.name: sc = (float(geom.name[3:]) + 1) / 4 scale = np.array((1.2, sc, 1)) bound_geom = physics.bind(geom) geom.parent.add( - 'geom', - name=geom.name + '_collision', + "geom", + name=geom.name + "_collision", pos=bound_geom.pos, size=bound_geom.size * scale, quat=bound_geom.quat, - dclass='collision_primitive', + dclass="collision_primitive", ) return caudal_joints diff --git a/dm_control/locomotion/walkers/assets/dog_v2/build_torso.py b/dm_control/locomotion/walkers/assets/dog_v2/build_torso.py index dcaf3bba..8fef6c1b 100644 --- a/dm_control/locomotion/walkers/assets/dog_v2/build_torso.py +++ b/dm_control/locomotion/walkers/assets/dog_v2/build_torso.py @@ -17,12 +17,13 @@ import collections -from dm_control import mjcf import numpy as np +from dm_control import mjcf + def create_torso( - model, bones, bone_position, lumbar_dofs_per_vertebra, side_sign, parent + model, bones, bone_position, lumbar_dofs_per_vertebra, side_sign, parent, composer ): """Add torso in the dog model. @@ -35,61 +36,62 @@ def create_torso( side_sign: a dictionary with two axis representing the signs of translations. parent: parent object on which we should start attaching new components. - + composer: boolean to determine if a model used by the composer is being created. Returns: The tuple `(pelvic_bones, lumbar_joints)`. """ # Lumbar Spine - def_lumbar_extend = model.default.find('default', 'lumbar_extend') - def_lumbar_bend = model.default.find('default', 'lumbar_bend') - def_lumbar_twist = model.default.find('default', 'lumbar_twist') + def_lumbar_extend = model.default.find("default", "lumbar_extend") + def_lumbar_bend = model.default.find("default", "lumbar_bend") + def_lumbar_twist = model.default.find("default", "lumbar_twist") lumbar_defaults = { - 'extend': def_lumbar_extend, - 'bend': def_lumbar_bend, - 'twist': def_lumbar_twist, + "extend": def_lumbar_extend, + "bend": def_lumbar_bend, + "twist": def_lumbar_twist, } - thoracic_spine = [m for m in bones if 'T_' in m] - ribs = [m for m in bones if 'Rib' in m and 'cage' not in m] - sternum = [m for m in bones if 'Sternum' in m] + thoracic_spine = [m for m in bones if "T_" in m] + ribs = [m for m in bones if "Rib" in m and "cage" not in m] + sternum = [m for m in bones if "Sternum" in m] torso_bones = thoracic_spine + ribs + sternum # + ['Xiphoid_cartilage'] - torso = parent.add('body', name='torso') - torso.add('freejoint', name='root') - torso.add('site', name='root', size=(0.01,), rgba=[0, 1, 0, 1]) - torso.add('light', name='light', mode='trackcom', pos=[0, 0, 3]) + torso = parent.add("body", name="torso") + + if not composer: + torso.add("freejoint", name="root") + + torso.add("site", name="root", size=(0.01,), rgba=[0, 1, 0, 1]) + torso.add("light", name="light", mode="trackcom", pos=[0, 0, 3]) torso.add( - 'camera', - name='y-axis', - mode='trackcom', + "camera", + name="y-axis", + mode="trackcom", pos=[0, -1.5, 0.8], xyaxes=[1, 0, 0, 0, 0.6, 1], ) torso.add( - 'camera', - name='x-axis', - mode='trackcom', + "camera", + name="x-axis", + mode="trackcom", pos=[2, 0, 0.5], xyaxes=[0, 1, 0, -0.3, 0, 1], ) torso_geoms = [] for bone in torso_bones: torso_geoms.append( - torso.add('geom', name=bone, mesh=bone, dclass='light_bone') - ) + torso.add("geom", name=bone, mesh=bone, dclass="light_bone")) # Reload, get CoM position, set pos physics = mjcf.Physics.from_mjcf_model(model) - torso_pos = np.array(physics.bind(model.find('body', 'torso')).xipos) + torso_pos = np.array(physics.bind(model.find("body", "torso")).xipos) torso.pos = torso_pos for geom in torso_geoms: geom.pos = -torso_pos - # Collision primitive for torso torso.add( - 'geom', - name='collision_torso', - dclass='nonself_collision_primitive', - type='ellipsoid', + "geom", + name="collision_torso", + dclass="nonself_collision_primitive", + type="ellipsoid", pos=[0, 0, 0], size=[0.2, 0.09, 0.11], euler=[0, 10, 0], @@ -97,23 +99,24 @@ def create_torso( ) # Lumbar spine bodies: - lumbar_bones = ['L_1', 'L_2', 'L_3', 'L_4', 'L_5', 'L_6', 'L_7'] + lumbar_bones = ["L_1", "L_2", "L_3", "L_4", "L_5", "L_6", "L_7"] parent = torso parent_pos = torso_pos lumbar_bodies = [] lumbar_geoms = [] for i, bone in enumerate(lumbar_bones): bone_pos = bone_position[bone] - child = parent.add('body', name=bone, pos=bone_pos - parent_pos) + child = parent.add("body", name=bone, pos=bone_pos - parent_pos) lumbar_bodies.append(child) - geom = child.add('geom', name=bone, mesh=bone, pos=-bone_pos, dclass='bone') + geom = child.add("geom", name=bone, mesh=bone, + pos=-bone_pos, dclass="bone") child.add( - 'geom', - name=bone + '_collision', - type='sphere', + "geom", + name=bone + "_collision", + type="sphere", size=[0.05], pos=[0, 0, -0.02], - dclass='nonself_collision_primitive', + dclass="nonself_collision_primitive", ) lumbar_geoms.append(geom) parent = child @@ -122,9 +125,9 @@ def create_torso( # Lumbar spine joints: lumbar_axis = collections.OrderedDict() - lumbar_axis['extend'] = np.array((0.0, 1.0, 0.0)) - lumbar_axis['bend'] = np.array((0.0, 0.0, 1.0)) - lumbar_axis['twist'] = np.array((1.0, 0.0, 0)) + lumbar_axis["extend"] = np.array((0.0, 1.0, 0.0)) + lumbar_axis["bend"] = np.array((0.0, 0.0, 1.0)) + lumbar_axis["twist"] = np.array((1.0, 0.0, 0)) num_dofs = 0 lumbar_joints = [] @@ -134,11 +137,11 @@ def create_torso( dof = num_dofs % 3 dof_name = list(lumbar_axis.keys())[dof] dof_axis = lumbar_axis[dof_name] - lumbar_joint_names.append(vertebra.name + '_' + dof_name) + lumbar_joint_names.append(vertebra.name + "_" + dof_name) joint = vertebra.add( - 'joint', + "joint", name=lumbar_joint_names[-1], - dclass='lumbar_' + dof_name, + dclass="lumbar_" + dof_name, axis=dof_axis, ) lumbar_joints.append(joint) @@ -146,36 +149,35 @@ def create_torso( # Scale joint defaults relative to 3 lumbar_dofs_per_veterbra for dof in lumbar_axis.keys(): - axis_scale = 7.0 / [dof in joint for joint in lumbar_joint_names].count( - True - ) + axis_scale = 7.0 / \ + [dof in joint for joint in lumbar_joint_names].count(True) lumbar_defaults[dof].joint.range *= axis_scale # Pelvis: pelvis = l_7.add( - 'body', name='pelvis', pos=bone_position['Pelvis'] - bone_position['L_7'] + "body", name="pelvis", pos=bone_position["Pelvis"] - bone_position["L_7"] ) - pelvic_bones = ['Sacrum', 'Pelvis'] + pelvic_bones = ["Sacrum", "Pelvis"] pelvic_geoms = [] for bone in pelvic_bones: geom = pelvis.add( - 'geom', + "geom", name=bone, mesh=bone, - pos=-bone_position['Pelvis'], - dclass='bone', + pos=-bone_position["Pelvis"], + dclass="bone", ) pelvic_geoms.append(geom) # Collision primitives for pelvis - for side in ['_L', '_R']: + for side in ["_L", "_R"]: pos = np.array((0.01, -0.02, -0.01)) * side_sign[side] pelvis.add( - 'geom', - name='collision_pelvis' + side, + "geom", + name="collision_pelvis" + side, pos=pos, size=[0.05, 0.05, 0], euler=[0, 70, 0], - dclass='nonself_collision_primitive', + dclass="nonself_collision_primitive", ) return pelvic_bones, lumbar_joints diff --git a/dm_control/locomotion/walkers/assets/dog_v2/create_skin.py b/dm_control/locomotion/walkers/assets/dog_v2/create_skin.py index 3ef74f0e..6a695e0b 100644 --- a/dm_control/locomotion/walkers/assets/dog_v2/create_skin.py +++ b/dm_control/locomotion/walkers/assets/dog_v2/create_skin.py @@ -17,56 +17,71 @@ import struct -from dm_control import mjcf -from dm_control.mujoco.wrapper.mjbindings import enums import numpy as np from scipy import spatial +from dm_control import mjcf +from dm_control.mujoco.wrapper.mjbindings import enums + -def create(model, skin_msh): +def create(model, mesh_file, mesh_name, asset_dir, tex_coords=True, transform=True, composer=False): """Create and add skin in the dog model. Args: model: model in which we want to add the skin. - skin_msh: a binary mesh format of the skin. + mesh_file: a binary mesh format of the skin. + mesh_name: the name of the skin to use in the xml + asset_dir: asset directory to load the .skn file + tex_coords: boolean to indicate if the mesh has texture coordinates. + transform: a boolean to rotate mesh orientation by 90 degrees along the z-axis. + composer: boolean to determine if a model used by the composer is being created. """ - print('Making Skin.') # Add skin mesh: - skinmesh = model.worldbody.add( - 'geom', - name='skinmesh', - mesh='skin_msh', - type='mesh', - contype=0, - conaffinity=0, - rgba=[1, 0.5, 0.5, 0.5], - group=1, - euler=(0, 0, 90), - ) + if transform: + skinmesh = model.worldbody.add( + "geom", + name=mesh_name, + mesh=mesh_name, + type="mesh", + contype=0, + conaffinity=0, + rgba=[1, 0.5, 0.5, 0.5], + group=1, + euler=(0, 0, 90), + ) + else: + skinmesh = model.worldbody.add( + "geom", + name=mesh_name, + mesh=mesh_name, + type="mesh", + contype=0, + conaffinity=0, + ) physics = mjcf.Physics.from_mjcf_model(model) # Get skinmesh vertices in global coordinates - vertadr = physics.named.model.mesh_vertadr['skin_msh'] - vertnum = physics.named.model.mesh_vertnum['skin_msh'] - skin_vertices = physics.model.mesh_vert[vertadr : vertadr + vertnum, :] + vertadr = physics.named.model.mesh_vertadr[mesh_name] + vertnum = physics.named.model.mesh_vertnum[mesh_name] + skin_vertices = physics.model.mesh_vert[vertadr: vertadr + vertnum, :] skin_vertices = skin_vertices.dot( - physics.named.data.geom_xmat['skinmesh'].reshape(3, 3).T + physics.named.data.geom_xmat[mesh_name].reshape(3, 3).T ) - skin_vertices += physics.named.data.geom_xpos['skinmesh'] - skin_normals = physics.model.mesh_normal[vertadr : vertadr + vertnum, :] + skin_vertices += physics.named.data.geom_xpos[mesh_name] + skin_normals = physics.model.mesh_normal[vertadr: vertadr + vertnum, :] skin_normals = skin_normals.dot( - physics.named.data.geom_xmat['skinmesh'].reshape(3, 3).T + physics.named.data.geom_xmat[mesh_name].reshape(3, 3).T ) - skin_normals += physics.named.data.geom_xpos['skinmesh'] + skin_normals += physics.named.data.geom_xpos[mesh_name] # Get skinmesh faces - faceadr = physics.named.model.mesh_faceadr['skin_msh'] - facenum = physics.named.model.mesh_facenum['skin_msh'] - skin_faces = physics.model.mesh_face[faceadr : faceadr + facenum, :] + faceadr = physics.named.model.mesh_faceadr[mesh_name] + facenum = physics.named.model.mesh_facenum[mesh_name] + skin_faces = physics.model.mesh_face[faceadr: faceadr + facenum, :] # Make skin skin = model.asset.add( - 'skin', name='skin', vertex=skin_vertices.ravel(), face=skin_faces.ravel() + "skin", name="skin_tmp", vertex=skin_vertices.ravel(), face=skin_faces.ravel() ) # Functions for capsule vertices @@ -117,8 +132,9 @@ def box(sx, sy, sz): # Find smallest distance between # each skin vertex and vertices of all meshes in body i - distance = np.zeros((skin_vertices.shape[0], physics.model.nbody)) - for i in range(1, physics.model.nbody): + distance = np.ones((skin_vertices.shape[0], physics.model.nbody)) * 1e6 + start = 1 if composer else 2 # in composer mode we don't have the floor + for i in range(start, physics.model.nbody): geom_id = np.argwhere(physics.model.geom_bodyid == i).ravel() mesh_id = physics.model.geom_dataid[geom_id] body_verts = [] @@ -127,7 +143,7 @@ def box(sx, sy, sz): if physics.model.geom_type[gid] == enums.mjtGeom.mjGEOM_MESH: vertadr = physics.model.mesh_vertadr[mesh_id[k]] vertnum = physics.model.mesh_vertnum[mesh_id[k]] - vertices = physics.model.mesh_vert[vertadr : vertadr + vertnum, :] + vertices = physics.model.mesh_vert[vertadr: vertadr + vertnum, :] elif physics.model.geom_type[gid] == enums.mjtGeom.mjGEOM_CAPSULE: radius = physics.model.geom_size[gid, 0] height = physics.model.geom_size[gid, 1] @@ -161,13 +177,13 @@ def box(sx, sy, sz): weights[weights < threshold] = 0 weights /= np.atleast_2d(np.sum(weights, axis=1)).T - for i in range(1, physics.model.nbody): + for i in range(start, physics.model.nbody): vertweight = weights[weights[:, i - 1] >= threshold, i - 1] vertid = np.argwhere(weights[:, i - 1] >= threshold).ravel() if vertid.any(): skin.add( - 'bone', - body=physics.model.id2name(i, 'body'), + "bone", + body=physics.model.id2name(i, "body"), bindquat=[1, 0, 0, 0], bindpos=physics.data.xpos[i, :], vertid=vertid, @@ -179,31 +195,38 @@ def box(sx, sy, sz): # Convert skin into *.skn file according to # https://mujoco.readthedocs.io/en/latest/XMLreference.html#asset-skin - f = open('dog_skin.skn', 'w+b') + f = open(asset_dir + "/skins/" + mesh_name + ".skn", "w+b") + nvert = skin.vertex.size // 3 - f.write( - struct.pack( - '4i', nvert, nvert, skin.face.size // 3, physics.model.nbody - 1 - ) - ) - f.write(struct.pack(str(skin.vertex.size) + 'f', *skin.vertex)) - assert physics.model.mesh_texcoord.shape[0] == physics.bind(skin_msh).vertnum - f.write( - struct.pack(str(2 * nvert) + 'f', *physics.model.mesh_texcoord.flatten()) - ) - f.write(struct.pack(str(skin.face.size) + 'i', *skin.face)) + + if tex_coords: + n_tex_coord = nvert + else: + n_tex_coord = 0 + f.write(struct.pack("4i", nvert, n_tex_coord, + skin.face.size // 3, len(skin.bone))) + f.write(struct.pack(str(skin.vertex.size) + "f", *skin.vertex)) + + if n_tex_coord: + assert physics.model.mesh_texcoord.shape[0] == physics.bind( + mesh_file).vertnum + f.write( + struct.pack(str(2 * nvert) + "f", * + physics.model.mesh_texcoord.flatten()) + ) + + f.write(struct.pack(str(skin.face.size) + "i", *skin.face)) for bone in skin.bone: name_length = len(bone.body) assert name_length <= 40 - f.write( - struct.pack(str(name_length) + 'c', *[s.encode() for s in bone.body]) - ) - f.write((40 - name_length) * b'\x00') - f.write(struct.pack('3f', *bone.bindpos)) - f.write(struct.pack('4f', *bone.bindquat)) - f.write(struct.pack('i', bone.vertid.size)) - f.write(struct.pack(str(bone.vertid.size) + 'i', *bone.vertid)) - f.write(struct.pack(str(bone.vertid.size) + 'f', *bone.vertweight)) + f.write(struct.pack(str(name_length) + "c", + *[s.encode() for s in bone.body])) + f.write((40 - name_length) * b"\x00") + f.write(struct.pack("3f", *bone.bindpos)) + f.write(struct.pack("4f", *bone.bindquat)) + f.write(struct.pack("i", bone.vertid.size)) + f.write(struct.pack(str(bone.vertid.size) + "i", *bone.vertid)) + f.write(struct.pack(str(bone.vertid.size) + "f", *bone.vertweight)) f.close() # Remove XML-based skin, add binary skin. diff --git a/dm_control/locomotion/walkers/assets/dog_v2/dog.xml b/dm_control/locomotion/walkers/assets/dog_v2/dog.xml index daf7f537..04e7441a 100644 --- a/dm_control/locomotion/walkers/assets/dog_v2/dog.xml +++ b/dm_control/locomotion/walkers/assets/dog_v2/dog.xml @@ -1,14 +1,15 @@ -