chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Drop find_ prefix in Python MjSpec methods accessing elements by name.

Browse Source 
PiperOrigin-RevId: 725594833
Change-Id: Ibf8a9918e98b298751d0ef2e657714e28e5fa460
This commit is contained in:
Alessio Quaglino
2025-02-11 05:49:48 -08:00
committed by Copybara-Service
parent a6dd5f086f
commit cef5b31a3f
4 changed files with 35 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mjspec.ipynb
+20 -18
View File
@@ -573,11 +573,11 @@
" self.spec = mj.MjSpec.from_string(leg_model)\n",
"\n",
" # Thigh:\n",
" thigh = self.spec.find_body('thigh')\n",
" thigh = self.spec.body('thigh')\n",
" thigh.add_geom(fromto=[0, 0, 0, length, 0, 0], size=[length/4, 0, 0], rgba=rgba)\n",
"\n",
" # Hip:\n",
" shin = self.spec.find_body('shin')\n",
" shin = self.spec.body('shin')\n",
" shin.add_geom(fromto=[0, 0, 0, 0, 0, -length], size=[length/5, 0, 0], rgba=rgba)\n",
" shin.pos[0] = length"
]
@@ -620,6 +620,7 @@
" \"\"\"Constructs a creature with `num_legs` legs.\"\"\"\n",
" rgba = random_state.uniform([0, 0, 0, 1], [1, 1, 1, 1])\n",
" spec = mj.MjSpec.from_string(creature_model)\n",
" spec.copy_during_attach = True\n",
"\n",
" # Attach legs to equidistant sites on the circumference.\n",
" spec.worldbody.first_geom().rgba = rgba\n",
@@ -628,7 +629,7 @@
" theta = 2 * i * np.pi / num_legs\n",
" hip_pos = BODY_RADIUS * np.array([np.cos(theta), np.sin(theta), 0])\n",
" hip_site = spec.worldbody.add_site(pos=hip_pos, euler=[0, 0, theta])\n",
" hip_site.attach_body(leg.spec.find_body('thigh'), '', '-' + str(i))\n",
" hip_site.attach_body(leg.spec.body('thigh'), '', '-' + str(i))\n",
"\n",
" return spec"
]
@@ -870,7 +871,7 @@
"# # Add another humanoid\n",
"spec_humanoid = mj.MjSpec.from_file(humanoid_file)\n",
"attachment_frame = spec.worldbody.add_frame(pos=[0, -1, 2])\n",
"attachment_frame.attach_body(spec_humanoid.find_body('torso'), 'a', 'b')\n",
"attachment_frame.attach_body(spec_humanoid.body('torso'), 'a', 'b')\n",
"\n",
"# Recompile preserving the state\n",
"new_model, new_data = spec.recompile(model, data)\n",
@@ -924,13 +925,14 @@
"#@title Humanoid with arms replaced by legs.{vertical-output: true}\n",
"\n",
"spec = mj.MjSpec.from_file(humanoid_file)\n",
"spec.copy_during_attach = True\n",
"\n",
"# Get the torso, arm, and leg bodies\n",
"arm_left = spec.find_body('upper_arm_left')\n",
"arm_right = spec.find_body('upper_arm_right')\n",
"leg_left = spec.find_body('thigh_left')\n",
"leg_right = spec.find_body('thigh_right')\n",
"torso = spec.find_body('torso')\n",
"arm_left = spec.body('upper_arm_left')\n",
"arm_right = spec.body('upper_arm_right')\n",
"leg_left = spec.body('thigh_left')\n",
"leg_right = spec.body('thigh_right')\n",
"torso = spec.body('torso')\n",
"\n",
"# Attach frames at the arm positions\n",
"shoulder_left = torso.add_frame(pos=arm_left.pos)\n",
@@ -976,13 +978,13 @@
" spec.degree = False # MuJoCo release\n",
"\n",
"# Replace right arm with frame\n",
"arm_right = spec.find_body('upper_arm_right')\n",
"torso = spec.find_body('torso')\n",
"arm_right = spec.body('upper_arm_right')\n",
"torso = spec.body('torso')\n",
"shoulder_right = torso.add_frame(pos=arm_right.pos, quat=[0, 0.8509035, 0, 0.525322])\n",
"spec.detach_body(arm_right)\n",
"\n",
"# Attach Franka arm to humanoid\n",
"franka_arm = franka.find_body('fr3_link2')\n",
"franka_arm = franka.body('fr3_link2')\n",
"shoulder_right.attach_body(franka_arm, 'franka', '')\n",
"\n",
"model = spec.compile()\n",
@@ -1033,6 +1035,7 @@
"\n",
"humanoid = mj.MjSpec.from_file(humanoid_file)\n",
"spec = mj.MjSpec()\n",
"spec.copy_during_attach = True\n",
"\n",
"# Delete all key frames to avoid name conflicts\n",
"while humanoid.keys:\n",
@@ -1044,18 +1047,18 @@
" humanoid.materials[0].rgba = [\n",
" np.random.uniform(), np.random.uniform(),\n",
" np.random.uniform(), 1] # Randomize color\n",
" humanoid.find_body('head').first_geom().size = [\n",
" humanoid.body('head').first_geom().size = [\n",
" .18*np.random.uniform(), 0, 0] # Randomize head size\n",
" humanoid.find_body('upper_arm_left').quat = [\n",
" humanoid.body('upper_arm_left').quat = [\n",
" np.random.uniform(), np.random.uniform(),\n",
" np.random.uniform(), np.random.uniform()] # Randomize left arm orientation\n",
" humanoid.find_body('upper_arm_right').quat = [\n",
" humanoid.body('upper_arm_right').quat = [\n",
" np.random.uniform(), np.random.uniform(),\n",
" np.random.uniform(), np.random.uniform()] # Randomize right arm orientation\n",
"\n",
" # attach randomized humanoid to parent spec\n",
" frame = spec.worldbody.add_frame(pos=[i, j, 0])\n",
" frame.attach_body(humanoid.find_body('torso'), str(i), str(j))\n",
" frame.attach_body(humanoid.body('torso'), str(i), str(j))\n",
"\n",
"spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n",
" targetbody='3torso3', diffuse=[.8, .8, .8],\n",
@@ -1072,8 +1075,7 @@
"sJFuNetilv4m"
],
"gpuClass": "premium",
"private_outputs": true,
"toc_visible": true
"private_outputs": true
},
"gpuClass": "premium",
"kernelspec": {