deploy: 089015f

Author: ooctipus

main/_modules/isaaclab/envs/mdp/events.html

@@ -1599,28 +1599,30 @@ <h1>Source code for isaaclab.envs.mdp.events</h1><div class="highlight"><pre>
 <span class="sd">    &quot;&quot;&quot;</span>
     <span class="c1"># extract the used quantities (to enable type-hinting)</span>
     <span class="n">asset</span><span class="p">:</span> <span class="n">Articulation</span> <span class="o">=</span> <span class="n">env</span><span class="o">.</span><span class="n">scene</span><span class="p">[</span><span class="n">asset_cfg</span><span class="o">.</span><span class="n">name</span><span class="p">]</span>
+
+    <span class="c1"># cast env_ids to allow broadcasting</span>
+    <span class="k">if</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span> <span class="o">!=</span> <span class="nb">slice</span><span class="p">(</span><span class="kc">None</span><span class="p">):</span>
+        <span class="n">iter_env_ids</span> <span class="o">=</span> <span class="n">env_ids</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">]</span>
+    <span class="k">else</span><span class="p">:</span>
+        <span class="n">iter_env_ids</span> <span class="o">=</span> <span class="n">env_ids</span>
+
     <span class="c1"># get default joint state</span>
-    <span class="n">joint_pos</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">default_joint_pos</span><span class="p">[</span><span class="n">env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
-    <span class="n">joint_vel</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">default_joint_vel</span><span class="p">[</span><span class="n">env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
+    <span class="n">joint_pos</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">default_joint_pos</span><span class="p">[</span><span class="n">iter_env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
+    <span class="n">joint_vel</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">default_joint_vel</span><span class="p">[</span><span class="n">iter_env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
 
     <span class="c1"># scale these values randomly</span>
     <span class="n">joint_pos</span> <span class="o">*=</span> <span class="n">math_utils</span><span class="o">.</span><span class="n">sample_uniform</span><span class="p">(</span><span class="o">*</span><span class="n">position_range</span><span class="p">,</span> <span class="n">joint_pos</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">joint_pos</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
     <span class="n">joint_vel</span> <span class="o">*=</span> <span class="n">math_utils</span><span class="o">.</span><span class="n">sample_uniform</span><span class="p">(</span><span class="o">*</span><span class="n">velocity_range</span><span class="p">,</span> <span class="n">joint_vel</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">joint_vel</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
 
     <span class="c1"># clamp joint pos to limits</span>
-    <span class="n">joint_pos_limits</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">soft_joint_pos_limits</span><span class="p">[</span><span class="n">env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span>
+    <span class="n">joint_pos_limits</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">soft_joint_pos_limits</span><span class="p">[</span><span class="n">iter_env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span>
     <span class="n">joint_pos</span> <span class="o">=</span> <span class="n">joint_pos</span><span class="o">.</span><span class="n">clamp_</span><span class="p">(</span><span class="n">joint_pos_limits</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">joint_pos_limits</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">1</span><span class="p">])</span>
     <span class="c1"># clamp joint vel to limits</span>
-    <span class="n">joint_vel_limits</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">soft_joint_vel_limits</span><span class="p">[</span><span class="n">env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span>
+    <span class="n">joint_vel_limits</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">soft_joint_vel_limits</span><span class="p">[</span><span class="n">iter_env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span>
     <span class="n">joint_vel</span> <span class="o">=</span> <span class="n">joint_vel</span><span class="o">.</span><span class="n">clamp_</span><span class="p">(</span><span class="o">-</span><span class="n">joint_vel_limits</span><span class="p">,</span> <span class="n">joint_vel_limits</span><span class="p">)</span>
 
     <span class="c1"># set into the physics simulation</span>
-    <span class="n">asset</span><span class="o">.</span><span class="n">write_joint_state_to_sim</span><span class="p">(</span>
-        <span class="n">joint_pos</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">env_ids</span><span class="p">),</span> <span class="o">-</span><span class="mi">1</span><span class="p">),</span>
-        <span class="n">joint_vel</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">env_ids</span><span class="p">),</span> <span class="o">-</span><span class="mi">1</span><span class="p">),</span>
-        <span class="n">env_ids</span><span class="o">=</span><span class="n">env_ids</span><span class="p">,</span>
-        <span class="n">joint_ids</span><span class="o">=</span><span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">,</span>
-    <span class="p">)</span></div>
+    <span class="n">asset</span><span class="o">.</span><span class="n">write_joint_state_to_sim</span><span class="p">(</span><span class="n">joint_pos</span><span class="p">,</span> <span class="n">joint_vel</span><span class="p">,</span> <span class="n">joint_ids</span><span class="o">=</span><span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">,</span> <span class="n">env_ids</span><span class="o">=</span><span class="n">env_ids</span><span class="p">)</span></div>
 
 
 <div class="viewcode-block" id="reset_joints_by_offset"><a class="viewcode-back" href="../../../../source/api/lab/isaaclab.envs.mdp.html#isaaclab.envs.mdp.events.reset_joints_by_offset">[docs]</a><span class="k">def</span><span class="w"> </span><span class="nf">reset_joints_by_offset</span><span class="p">(</span>
@@ -1638,28 +1640,29 @@ <h1>Source code for isaaclab.envs.mdp.events</h1><div class="highlight"><pre>
     <span class="c1"># extract the used quantities (to enable type-hinting)</span>
     <span class="n">asset</span><span class="p">:</span> <span class="n">Articulation</span> <span class="o">=</span> <span class="n">env</span><span class="o">.</span><span class="n">scene</span><span class="p">[</span><span class="n">asset_cfg</span><span class="o">.</span><span class="n">name</span><span class="p">]</span>
 
+    <span class="c1"># cast env_ids to allow broadcasting</span>
+    <span class="k">if</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span> <span class="o">!=</span> <span class="nb">slice</span><span class="p">(</span><span class="kc">None</span><span class="p">):</span>
+        <span class="n">iter_env_ids</span> <span class="o">=</span> <span class="n">env_ids</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">]</span>
+    <span class="k">else</span><span class="p">:</span>
+        <span class="n">iter_env_ids</span> <span class="o">=</span> <span class="n">env_ids</span>
+
     <span class="c1"># get default joint state</span>
-    <span class="n">joint_pos</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">default_joint_pos</span><span class="p">[</span><span class="n">env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
-    <span class="n">joint_vel</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">default_joint_vel</span><span class="p">[</span><span class="n">env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
+    <span class="n">joint_pos</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">default_joint_pos</span><span class="p">[</span><span class="n">iter_env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
+    <span class="n">joint_vel</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">default_joint_vel</span><span class="p">[</span><span class="n">iter_env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
 
     <span class="c1"># bias these values randomly</span>
     <span class="n">joint_pos</span> <span class="o">+=</span> <span class="n">math_utils</span><span class="o">.</span><span class="n">sample_uniform</span><span class="p">(</span><span class="o">*</span><span class="n">position_range</span><span class="p">,</span> <span class="n">joint_pos</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">joint_pos</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
     <span class="n">joint_vel</span> <span class="o">+=</span> <span class="n">math_utils</span><span class="o">.</span><span class="n">sample_uniform</span><span class="p">(</span><span class="o">*</span><span class="n">velocity_range</span><span class="p">,</span> <span class="n">joint_vel</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">joint_vel</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
 
     <span class="c1"># clamp joint pos to limits</span>
-    <span class="n">joint_pos_limits</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">soft_joint_pos_limits</span><span class="p">[</span><span class="n">env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span>
+    <span class="n">joint_pos_limits</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">soft_joint_pos_limits</span><span class="p">[</span><span class="n">iter_env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span>
     <span class="n">joint_pos</span> <span class="o">=</span> <span class="n">joint_pos</span><span class="o">.</span><span class="n">clamp_</span><span class="p">(</span><span class="n">joint_pos_limits</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">joint_pos_limits</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">1</span><span class="p">])</span>
     <span class="c1"># clamp joint vel to limits</span>
-    <span class="n">joint_vel_limits</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">soft_joint_vel_limits</span><span class="p">[</span><span class="n">env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span>
+    <span class="n">joint_vel_limits</span> <span class="o">=</span> <span class="n">asset</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">soft_joint_vel_limits</span><span class="p">[</span><span class="n">iter_env_ids</span><span class="p">,</span> <span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">]</span>
     <span class="n">joint_vel</span> <span class="o">=</span> <span class="n">joint_vel</span><span class="o">.</span><span class="n">clamp_</span><span class="p">(</span><span class="o">-</span><span class="n">joint_vel_limits</span><span class="p">,</span> <span class="n">joint_vel_limits</span><span class="p">)</span>
 
     <span class="c1"># set into the physics simulation</span>
-    <span class="n">asset</span><span class="o">.</span><span class="n">write_joint_state_to_sim</span><span class="p">(</span>
-        <span class="n">joint_pos</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">env_ids</span><span class="p">),</span> <span class="o">-</span><span class="mi">1</span><span class="p">),</span>
-        <span class="n">joint_vel</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">env_ids</span><span class="p">),</span> <span class="o">-</span><span class="mi">1</span><span class="p">),</span>
-        <span class="n">env_ids</span><span class="o">=</span><span class="n">env_ids</span><span class="p">,</span>
-        <span class="n">joint_ids</span><span class="o">=</span><span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">,</span>
-    <span class="p">)</span></div>
+    <span class="n">asset</span><span class="o">.</span><span class="n">write_joint_state_to_sim</span><span class="p">(</span><span class="n">joint_pos</span><span class="p">,</span> <span class="n">joint_vel</span><span class="p">,</span> <span class="n">joint_ids</span><span class="o">=</span><span class="n">asset_cfg</span><span class="o">.</span><span class="n">joint_ids</span><span class="p">,</span> <span class="n">env_ids</span><span class="o">=</span><span class="n">env_ids</span><span class="p">)</span></div>
 
 
 <div class="viewcode-block" id="reset_nodal_state_uniform"><a class="viewcode-back" href="../../../../source/api/lab/isaaclab.envs.mdp.html#isaaclab.envs.mdp.events.reset_nodal_state_uniform">[docs]</a><span class="k">def</span><span class="w"> </span><span class="nf">reset_nodal_state_uniform</span><span class="p">(</span>