116 lines
3.8 KiB
Python
116 lines
3.8 KiB
Python
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
|
|
#
|
|
# 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.
|
|
# ==============================================================================
|
|
|
|
|
|
"""Benchmark environment step performance."""
|
|
|
|
import time
|
|
|
|
import numpy as np
|
|
from absl import app, flags
|
|
|
|
from motrix_envs import registry
|
|
|
|
FLAGS = flags.FLAGS
|
|
|
|
flags.DEFINE_string("env", "cartpole", "Environment name to benchmark")
|
|
flags.DEFINE_integer("num_steps", 1000, "Number of steps to benchmark")
|
|
flags.DEFINE_boolean("random_actions", True, "Use random actions (True) or zero actions (False)")
|
|
flags.DEFINE_integer("num_envs", 1, "Number of parallel environments")
|
|
flags.DEFINE_string("sim_backend", None, "Simulation backend (auto-select if None)")
|
|
|
|
|
|
def generate_action(env, random_actions: bool) -> np.ndarray:
|
|
"""Generate action for the environment.
|
|
|
|
Args:
|
|
env: The environment instance.
|
|
random_actions: If True, sample random actions from action space.
|
|
If False, use zero actions.
|
|
|
|
Returns:
|
|
Action array with shape (num_envs, *action_space.shape).
|
|
"""
|
|
action_space = env.action_space
|
|
|
|
if random_actions:
|
|
# Sample random actions within bounds
|
|
low = action_space.low
|
|
high = action_space.high
|
|
|
|
# Handle infinite bounds
|
|
low = np.where(np.isneginf(low), -1e6, low)
|
|
high = np.where(np.isposinf(high), 1e6, high)
|
|
|
|
size = (env.num_envs, *action_space.shape)
|
|
return np.random.uniform(low=low, high=high, size=size).astype(action_space.dtype)
|
|
else:
|
|
# Use zero actions
|
|
return np.zeros((env.num_envs, *action_space.shape), dtype=action_space.dtype)
|
|
|
|
|
|
def main(argv):
|
|
"""Main benchmark function."""
|
|
del argv # Unused
|
|
|
|
env_name = FLAGS.env
|
|
num_steps = FLAGS.num_steps
|
|
random_actions = FLAGS.random_actions
|
|
num_envs = FLAGS.num_envs
|
|
sim_backend = FLAGS.sim_backend
|
|
|
|
# Create environment
|
|
print(f"Creating environment: {env_name}")
|
|
env = registry.make(env_name, sim_backend=sim_backend, num_envs=num_envs)
|
|
|
|
# Print environment info
|
|
print("\nBenchmark configuration:")
|
|
print(f" Environment: {env_name}")
|
|
print(f" Number of parallel environments: {num_envs}")
|
|
print(f" Action space shape: {env.action_space.shape}")
|
|
print(f" Random actions: {random_actions}")
|
|
print(f" Number of Batch steps: {num_steps}")
|
|
print("\nRunning benchmark...\n")
|
|
|
|
# Generate action
|
|
action = generate_action(env, random_actions)
|
|
|
|
# Warmup run (to reduce cold start effects)
|
|
# Note: step() automatically initializes state if needed
|
|
for _ in range(10):
|
|
env.step(action)
|
|
|
|
# Benchmark loop
|
|
start_time = time.perf_counter()
|
|
for _ in range(num_steps):
|
|
env.step(action)
|
|
end_time = time.perf_counter()
|
|
|
|
# Calculate metrics
|
|
total_time = end_time - start_time
|
|
steps_per_second = num_steps / total_time
|
|
time_per_step_ms = (total_time / num_steps) * 1000
|
|
|
|
# Print results
|
|
print("Results:")
|
|
print(f" Total time: {total_time:.4f} seconds")
|
|
print(f" Batch Steps per second: {steps_per_second:.2f}")
|
|
print(f" Total Steps per seconds: {steps_per_second * num_envs:.2f}")
|
|
print(f" Time per batch step: {time_per_step_ms:.4f} ms")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run(main)
|