chore: release v0.0.1
This commit is contained in:
motphys-developers
40
scripts/gpu_utils.py
Normal file
40
scripts/gpu_utils.py
Normal file
@@ -0,0 +1,40 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
|
||||
import jax.numpy as jnp
|
||||
import pynvml
|
||||
|
||||
|
||||
def monitor_gpu_utilization(stop_event, gpu_index=0, interval=1.0):
|
||||
pynvml.nvmlInit()
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_index)
|
||||
utilization_samples = []
|
||||
|
||||
while not stop_event.is_set():
|
||||
util = pynvml.nvmlDeviceGetUtilizationRates(handle)
|
||||
utilization_samples.append(util.gpu)
|
||||
stop_event.wait(interval)
|
||||
|
||||
pynvml.nvmlShutdown()
|
||||
|
||||
if utilization_samples:
|
||||
data = jnp.array(utilization_samples)
|
||||
print(f"GPU utilization statistics over {len(data)} samples:")
|
||||
print(f" Mean: {jnp.mean(data):.2f}%")
|
||||
print(f" Max : {jnp.max(data):.2f}%")
|
||||
print(f" Min : {jnp.min(data):.2f}%")
|
||||
print(f" Median : {jnp.median(data):.2f}%")
|
||||
else:
|
||||
print("No GPU utilization samples recorded.")
|
||||
161
scripts/play.py
Normal file
161
scripts/play.py
Normal file
@@ -0,0 +1,161 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from absl import app, flags
|
||||
from skrl import config
|
||||
|
||||
from motrix_rl import utils
|
||||
from motrix_rl.skrl import get_log_dir
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_ENV = flags.DEFINE_string("env", "cartpole", "The env to play")
|
||||
_SIM_BACKEND = flags.DEFINE_string(
|
||||
"sim-backend",
|
||||
None,
|
||||
"The simulation backend to use.(If not specified, it will be choosen automatically)",
|
||||
)
|
||||
_POLICY = flags.DEFINE_string("policy", None, "The policy to load")
|
||||
_NUM_ENVS = flags.DEFINE_integer("num-envs", 2048, "Number of envs to play")
|
||||
_SEED = flags.DEFINE_integer("seed", None, "Random seed for reproducibility")
|
||||
_RAND_SEED = flags.DEFINE_bool("rand-seed", False, "Generate random seed")
|
||||
|
||||
|
||||
def get_inference_backend(policy_path: str):
|
||||
if policy_path.endswith(".pt"):
|
||||
return "torch"
|
||||
if policy_path.endswith(".pickle"):
|
||||
return "jax"
|
||||
else:
|
||||
raise Exception(f"Unknown policy format: {policy_path}")
|
||||
|
||||
|
||||
def find_best_policy(env_name: str) -> str:
|
||||
"""
|
||||
Find the most recent best policy for the given environment.
|
||||
|
||||
Args:
|
||||
env_name: The name of the environment
|
||||
|
||||
Returns:
|
||||
Path to the best policy file
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If no policy files are found
|
||||
"""
|
||||
# Base runs directory
|
||||
|
||||
env_dir = Path(get_log_dir(env_name))
|
||||
|
||||
if not env_dir.exists():
|
||||
raise FileNotFoundError(f"No training results found for environment '{env_name}' in {env_dir}")
|
||||
|
||||
# Find all training run directories (pattern: YY-MM-DD_HH-MM-SS-_XXXXX_PPO)
|
||||
training_runs = [d for d in env_dir.iterdir() if d.is_dir()]
|
||||
|
||||
if not training_runs:
|
||||
raise FileNotFoundError(f"No training runs found for environment '{env_name}'")
|
||||
|
||||
# Sort by modification time to get the most recent
|
||||
latest_run = max(training_runs, key=lambda x: x.stat().st_mtime)
|
||||
checkpoints_dir = latest_run / "checkpoints"
|
||||
|
||||
if not checkpoints_dir.exists():
|
||||
raise FileNotFoundError(f"No checkpoints directory found in {latest_run}")
|
||||
|
||||
# First, try to find best_agent files (highest performance models)
|
||||
best_files = list(checkpoints_dir.glob("best_agent.*"))
|
||||
|
||||
if best_files:
|
||||
# Return the first best_agent file found (there should only be one)
|
||||
return str(best_files[0])
|
||||
|
||||
# If no best_agent files, find the checkpoint with the highest timestep
|
||||
checkpoint_files = list(checkpoints_dir.glob("agent_*.pt")) + list(checkpoints_dir.glob("agent_*.pickle"))
|
||||
|
||||
if not checkpoint_files:
|
||||
raise FileNotFoundError(f"No policy files found in {checkpoints_dir}")
|
||||
|
||||
# Extract timestep from filename and find the highest
|
||||
def extract_timestep(filename):
|
||||
# Pattern: agent_{timestep}.ext
|
||||
stem = Path(filename).stem # agent_{timestep}
|
||||
parts = stem.split("_")
|
||||
if len(parts) >= 2:
|
||||
try:
|
||||
return int(parts[1])
|
||||
except ValueError:
|
||||
return 0
|
||||
return 0
|
||||
|
||||
latest_checkpoint = max(checkpoint_files, key=extract_timestep)
|
||||
return str(latest_checkpoint)
|
||||
|
||||
|
||||
def main(argv):
|
||||
device_supports = utils.get_device_supports()
|
||||
logger.info(device_supports)
|
||||
env_name = _ENV.value
|
||||
enable_render = True
|
||||
|
||||
rl_override = {}
|
||||
|
||||
if _NUM_ENVS.present:
|
||||
rl_override["play_num_envs"] = _NUM_ENVS.value
|
||||
|
||||
if _RAND_SEED.value:
|
||||
rl_override["seed"] = None
|
||||
elif _SEED.present:
|
||||
rl_override["seed"] = _SEED.value
|
||||
|
||||
sim_backend = _SIM_BACKEND.value
|
||||
|
||||
# Determine policy path: use explicit policy if provided, otherwise auto-discover
|
||||
if _POLICY.present:
|
||||
policy_path = _POLICY.value
|
||||
logger.info(f"Using specified policy: {policy_path}")
|
||||
else:
|
||||
try:
|
||||
policy_path = find_best_policy(env_name)
|
||||
logger.info(f"Auto-discovered best policy: {policy_path}")
|
||||
except FileNotFoundError as e:
|
||||
logger.error(f"Error: {e}")
|
||||
logger.error("Please specify a policy using --policy flag or train a model first")
|
||||
return
|
||||
|
||||
backend = get_inference_backend(policy_path)
|
||||
|
||||
if backend == "jax":
|
||||
assert device_supports.jax, "jax is not avaliable on your device "
|
||||
from motrix_rl.skrl.jax.train import ppo
|
||||
|
||||
config.jax.backend = "jax" # or "numpy"
|
||||
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
|
||||
trainer.play(policy_path)
|
||||
|
||||
elif backend == "torch":
|
||||
assert device_supports.torch, "torch is not avaliable on your device"
|
||||
from motrix_rl.skrl.torch.train import ppo
|
||||
|
||||
config.torch.backend = "torch"
|
||||
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
|
||||
trainer.play(policy_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app.run(main)
|
||||
94
scripts/train.py
Normal file
94
scripts/train.py
Normal file
@@ -0,0 +1,94 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
|
||||
|
||||
import logging
|
||||
|
||||
from absl import app, flags
|
||||
from skrl import config
|
||||
|
||||
from motrix_rl import utils
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_ENV = flags.DEFINE_string("env", "cartpole", "The env to train")
|
||||
_SIM_BACKEND = flags.DEFINE_string(
|
||||
"sim-backend",
|
||||
None,
|
||||
"The simulation backend to use.(If not specified, it will be choosen automatically)",
|
||||
)
|
||||
_NUM_ENVS = flags.DEFINE_integer("num-envs", 2048, "Number of envs to train")
|
||||
_RENDER = flags.DEFINE_bool("render", False, "Render the env")
|
||||
_TRAIN_BACKEND = flags.DEFINE_string("train-backend", "jax", "The learning backend. (jax/torch)")
|
||||
_SEED = flags.DEFINE_integer("seed", None, "Random seed for reproducibility")
|
||||
_RAND_SEED = flags.DEFINE_bool("rand-seed", False, "Generate random seed")
|
||||
|
||||
|
||||
def get_train_backend(supports: utils.DeviceSupports):
|
||||
if supports.jax and supports.jax_gpu:
|
||||
return "jax"
|
||||
elif supports.torch and supports.torch_gpu:
|
||||
return "torch"
|
||||
elif supports.jax:
|
||||
return "jax"
|
||||
elif supports.torch:
|
||||
return "torch"
|
||||
else:
|
||||
raise Exception("neither jax nor torch not avaliable on the device.")
|
||||
|
||||
|
||||
def main(argv):
|
||||
device_supports = utils.get_device_supports()
|
||||
logger.info(device_supports)
|
||||
env_name = _ENV.value
|
||||
enable_render = _RENDER.value
|
||||
|
||||
rl_override = {}
|
||||
|
||||
if _NUM_ENVS.present:
|
||||
rl_override["num_envs"] = _NUM_ENVS.value
|
||||
|
||||
if _RAND_SEED.value:
|
||||
rl_override["seed"] = None
|
||||
elif _SEED.present:
|
||||
rl_override["seed"] = _SEED.value
|
||||
|
||||
sim_backend = _SIM_BACKEND.value
|
||||
train_backend = "jax"
|
||||
if not _TRAIN_BACKEND.present:
|
||||
train_backend = get_train_backend(device_supports)
|
||||
else:
|
||||
train_backend = _TRAIN_BACKEND.value
|
||||
|
||||
trainer = None
|
||||
if train_backend == "jax":
|
||||
from motrix_rl.skrl.jax.train import ppo
|
||||
|
||||
config.jax.backend = "jax" # or "numpy"
|
||||
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
|
||||
|
||||
elif train_backend == "torch":
|
||||
from motrix_rl.skrl.torch.train import ppo
|
||||
|
||||
config.torch.backend = "torch"
|
||||
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
|
||||
else:
|
||||
raise Exception(f"Unknown train backend: {train_backend}")
|
||||
|
||||
trainer.train()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app.run(main)
|
||||
78
scripts/view.py
Normal file
78
scripts/view.py
Normal file
@@ -0,0 +1,78 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
from absl import app, flags
|
||||
|
||||
from motrix_envs import registry
|
||||
from motrix_envs.np.env import NpEnv
|
||||
from motrix_envs.np.renderer import NpRenderer
|
||||
|
||||
_ENV = flags.DEFINE_string("env", "cartpole", "The env to view")
|
||||
_SIM_BACKEND = flags.DEFINE_string("sim-backend", None, "The simulation backend to use.")
|
||||
_NUM_ENVS = flags.DEFINE_integer("num-envs", 1, "Number of parallel environments.")
|
||||
|
||||
|
||||
class NpEnvRunner:
|
||||
_renderer: NpRenderer
|
||||
|
||||
def __init__(self, env: NpEnv):
|
||||
self._env = env
|
||||
self._renderer = NpRenderer(env)
|
||||
|
||||
def _sample_random_action(self):
|
||||
action_space = self._env.action_space
|
||||
if isinstance(action_space, gym.spaces.Box):
|
||||
size = (self._env.num_envs, *action_space.shape)
|
||||
return np.random.uniform(
|
||||
low=action_space.low,
|
||||
high=action_space.high,
|
||||
size=size,
|
||||
).astype(action_space.dtype)
|
||||
else:
|
||||
raise NotImplementedError("Only Box action space is supported")
|
||||
|
||||
def step(self):
|
||||
actions = self._sample_random_action()
|
||||
self._env.step(actions)
|
||||
|
||||
def start(self):
|
||||
import time
|
||||
|
||||
env_dt = self._env.cfg.ctrl_dt
|
||||
while True:
|
||||
t0 = time.monotonic()
|
||||
actions = self._sample_random_action()
|
||||
self._env.step(actions)
|
||||
self._renderer.render()
|
||||
real_dt = time.monotonic() - t0
|
||||
sleep_dt = env_dt - real_dt
|
||||
if sleep_dt > 0:
|
||||
time.sleep(sleep_dt)
|
||||
|
||||
|
||||
def main(argv):
|
||||
env_name = _ENV.value
|
||||
sim_backend = _SIM_BACKEND.value
|
||||
num_envs = _NUM_ENVS.value
|
||||
env = registry.make(env_name, sim_backend=sim_backend, num_envs=num_envs)
|
||||
|
||||
runner = NpEnvRunner(env)
|
||||
runner.start()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app.run(main)
|
||||
Reference in New Issue
Block a user