Working 1 bot per env 20k sps

pufferlib/environments/isaacgym/environment.py

@@ -3,6 +3,7 @@
 import gymnasium as gym
 import numpy as np
 import functools
+import math
 import sys
 import os
 
@@ -297,7 +298,7 @@ def __init__(self, num_envs=8, control_freq_inv=2, headless=True, device='cuda',
             char_h = 0.89
             start_pose = gymapi.Transform()
             start_pose.p = gymapi.Vec3(*get_axis_params(char_h, 2))
-            start_pose.r = gymapi.Quat(0.0, 0.0, 0.0, 1.0)
+            start_pose.r = gymapi.Quat.from_euler_zyx(math.pi / 2, 0, 0)
 
             humanoid_handle = self.gym.create_actor(env_ptr, humanoid_asset,
                 start_pose, "humanoid", group=env_id, filter=0, segmentationId=0)
@@ -538,9 +539,9 @@ def reset(self, env_ids=slice(None)):
         #self._compute_observations(env_ids)
 
     def step(self, actions):
-        #self.actions = actions.to(self.device).clone()
-        #pd_tar_tensor = gymtorch.unwrap_tensor(self._pd_action_offset + self._pd_action_scale*self.actions)
-        #self.gym.set_dof_position_target_tensor(self.sim, pd_tar_tensor)
+        self.actions = actions.to(self.device).clone()
+        pd_tar_tensor = gymtorch.unwrap_tensor(self._pd_action_offset + self._pd_action_scale*self.actions)
+        self.gym.set_dof_position_target_tensor(self.sim, pd_tar_tensor)
 
         self._refresh_sim_tensors()
         self.env.step()

test_isaac.py

@@ -9,14 +9,32 @@ def test_base_env():
 def test_humanoid():
     env = HumanoidSMPLX(headless=False)
     env.reset()
-    for i in range(10):
+    env.render()
+    while True:
         action = torch.from_numpy(env.action_space.sample())
         env.step(action)
         env.render()
-    breakpoint()
-    pass
+
+def test_humanoid_perf(timeout=10, num_envs=4096):
+    env = HumanoidSMPLX(num_envs=num_envs, headless=True)
+    env.reset()
+
+    import time
+    start = time.time()
+    action = torch.from_numpy(env.action_space.sample())
+
+    steps = 0
+    while time.time() - start < timeout:
+        env.step(action)
+        steps += num_envs
+
+    end = time.time()
+    sps = steps / (end - start)
+    print(f"Steps: {steps}, SPS: {sps}")
+
 
 if __name__ == '__main__':
     #test_base_env()
-    test_humanoid()
+    #test_humanoid()
+    test_humanoid_perf()