From 7513a117562e5d8f8de9077992f1bcab38660bf5 Mon Sep 17 00:00:00 2001 From: Joseph Suarez Date: Thu, 23 Jan 2025 20:39:48 +0000 Subject: [PATCH] Working 1 bot per env 20k sps --- .../environments/isaacgym/environment.py | 9 ++++--- test_isaac.py | 26 ++++++++++++++++--- 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/pufferlib/environments/isaacgym/environment.py b/pufferlib/environments/isaacgym/environment.py index 0ae10bbd..6361d116 100644 --- a/pufferlib/environments/isaacgym/environment.py +++ b/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() diff --git a/test_isaac.py b/test_isaac.py index 044328e5..00e9d448 100644 --- a/test_isaac.py +++ b/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()