From ba664148b2950b17a2e6fa1cc236fd46fe5e2a2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Budzianowski?= Date: Thu, 5 Dec 2024 13:59:12 -0800 Subject: [PATCH] update links to gpr (#125) update links, remove unused scripts, fix ci/cd --- Makefile | 8 +- pyproject.toml | 16 +- sim/compare_h5.py | 119 ------- sim/envs/__init__.py | 8 +- sim/envs/humanoids/grp_config.py | 4 +- sim/h5_logger.py | 35 +- sim/model_export.py | 80 +---- sim/play.py | 90 +++-- sim/play_old.py | 15 +- sim/produce_sim_data.py | 3 +- sim/requirements-dev.txt | 6 +- sim/requirements.txt | 4 +- sim/resources/gpr/joints.py | 7 +- sim/scripts/calibration_mujoco.py | 6 +- sim/scripts/create_fixed_torso.py | 17 +- sim/scripts/create_mjcf.py | 397 ---------------------- sim/scripts/download_assets.sh | 19 +- sim/scripts/mjcf.py | 524 ------------------------------ sim/sim2sim.py | 66 ++-- sim/sim2sim_old.py | 39 +-- 20 files changed, 172 insertions(+), 1291 deletions(-) delete mode 100644 sim/compare_h5.py delete mode 100755 sim/scripts/create_mjcf.py delete mode 100755 sim/scripts/mjcf.py diff --git a/Makefile b/Makefile index 3c344268..9ded39e0 100755 --- a/Makefile +++ b/Makefile @@ -27,16 +27,16 @@ all: # ------------------------ # train-one-vis: - @python -m sim.train --task stompymini --num_envs 1 + @python -m sim.train --task gpr --num_envs 1 train-many-vis: - @python -m sim.train --task stompymini --num_envs 16 + @python -m sim.train --task gpr --num_envs 16 train: - @python -m sim.train --task stompymini --num_envs 4096 --headless + @python -m sim.train --task gpr --num_envs 4096 --headless play: - @python -m sim.play --task stompymini + @python -m sim.play --task gpr # ------------------------ # # Build # diff --git a/pyproject.toml b/pyproject.toml index eb7692ca..13ba2009 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [tool.black] line-length = 120 -target-version = ["py311"] +target-version = ["py38"] include = '\.pyi?$' [tool.pytest.ini_options] @@ -30,16 +30,21 @@ warn_redundant_casts = true incremental = true namespace_packages = false -exclude = ["sim/envs/", "sim/algo/*", "sim/utils/*", "sim/sim2sim.py", - "sim/scripts/create_mjcf.py", "sim/scripts/create_fixed_torso.py"] +exclude = ["sim/envs/", "sim/algo/", "sim/utils/", "sim/sim2sim.py", + "sim/scripts/create_mjcf.py", "sim/scripts/create_fixed_torso.py", + "sim/sim2sim_old.py", "sim/play_old.py"] [[tool.mypy.overrides]] module = [ "isaacgym.*", "mujoco.*", + "sim.envs.*", + "sim.utils.*", ] +ignore_errors = true + [tool.isort] profile = "black" @@ -48,10 +53,11 @@ known_third_party = ["wandb"] [tool.ruff] line-length = 120 -target-version = "py310" +target-version = "py38" exclude = ["sim/envs/", "sim/algo/", "sim/play.py", - "sim/sim2sim.py","sim/utils", "sim/scripts/create_mjcf.py"] + "sim/sim2sim.py","sim/utils", "sim/scripts/create_mjcf.py", + "sim/play_old.py", "sim/sim2sim_old.py"] [tool.ruff.lint] diff --git a/sim/compare_h5.py b/sim/compare_h5.py deleted file mode 100644 index 5da44b34..00000000 --- a/sim/compare_h5.py +++ /dev/null @@ -1,119 +0,0 @@ -""" -Compare two HDF5 files and analyze their differences. - -Usage: - python compare_h5.py --isaac-file path/to/isaac.h5 --mujoco-file path/to/mujoco.h5 - -Example: - python sim/compare_h5.py \ - --isaac-file runs/h5_out/stompypro/2024-12-02_20-04-51/env_0/stompypro_env_0/h5_out/2024-12-02_20-04-51__053b2b5b-21c9-497c-b637-b66935dfe475.h5 \ - --mujoco-file sim/resources/stompypro/h5_out/2024-12-02_21-10-41__5820ade7-9fc0-46df-8469-2f305480bcae.h5 - - python sim/compare_h5.py \ - --isaac-file runs/h5_out/stompypro/2024-12-02_20-04-51/env_1/stompypro_env_1/h5_out/2024-12-02_20-04-51__d2016a6f-a486-4e86-8c5e-c9addd2cc13e.h5 \ - --mujoco-file sim/resources/stompypro/h5_out/2024-12-02_21-10-41__5820ade7-9fc0-46df-8469-2f305480bcae.h5 - -""" - -import argparse -from pathlib import Path - -import h5py -import numpy as np - - -def load_h5_file(file_path): - """Load H5 file and return a dictionary of datasets""" - data = {} - with h5py.File(file_path, "r") as f: - # Recursively load all datasets - def load_group(group, prefix=""): - for key in group.keys(): - path = f"{prefix}/{key}" if prefix else key - if isinstance(group[key], h5py.Group): - load_group(group[key], path) - else: - data[path] = group[key][:] - - load_group(f) - return data - - -def compare_h5_files(issac_path, mujoco_path): - """Compare two H5 files and print differences""" - print(f"\nLoading files:") - print(f"Isaac: {issac_path}") - print(f"Mujoco: {mujoco_path}") - - # Load both files - issac_data = load_h5_file(issac_path) - mujoco_data = load_h5_file(mujoco_path) - - print("\nFile lengths:") - print(f"Isaac datasets: {len(issac_data)}") - print(f"Mujoco datasets: {len(mujoco_data)}") - - print("\nDataset shapes:") - print("\nIsaac shapes:") - for key, value in issac_data.items(): - print(f"{key}: {value.shape}") - - print("\nMujoco shapes:") - for key, value in mujoco_data.items(): - print(f"{key}: {value.shape}") - - # Find common keys - common_keys = set(issac_data.keys()) & set(mujoco_data.keys()) - print(f"\nCommon datasets: {len(common_keys)}") - - # Find uncommon keys - issac_only_keys = set(issac_data.keys()) - common_keys - mujoco_only_keys = set(mujoco_data.keys()) - common_keys - print(f"\nIsaac only datasets: {len(issac_only_keys)}") - for key in issac_only_keys: - print(f"Isaac only: {key}") - print(f"\nMujoco only datasets: {len(mujoco_only_keys)}") - for key in mujoco_only_keys: - print(f"Mujoco only: {key}") - - # Compare data for common keys - for key in common_keys: - issac_arr = issac_data[key] - mujoco_arr = mujoco_data[key] - - print(f"\n========== For {key} ===============") - - if issac_arr.shape != mujoco_arr.shape: - print(f"\n{key} - Shape mismatch: Isaac {issac_arr.shape} vs Mujoco {mujoco_arr.shape}") - - # Calculate differences - min_shape = min(issac_arr.shape[0], mujoco_arr.shape[0]) - if issac_arr.shape != mujoco_arr.shape: - raise ValueError(f"Shapes do not match for {key}. Cannot compare datasets with different shapes.") - diff = np.abs(issac_arr[:min_shape] - mujoco_arr[:min_shape]) - max_diff = np.max(diff) - mean_diff = np.mean(diff) - - print(f"Max difference: {max_diff:.6f}") - print(f"Mean difference: {mean_diff:.6f}\n") - - start_idx = 0 - display_timesteps = 10 - end_idx = start_idx + display_timesteps - - np.set_printoptions(formatter={"float": "{:0.6f}".format}, suppress=True) - print("Isaac:\n", issac_arr[start_idx:end_idx]) - print("Mujoco:\n", mujoco_arr[start_idx:end_idx]) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Compare two H5 files from Isaac and Mujoco simulations") - parser.add_argument("--isaac-file", required=True, help="Path to Isaac simulation H5 file") - parser.add_argument("--mujoco-file", required=True, help="Path to Mujoco simulation H5 file") - - args = parser.parse_args() - - print(f"Isaac path: {args.isaac_file}") - print(f"Mujoco path: {args.mujoco_file}") - - compare_h5_files(args.isaac_file, args.mujoco_file) diff --git a/sim/envs/__init__.py b/sim/envs/__init__.py index ecd72198..639853fc 100755 --- a/sim/envs/__init__.py +++ b/sim/envs/__init__.py @@ -12,16 +12,12 @@ from sim.envs.humanoids.dora_env import DoraFreeEnv from sim.envs.humanoids.g1_config import G1Cfg, G1CfgPPO from sim.envs.humanoids.g1_env import G1FreeEnv +from sim.envs.humanoids.gpr_env import GprFreeEnv +from sim.envs.humanoids.grp_config import GprCfg, GprCfgPPO, GprStandingCfg from sim.envs.humanoids.h1_config import H1Cfg, H1CfgPPO from sim.envs.humanoids.h1_env import H1FreeEnv from sim.envs.humanoids.stompymicro_config import StompyMicroCfg, StompyMicroCfgPPO from sim.envs.humanoids.stompymicro_env import StompyMicroEnv -from sim.envs.humanoids.grp_config import ( - GprCfg, - GprCfgPPO, - GprStandingCfg, -) -from sim.envs.humanoids.gpr_env import GprFreeEnv from sim.envs.humanoids.xbot_config import XBotCfg, XBotCfgPPO from sim.envs.humanoids.xbot_env import XBotLFreeEnv from sim.utils.task_registry import TaskRegistry # noqa: E402 diff --git a/sim/envs/humanoids/grp_config.py b/sim/envs/humanoids/grp_config.py index 5ab53816..35a643dc 100644 --- a/sim/envs/humanoids/grp_config.py +++ b/sim/envs/humanoids/grp_config.py @@ -152,7 +152,7 @@ class rewards: target_joint_pos_scale = 0.17 # rad target_feet_height = 0.06 # m - cycle_time = 0.25 # sec + cycle_time = 0.64 # sec # if true negative total rewards are clipped at zero (avoids early termination problems) only_positive_rewards = True # tracking reward = exp(error*sigma) @@ -180,7 +180,7 @@ class scales: # base pos default_joint_pos = 0.5 - orientation = 1. + orientation = 1.0 base_height = 0.2 base_acc = 0.2 # energy diff --git a/sim/h5_logger.py b/sim/h5_logger.py index 69f46271..8e96a554 100644 --- a/sim/h5_logger.py +++ b/sim/h5_logger.py @@ -1,8 +1,9 @@ -""" Logger for logging data to HDF5 files """ +"""Logger for logging data to HDF5 files.""" + import os import uuid from datetime import datetime -from typing import Dict +from typing import Dict, Tuple import h5py import matplotlib.pyplot as plt # dependency issues with python 3.8 @@ -17,7 +18,7 @@ def __init__( max_timesteps: int, num_observations: int, h5_out_dir: str = "sim/resources/", - ): + ) -> None: self.data_name = data_name self.num_actions = num_actions self.max_timesteps = max_timesteps @@ -27,7 +28,7 @@ def __init__( self.h5_file, self.h5_dict = self._create_h5_file() self.current_timestep = 0 - def _create_h5_file(self): + def _create_h5_file(self) -> Tuple[h5py.File, Dict[str, h5py.Dataset]]: # Create a unique file ID idd = str(uuid.uuid4()) timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") @@ -43,8 +44,8 @@ def _create_h5_file(self): dset_prev_actions = h5_file.create_dataset( "prev_actions", (self.max_timesteps, self.num_actions), dtype=np.float32 ) - dset_2D_command = h5_file.create_dataset("observations/2D_command", (self.max_timesteps, 2), dtype=np.float32) - dset_3D_command = h5_file.create_dataset("observations/3D_command", (self.max_timesteps, 3), dtype=np.float32) + dset_2d_command = h5_file.create_dataset("observations/2D_command", (self.max_timesteps, 2), dtype=np.float32) + dset_3d_command = h5_file.create_dataset("observations/3D_command", (self.max_timesteps, 3), dtype=np.float32) dset_q = h5_file.create_dataset("observations/q", (self.max_timesteps, self.num_actions), dtype=np.float32) dset_dq = h5_file.create_dataset("observations/dq", (self.max_timesteps, self.num_actions), dtype=np.float32) dset_ang_vel = h5_file.create_dataset("observations/ang_vel", (self.max_timesteps, 3), dtype=np.float32) @@ -61,8 +62,8 @@ def _create_h5_file(self): h5_dict = { "prev_actions": dset_prev_actions, "curr_actions": dset_curr_actions, - "2D_command": dset_2D_command, - "3D_command": dset_3D_command, + "2D_command": dset_2d_command, + "3D_command": dset_3d_command, "joint_pos": dset_q, "joint_vel": dset_dq, "ang_vel": dset_ang_vel, @@ -72,7 +73,7 @@ def _create_h5_file(self): } return h5_file, h5_dict - def log_data(self, data: Dict[str, np.ndarray]): + def log_data(self, data: Dict[str, np.ndarray]) -> None: if self.current_timestep >= self.max_timesteps: print(f"Warning: Exceeded maximum timesteps ({self.max_timesteps})") return @@ -83,7 +84,7 @@ def log_data(self, data: Dict[str, np.ndarray]): self.current_timestep += 1 - def close(self): + def close(self) -> None: for key, dataset in self.h5_dict.items(): max_val = np.max(np.abs(dataset[:])) if max_val > self.max_threshold: @@ -97,9 +98,8 @@ def close(self): self.h5_file.close() @staticmethod - def visualize_h5(h5_file_path: str): - """ - Visualizes the data from an HDF5 file by plotting each variable one by one. + def visualize_h5(h5_file_path: str) -> None: + """Visualizes the data from an HDF5 file by plotting each variable one by one. Args: h5_file_path (str): Path to the HDF5 file. @@ -122,9 +122,8 @@ def visualize_h5(h5_file_path: str): print(f"Failed to visualize HDF5 file: {e}") @staticmethod - def _plot_dataset(name: str, data: np.ndarray): - """ - Helper method to plot a single dataset. + def _plot_dataset(name: str, data: np.ndarray) -> None: + """Helper method to plot a single dataset. Args: name (str): Name of the dataset. @@ -144,7 +143,3 @@ def _plot_dataset(name: str, data: np.ndarray): plt.grid(True) plt.tight_layout() plt.show() - - -if __name__ == "__main__": - HDF5Logger.visualize_h5("stompypro/5a7dc371-445c-4f56-be05-4e65c5cc38bc.h5") diff --git a/sim/model_export.py b/sim/model_export.py index c5184817..09c190a4 100644 --- a/sim/model_export.py +++ b/sim/model_export.py @@ -8,7 +8,7 @@ import onnx import onnxruntime as ort import torch -from scripts.create_mjcf import load_embodiment +from scripts.create_fixed_torso import load_embodiment from torch import Tensor, nn from torch.distributions import Normal @@ -47,25 +47,25 @@ def __init__( # Policy function. actor_layers = [] actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0])) - actor_layers.append(activation) + actor_layers.append(activation) # type: ignore[arg-type] for dim_i in range(len(actor_hidden_dims)): if dim_i == len(actor_hidden_dims) - 1: actor_layers.append(nn.Linear(actor_hidden_dims[dim_i], num_actions)) else: actor_layers.append(nn.Linear(actor_hidden_dims[dim_i], actor_hidden_dims[dim_i + 1])) - actor_layers.append(activation) + actor_layers.append(activation) # type: ignore[arg-type] self.actor = nn.Sequential(*actor_layers) # Value function. critic_layers = [] critic_layers.append(nn.Linear(mlp_input_dim_c, critic_hidden_dims[0])) - critic_layers.append(activation) + critic_layers.append(activation) # type: ignore[arg-type] for dim_i in range(len(critic_hidden_dims)): if dim_i == len(critic_hidden_dims) - 1: critic_layers.append(nn.Linear(critic_hidden_dims[dim_i], 1)) else: critic_layers.append(nn.Linear(critic_hidden_dims[dim_i], critic_hidden_dims[dim_i + 1])) - critic_layers.append(activation) + critic_layers.append(activation) # type: ignore[arg-type] self.critic = nn.Sequential(*critic_layers) # Action noise. @@ -73,7 +73,7 @@ def __init__( self.distribution = None # Disable args validation for speedup. - Normal.set_default_validate_args = False + Normal.set_default_validate_args = False # type: ignore[unused-ignore,assignment,method-assign,attr-defined] class Actor(nn.Module): @@ -205,52 +205,6 @@ def forward( return actions_scaled, actions, x -def get_actor_policy(model_path: str, cfg: ActorCfg) -> Tuple[nn.Module, dict, Tuple[Tensor, ...]]: - all_weights = torch.load(model_path, map_location="cpu", weights_only=True) - weights = all_weights["model_state_dict"] - num_actor_obs = weights["actor.0.weight"].shape[1] - num_critic_obs = weights["critic.0.weight"].shape[1] - num_actions = weights["std"].shape[0] - actor_hidden_dims = [v.shape[0] for k, v in weights.items() if re.match(r"actor\.\d+\.weight", k)] - critic_hidden_dims = [v.shape[0] for k, v in weights.items() if re.match(r"critic\.\d+\.weight", k)] - actor_hidden_dims = actor_hidden_dims[:-1] - critic_hidden_dims = critic_hidden_dims[:-1] - - ac_model = ActorCritic(num_actor_obs, num_critic_obs, num_actions, actor_hidden_dims, critic_hidden_dims) - ac_model.load_state_dict(weights) - - a_model = Actor(ac_model.actor, cfg) - - # Gets the model input tensors. - x_vel = torch.randn(1) - y_vel = torch.randn(1) - rot = torch.randn(1) - t = torch.randn(1) - dof_pos = torch.randn(a_model.num_actions) - dof_vel = torch.randn(a_model.num_actions) - prev_actions = torch.randn(a_model.num_actions) - imu_ang_vel = torch.randn(3) - imu_euler_xyz = torch.randn(3) - buffer = a_model.get_init_buffer() - input_tensors = (x_vel, y_vel, rot, t, dof_pos, dof_vel, prev_actions, imu_ang_vel, imu_euler_xyz, buffer) - - jit_model = torch.jit.script(a_model) - - # Add sim2sim metadata - robot_effort = list(a_model.robot.effort().values()) - robot_stiffness = list(a_model.robot.stiffness().values()) - robot_damping = list(a_model.robot.damping().values()) - num_actions = a_model.num_actions - num_observations = a_model.num_observations - - return a_model, { - "robot_effort": robot_effort, - "robot_stiffness": robot_stiffness, - "robot_damping": robot_damping, - "num_actions": num_actions, - "num_observations": num_observations, - }, input_tensors - def get_actor_policy(model_path: str, cfg: ActorCfg) -> Tuple[nn.Module, dict, Tuple[Tensor, ...]]: all_weights = torch.load(model_path, map_location="cpu", weights_only=True) @@ -281,24 +235,24 @@ def get_actor_policy(model_path: str, cfg: ActorCfg) -> Tuple[nn.Module, dict, T buffer = a_model.get_init_buffer() input_tensors = (x_vel, y_vel, rot, t, dof_pos, dof_vel, prev_actions, imu_ang_vel, imu_euler_xyz, buffer) - jit_model = torch.jit.script(a_model) - # Add sim2sim metadata robot_effort = list(a_model.robot.effort().values()) robot_stiffness = list(a_model.robot.stiffness().values()) robot_damping = list(a_model.robot.damping().values()) - default_standing = list(a_model.robot.default_standing().values()) num_actions = a_model.num_actions num_observations = a_model.num_observations - return a_model, { - "robot_effort": robot_effort, - "robot_stiffness": robot_stiffness, - "robot_damping": robot_damping, - "default_standing": default_standing, - "num_actions": num_actions, - "num_observations": num_observations, - }, input_tensors + return ( + a_model, + { + "robot_effort": robot_effort, + "robot_stiffness": robot_stiffness, + "robot_damping": robot_damping, + "num_actions": num_actions, + "num_observations": num_observations, + }, + input_tensors, + ) def convert_model_to_onnx(model_path: str, cfg: ActorCfg, save_path: Optional[str] = None) -> ort.InferenceSession: diff --git a/sim/play.py b/sim/play.py index 4881d6ba..9abe5116 100755 --- a/sim/play.py +++ b/sim/play.py @@ -18,23 +18,22 @@ import h5py import numpy as np from isaacgym import gymapi +from kinfer.export.pytorch import export_to_onnx from tqdm import tqdm # Local imports third -from sim.env import run_dir -from sim.envs import task_registry -from sim.h5_logger import HDF5Logger -from sim.model_export import ActorCfg, convert_model_to_onnx -from sim.utils.helpers import get_args -from sim.utils.logger import Logger - -import torch # special case with isort: skip comment from sim.env import run_dir # noqa: E402 from sim.envs import task_registry # noqa: E402 -from sim.model_export import ActorCfg, get_actor_policy # noqa: E402 +from sim.h5_logger import HDF5Logger +from sim.model_export import ( # noqa: E402 + ActorCfg, + convert_model_to_onnx, + get_actor_policy, +) from sim.utils.helpers import get_args # noqa: E402 from sim.utils.logger import Logger # noqa: E402 -from kinfer.export.pytorch import export_to_onnx + +import torch # special case with isort: skip comment logger = logging.getLogger(__name__) @@ -92,7 +91,7 @@ def play(args: argparse.Namespace) -> None: # export policy as a onnx module (used to run it on web) if args.export_onnx: path = ppo_runner.load_path - embodiment = ppo_runner.cfg['experiment_name'].lower() + embodiment = ppo_runner.cfg["experiment_name"].lower() policy_cfg = ActorCfg(embodiment=embodiment) if embodiment == "stompypro": @@ -101,18 +100,13 @@ def play(args: argparse.Namespace) -> None: policy_cfg.cycle_time = 0.2 else: print(f"Specific policy cfg for {embodiment} not implemented") - + actor_model, sim2sim_info, input_tensors = get_actor_policy(path, policy_cfg) # Merge policy_cfg and sim2sim_info into a single config object export_config = {**vars(policy_cfg), **sim2sim_info} - export_to_onnx( - actor_model, - input_tensors=input_tensors, - config=export_config, - save_path="kinfer_policy.onnx" - ) + export_to_onnx(actor_model, input_tensors=input_tensors, config=export_config, save_path="kinfer_policy.onnx") print("Exported policy as kinfer-compatible onnx to: ", path) # Prepare for logging @@ -126,7 +120,7 @@ def play(args: argparse.Namespace) -> None: # Create directory for HDF5 files h5_dir = run_dir() / "h5_out" / args.task / now h5_dir.mkdir(parents=True, exist_ok=True) - + # Get observation dimensions num_actions = env.num_dof obs_buffer = env.obs_buf.shape[1] @@ -136,14 +130,16 @@ def play(args: argparse.Namespace) -> None: for env_idx in range(env_cfg.env.num_envs): h5_dir = run_dir() / "h5_out" / args.task / now / f"env_{env_idx}" h5_dir.mkdir(parents=True, exist_ok=True) - - h5_loggers.append(HDF5Logger( - data_name=f"{args.task}_env_{env_idx}", - num_actions=num_actions, - max_timesteps=env_steps_to_run, - num_observations=obs_buffer, - h5_out_dir=str(h5_dir) - )) + + h5_loggers.append( + HDF5Logger( + data_name=f"{args.task}_env_{env_idx}", + num_actions=num_actions, + max_timesteps=env_steps_to_run, + num_observations=obs_buffer, + h5_out_dir=str(h5_dir), + ) + ) if args.render: camera_properties = gymapi.CameraProperties() @@ -228,27 +224,29 @@ def play(args: argparse.Namespace) -> None: if args.log_h5: # Extract the current observation for env_idx in range(env_cfg.env.num_envs): - h5_loggers[env_idx].log_data({ - "t": np.array([t * env.dt], dtype=np.float32), - "2D_command": np.array( - [ - np.sin(2 * math.pi * t * env.dt / env.cfg.rewards.cycle_time), - np.cos(2 * math.pi * t * env.dt / env.cfg.rewards.cycle_time), - ], - dtype=np.float32, - ), - "3D_command": np.array(env.commands[env_idx, :3].cpu().numpy(), dtype=np.float32), - "joint_pos": np.array(env.dof_pos[env_idx].cpu().numpy(), dtype=np.float32), - "joint_vel": np.array(env.dof_vel[env_idx].cpu().numpy(), dtype=np.float32), - "prev_actions": prev_actions[env_idx].astype(np.float32), - "curr_actions": actions[env_idx].astype(np.float32), - "ang_vel": env.base_ang_vel[env_idx].cpu().numpy().astype(np.float32), - "euler_rotation": env.base_euler_xyz[env_idx].cpu().numpy().astype(np.float32), - "buffer": env.obs_buf[env_idx].cpu().numpy().astype(np.float32) - }) + h5_loggers[env_idx].log_data( + { + "t": np.array([t * env.dt], dtype=np.float32), + "2D_command": np.array( + [ + np.sin(2 * math.pi * t * env.dt / env.cfg.rewards.cycle_time), + np.cos(2 * math.pi * t * env.dt / env.cfg.rewards.cycle_time), + ], + dtype=np.float32, + ), + "3D_command": np.array(env.commands[env_idx, :3].cpu().numpy(), dtype=np.float32), + "joint_pos": np.array(env.dof_pos[env_idx].cpu().numpy(), dtype=np.float32), + "joint_vel": np.array(env.dof_vel[env_idx].cpu().numpy(), dtype=np.float32), + "prev_actions": prev_actions[env_idx].astype(np.float32), + "curr_actions": actions[env_idx].astype(np.float32), + "ang_vel": env.base_ang_vel[env_idx].cpu().numpy().astype(np.float32), + "euler_rotation": env.base_euler_xyz[env_idx].cpu().numpy().astype(np.float32), + "buffer": env.obs_buf[env_idx].cpu().numpy().astype(np.float32), + } + ) prev_actions = actions - + if infos["episode"]: num_episodes = env.reset_buf.sum().item() if num_episodes > 0: diff --git a/sim/play_old.py b/sim/play_old.py index d842fb83..6747abb4 100755 --- a/sim/play_old.py +++ b/sim/play_old.py @@ -5,6 +5,7 @@ python sim/play.py --task g1 --log_h5 python sim/play.py --task stompymini --log_h5 """ + import argparse import copy import logging @@ -22,7 +23,7 @@ from sim.env import run_dir # noqa: E402 from sim.envs import task_registry # noqa: E402 -from sim.model_export import ActorCfg, convert_model_to_onnx # noqa: E402 +from sim.model_export import convert_model_to_onnx # noqa: E402 from sim.utils.helpers import get_args # noqa: E402 from sim.utils.logger import Logger # noqa: E402 @@ -72,18 +73,6 @@ def play(args: argparse.Namespace) -> None: ppo_runner, train_cfg = task_registry.make_alg_runner(env=env, name=args.task, args=args, train_cfg=train_cfg) policy = ppo_runner.get_inference_policy(device=env.device) - # # Export policy if needed - # if args.export_policy: - # path = os.path.join(".") - # export_policy_as_jit(ppo_runner.alg.actor_critic, path) - # print("Exported policy as jit script to: ", path) - - # export policy as a onnx module (used to run it on web) - # if args.export_onnx: - # path = ppo_runner.alg.actor_critic - # convert_model_to_onnx(path, ActorCfg(), save_path="policy.onnx") - # print("Exported policy as onnx to: ", path) - # Prepare for logging env_logger = Logger(env.dt) robot_index = 0 diff --git a/sim/produce_sim_data.py b/sim/produce_sim_data.py index a85f8f33..a9fd0f55 100644 --- a/sim/produce_sim_data.py +++ b/sim/produce_sim_data.py @@ -1,4 +1,5 @@ -""" Produce simulation data for training.""" +"""Produce simulation data for training.""" + import argparse import multiprocessing as mp import subprocess diff --git a/sim/requirements-dev.txt b/sim/requirements-dev.txt index d2e78d59..c55c16ba 100755 --- a/sim/requirements-dev.txt +++ b/sim/requirements-dev.txt @@ -6,11 +6,10 @@ h5py==3.11.0 imageio==2.35.1 matplotlib==3.3.4 mediapy==1.2.2 -mujoco==3.2.3 +mujoco==2.3.6 mujoco_python_viewer==0.1.4 -numpy<2.0.0 onnx==1.15.0 -onnxruntime==1.19.2 +onnxruntime pandas==1.4.4 Pillow>6.2.0 poselib==2.0.4 @@ -19,7 +18,6 @@ python-dotenv==1.0.1 ruamel.base==1.0.0 scipy>1.10.0 setuptools==75.1.0 -torch==2.4.1 tqdm==4.67.0 wandb==0.18.7 mypy==1.10.0 diff --git a/sim/requirements.txt b/sim/requirements.txt index ff8a57a2..ab18cb34 100755 --- a/sim/requirements.txt +++ b/sim/requirements.txt @@ -6,11 +6,9 @@ python-dotenv h5py gdown matplotlib==3.3.4 -numpy<2.0.0 +numpy==1.20.0 wandb tensorboard==2.14.0 onnxscript mujoco==2.3.6 -# onnxruntime - kinfer diff --git a/sim/resources/gpr/joints.py b/sim/resources/gpr/joints.py index d1f966eb..ccfb301a 100755 --- a/sim/resources/gpr/joints.py +++ b/sim/resources/gpr/joints.py @@ -128,7 +128,7 @@ def isaac_to_mujoco_signs(cls) -> Dict[str, int]: def default_positions(cls) -> Dict[str, float]: return {} - # CONTRACT - this should be ordered according to how the policy is trained. + # CONTRACT - this should be ordered according to how the policy is trained. # E.g. the first entry should be the angle of the first joint in the policy. @classmethod def default_standing(cls) -> Dict[str, float]: @@ -149,7 +149,7 @@ def default_standing(cls) -> Dict[str, float]: def default_limits(cls) -> Dict[str, Dict[str, float]]: return { Robot.legs.left.knee_pitch: {"lower": -1.57, "upper": 0}, - Robot.legs.right.knee_pitch: {"lower": 0, "upper": 1.57} + Robot.legs.right.knee_pitch: {"lower": 0, "upper": 1.57}, } # p_gains @@ -184,6 +184,7 @@ def effort(cls) -> Dict[str, float]: "knee": 120, "ankle_y": 50, } + # # vel_limits @classmethod def velocity(cls) -> Dict[str, float]: @@ -207,4 +208,4 @@ def print_joints() -> None: if __name__ == "__main__": # python -m sim.Robot.joints - print_joints() \ No newline at end of file + print_joints() diff --git a/sim/scripts/calibration_mujoco.py b/sim/scripts/calibration_mujoco.py index 7045d44b..0d2447d2 100755 --- a/sim/scripts/calibration_mujoco.py +++ b/sim/scripts/calibration_mujoco.py @@ -10,13 +10,13 @@ from copy import deepcopy from typing import Any -import matplotlib.pyplot as plt # Add this import for plotting +import matplotlib.pyplot as plt import mujoco import mujoco_viewer import numpy as np from scipy.spatial.transform import Rotation as R -from sim.scripts.create_mjcf import load_embodiment +from sim.scripts.create_fixed_torso import load_embodiment import torch # isort: skip @@ -111,7 +111,7 @@ def run_mujoco(cfg: Any, joint_id: int = 0, steps: int = 1000) -> None: # noqa: position_data = [] while step < steps: - q, dq, _, _, _, _ = get_obs(data) + q, dq, _, _, _, _ = get_obs(data) # type: ignore[misc] q = q[-cfg.num_actions :] dq = dq[-cfg.num_actions :] diff --git a/sim/scripts/create_fixed_torso.py b/sim/scripts/create_fixed_torso.py index 0a1cf123..f93a4960 100755 --- a/sim/scripts/create_fixed_torso.py +++ b/sim/scripts/create_fixed_torso.py @@ -2,10 +2,19 @@ """This script updates the URDF file to fix the joints of the robot.""" import argparse +import importlib import xml.etree.ElementTree as ET from pathlib import Path +from typing import Any -from sim.scripts.create_mjcf import load_embodiment + +def load_embodiment(embodiment: str) -> Any: # noqa: ANN401 + # Dynamically import embodiment based on MODEL_DIR + module_name = f"sim.resources.{embodiment}.joints" + module = importlib.import_module(module_name) + robot = getattr(module, "Robot") + print(robot) + return robot def update_urdf(model_path: str, embodiment: str) -> None: @@ -35,15 +44,15 @@ def update_urdf(model_path: str, embodiment: str) -> None: limit.set("lower", lower) limit.set("upper", upper) for key, value in effort.items(): - if key in joint_name: + if key in joint_name: # type: ignore[operator] limit.set("effort", str(value)) for key, value in velocity.items(): - if key in joint_name: + if key in joint_name: # type: ignore[operator] limit.set("velocity", str(value)) dynamics = joint.find("dynamics") if dynamics is not None: for key, value in friction.items(): - if key in joint_name: + if key in joint_name: # type: ignore[operator] dynamics.set("friction", str(value)) # Save the modified URDF to a new file diff --git a/sim/scripts/create_mjcf.py b/sim/scripts/create_mjcf.py deleted file mode 100755 index 123dcbbd..00000000 --- a/sim/scripts/create_mjcf.py +++ /dev/null @@ -1,397 +0,0 @@ -# mypy: disable-error-code="operator,union-attr" -"""Defines common types and functions for exporting MJCF files. - -Run: - python sim/scripts/create_mjcf.py /path/to/robot.xml - -Todo: - 1. Add IMU to the right position - 2. Armature damping setup for different parts of body -""" - -import argparse -import importlib -import logging -import os -import xml.dom.minidom -import xml.etree.ElementTree as ET -from pathlib import Path -from typing import Any, List, OrderedDict, Union - -from sim.scripts import mjcf - -logger = logging.getLogger(__name__) - -DAMPING_DEFAULT = 0.02 - - -def load_embodiment(embodiment: str) -> Any: - # Dynamically import embodiment based on MODEL_DIR - module_name = f"sim.resources.{embodiment}.joints" - module = importlib.import_module(module_name) - robot = getattr(module, "Robot") - print(robot) - return robot - - -def load_config() -> Any: - # Dynamically import config based on MODEL_DIR - model_dir = os.environ.get("MODEL_DIR", "stompypro") - if "sim/" in model_dir: - model_dir = model_dir.split("sim/")[1] - model_dir = "stompypro" - module_name = f"sim.{model_dir}.config" - module = importlib.import_module(module_name) - config = getattr(module, "Config") - return config - - -def _pretty_print_xml(xml_string: str) -> str: - """Formats the provided XML string into a pretty-printed version.""" - parsed_xml = xml.dom.minidom.parseString(xml_string) - pretty_xml = parsed_xml.toprettyxml(indent=" ") - - # Split the pretty-printed XML into lines and filter out empty lines - lines = pretty_xml.split("\n") - non_empty_lines = [line for line in lines if line.strip() != ""] - # Remove declaration - return "\n".join(non_empty_lines[1:]) - - -class Sim2SimRobot(mjcf.Robot): - """A class to adapt the world in a Mujoco XML file.""" - - def update_joints(self, root: ET.Element, damping: float = DAMPING_DEFAULT) -> ET.Element: - joint_limits = robot.default_limits() - default_standing = robot.default_standing() - - for body in root.findall(".//body"): - joints_to_remove = [] - for joint in body.findall("joint"): - joint_name = joint.get("name") - if joint_name in joint_limits: - limits = joint_limits.get(joint_name) - lower = str(limits.get("lower", 0.0)) - upper = str(limits.get("upper", 0.0)) - joint.set("range", f"{lower} {upper}") - - # Comment to use Mujoco defaults - # keys = robot.damping().keys() - # for key in keys: - # if key in joint_name: - # joint_damping = damping - # joint.set("damping", str(joint_damping)) - # print(f"Damping for {joint_name}: {joint_damping}") - - # keys = robot.stiffness().keys() - # for key in keys: - # if key in joint_name: - # stiffness = robot.stiffness()[key] - # stiffness = 0.1 - # joint.set("stiffness", str(stiffness)) - # print(f"Stiffness for {joint_name}: {stiffness}") - - # Check if the joint is not in default_standing - if joint_name not in default_standing: - joints_to_remove.append(joint) - - # Remove joints not in default_standing - for joint in joints_to_remove: - body.remove(joint) - print(f"Removed joint: {joint.get('name')}") - - return root - - def swap_bodies(self, root: ET.Element, body1_name: str, body2_name: str) -> None: - """Swap the positions of two bodies in the XML tree.""" - parent_body = root.find(".//body[@name='root']") - if parent_body is not None: - body1 = parent_body.find(f".//body[@name='{body1_name}']") - body2 = parent_body.find(f".//body[@name='{body2_name}']") - if body1 is not None and body2 is not None: - body1_index = list(parent_body).index(body1) - body2_index = list(parent_body).index(body2) - # Swap the bodies - parent_body[body1_index], parent_body[body2_index] = parent_body[body2_index], parent_body[body1_index] - print(f"Swapped bodies: {body1_name} and {body2_name}") - else: - print(f"One or both bodies not found: {body1_name}, {body2_name}") - else: - print("Root body not found") - - def adapt_world(self, add_floor: bool = True, remove_frc_range: bool = True) -> None: - root: ET.Element = self.tree.getroot() - - worldbody = root.find("worldbody") - new_root_body = mjcf.Body(name="root", pos=(0, 0, 1), quat=(0, 0, 0, 1)).to_xml() - # add freejoint to root - freejoint = ET.Element("freejoint", name="root") - new_root_body.insert(0, freejoint) - items_to_move = [] - # Gather all children (geoms and bodies) that need to be moved under the new root body - for element in worldbody: - items_to_move.append(element) - # Move gathered elements to the new root body - for item in items_to_move: - worldbody.remove(item) - new_root_body.append(item) - # Add the new root body to the worldbody - worldbody.append(new_root_body) - - if add_floor: - asset = root.find("asset") - asset.append( - ET.Element( - "texture", - name="texplane", - type="2d", - builtin="checker", - rgb1=".0 .0 .0", - rgb2=".8 .8 .8", - width="100", - height="108", - ) - ) - asset.append( - ET.Element( - "material", - name="matplane", - reflectance="0.", - texture="texplane", - texrepeat="1 1", - texuniform="true", - ) - ) - asset.append(ET.Element("material", name="visualgeom", rgba="0.5 0.9 0.2 1")) - - compiler = root.find("compiler") - if self.compiler is not None: - compiler = self.compiler.to_xml(compiler) - - worldbody.insert( - 0, - mjcf.Light( - directional=True, - diffuse=(0.4, 0.4, 0.4), - specular=(0.1, 0.1, 0.1), - pos=(0, 0, 5.0), - dir=(0, 0, -1), - castshadow=False, - ).to_xml(), - ) - worldbody.insert( - 0, - mjcf.Light( - directional=True, diffuse=(0.6, 0.6, 0.6), specular=(0.2, 0.2, 0.2), pos=(0, 0, 4), dir=(0, 0, -1) - ).to_xml(), - ) - if add_floor: - worldbody.insert( - 0, - mjcf.Geom( - name="ground", - type="plane", - size=(0, 0, 1), - pos=(0.001, 0, 0), - quat=(1, 0, 0, 0), - material="matplane", - condim=1, - conaffinity=15, - ).to_xml(), - ) - worldbody = root.find("worldbody") - - motors: List[mjcf.Motor] = [] - sensor_pos: List[mjcf.Actuatorpos] = [] - sensor_vel: List[mjcf.Actuatorvel] = [] - sensor_frc: List[mjcf.Actuatorfrc] = [] - # Create motors and sensors for the joints - joints = list(root.findall(".//joint")) - original_joints = joints.copy() - for joint_xml in joints: # robot.all_joints(): - joint = joint_xml.get("name") - if joint in robot.default_standing().keys(): - joint_name = joint - limit = 1000.0 # 200.0 # Ensure limit is a float - keys = robot.effort().keys() - for key in keys: - if key in joint_name: - limit = robot.effort()[key] - print(f"Joint: {joint}, limit: {limit}") - motors.append( - mjcf.Motor( - name=joint, - joint=joint, - gear=1, - ctrlrange=(-limit, limit), - ctrllimited=True, - ) - ) - sensor_pos.append(mjcf.Actuatorpos(name=joint + "_p", actuator=joint, user="13")) - sensor_vel.append(mjcf.Actuatorvel(name=joint + "_v", actuator=joint, user="13")) - sensor_frc.append(mjcf.Actuatorfrc(name=joint + "_f", actuator=joint, user="13", noise=0.001)) - else: - print(f"Joint: {joint} not in default_standing") - root = self.update_joints(root) - # Add motors and sensors - root.append(mjcf.Actuator(motors).to_xml()) - root.append(mjcf.Sensor(sensor_pos, sensor_vel, sensor_frc).to_xml()) - - # TODO: Add additional sensors when necessary - sensors = root.find("sensor") - sensors.extend( - [ - ET.Element("framequat", name="orientation", objtype="site", noise="0.001", objname="imu"), - ET.Element("gyro", name="angular-velocity", site="imu", noise="0.005", cutoff="34.9"), - # ET.Element("framepos", name="position", objtype="site", noise="0.001", objname="imu"), - # ET.Element("velocimeter", name="linear-velocity", site="imu", noise="0.001", cutoff="30"), - # ET.Element("accelerometer", name="linear-acceleration", site="imu", noise="0.005", cutoff="157"), - # ET.Element("magnetometer", name="magnetometer", site="imu"), - ] - ) - - root.insert( - 1, - mjcf.Option( - timestep=0.001, - iterations=50, - solver="PGS", - gravity=(0, 0, -9.81), - ).to_xml(), - ) - - visual_geom = ET.Element("default", {"class": "visualgeom"}) - geom_attributes = {"material": "visualgeom", "condim": "1", "contype": "0", "conaffinity": "0"} - ET.SubElement(visual_geom, "geom", geom_attributes) - - root.insert( - 1, - mjcf.Default( - joint=mjcf.Joint(armature=0.01, damping=0.01, limited=True, frictionloss=0.01), - motor=mjcf.Motor(ctrllimited=True), - equality=mjcf.Equality(solref=(0.001, 2)), - geom=mjcf.Geom( - solref=(0.001, 2), - friction=(0.9, 0.2, 0.2), - condim=4, - contype=1, - conaffinity=15, - ), - visual_geom=visual_geom, - ).to_xml(), - ) - - # Locate actual root body inside of worldbody - root_body = worldbody.find(".//body") - root_body.set("pos", "0 0 0") - root_body.set("quat", " ".join(map(str, robot.rotation))) - - # Add cameras and imu - root_body.insert(1, ET.Element("camera", name="front", pos="0 -3 1", xyaxes="1 0 0 0 1 2", mode="trackcom")) - root_body.insert( - 2, - ET.Element( - "camera", - name="side", - pos="-2.893 -1.330 0.757", - xyaxes="0.405 -0.914 0.000 0.419 0.186 0.889", - mode="trackcom", - ), - ) - root_body.insert(3, ET.Element("site", name="imu", size="0.01", pos="0 0 0")) - - # add visual geom logic - for body in root.findall(".//body"): - original_geoms = list(body.findall("geom")) - for geom in original_geoms: - geom.set("class", "visualgeom") - # Create a new geom element - new_geom = ET.Element("geom") - new_geom.set("type", geom.get("type") or "") # Ensure type is not None - new_geom.set("rgba", geom.get("rgba") or "1 0.5 0.75 1") # Ensure rgba is not None - - # Check if geom has mesh or is a box - if geom.get("mesh") is None: - if geom.get("type") == "box": - new_geom.set("type", "box") - new_geom.set("size", geom.get("size") or "") - else: - print(f"Unknown geom type: {geom.get('type')}") - else: - new_geom.set("mesh", geom.get("mesh")) - if geom.get("pos"): - new_geom.set("pos", geom.get("pos") or "") - if geom.get("quat"): - new_geom.set("quat", geom.get("quat") or "") - try: - # Exclude collision meshes when setting contact - if geom.get("mesh") not in robot.collision_links and geom.get("mesh") is not None: - new_geom.set("contype", "0") - new_geom.set("conaffinity", "0") - new_geom.set("group", "1") - new_geom.set("density", "0") - except Exception as e: - print(e) - - # Append the new geom to the body - index = list(body).index(geom) - body.insert(index + 1, new_geom) - - for body in root.findall(".//body"): - joints = list(body.findall("joint")) - for join in joints: - if "actuatorfrcrange" in join.attrib: - join.attrib.pop("actuatorfrcrange") - - default_standing = robot.default_standing() - joint_defaults = list(default_standing.values()) - - qpos = ( - [0, 0, robot.height] - + [robot.rotation[3], robot.rotation[0], robot.rotation[1], robot.rotation[2]] - + joint_defaults - ) - - default_key = mjcf.Key(name="default", qpos=" ".join(map(str, qpos))) - keyframe = mjcf.Keyframe(keys=[default_key]) - root.append(keyframe.to_xml()) - - # Swap left and right clavicle (not necessary for current setup - wesley) - if False: - self.swap_bodies(root, "L_clav", "R_clav") - - def save(self, path: Union[str, Path]) -> None: - rough_string = ET.tostring(self.tree.getroot(), "utf-8", xml_declaration=False) - # Pretty print the XML - formatted_xml = _pretty_print_xml(rough_string) - logger.info("XML:\n%s", formatted_xml) - with open(path, "w") as f: - f.write(formatted_xml) - - -def create_mjcf(filepath: Path) -> None: - """Create a MJCF file for the Stompy robot.""" - path = Path(filepath) - robot_name = path.stem - path = path.parent - robot = Sim2SimRobot( - robot_name, - path, - mjcf.Compiler(angle="radian", meshdir="meshes", autolimits=True, eulerseq="zyx"), - ) - robot.adapt_world() - - robot.save(path / f"{robot_name}_fixed.xml") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Create a MJCF file for the robot.") - parser.add_argument("filepath", type=str, help="The path to load and save the MJCF file.") - parser.add_argument("--robot", type=str, help="The robot name to load.") - args = parser.parse_args() - # Robot name is whatever string comes right before ".urdf" extension - robot_name = args.filepath.split("/")[-1].split(".")[0] - - # Load the robot and config - robot = load_embodiment(args.robot) - create_mjcf(args.filepath) diff --git a/sim/scripts/download_assets.sh b/sim/scripts/download_assets.sh index 34e54529..03ff11a4 100755 --- a/sim/scripts/download_assets.sh +++ b/sim/scripts/download_assets.sh @@ -2,17 +2,10 @@ echo "If any of the downloads fail, you can manually download the assets and place them in the respective directories." -# Stompymini -echo -echo "Downloading Stompymini assets..." -gdown --id 1Ah92CTN63M2h4uKPLf8eXkX2H9Bg3Kji -unzip meshes.zip -d sim/resources/stompymini/ -rm meshes.zip - -# Stompypro +# Stompymicro echo -echo "Downloading Stompypro assets..." -gdown --folder https://drive.google.com/drive/folders/1-iIqy8j4gF6JeuMc_MjxkRe4vSZl8Ozp -O sim/resources/stompypro/ +echo "Downloading Stompymicro assets..." +gdown --folder https://drive.google.com/drive/folders/1C_v0FKoc6um0tUK2f1e6cWXtfvuc-xsD -O sim/resources/stompymicro/ # Xbot echo @@ -37,9 +30,3 @@ echo "Downloading H1_2 assets..." gdown --id 19ih7zG6Ky8xJVJD5M1th2hmqtxaNiZyh tar -xzvf meshes.zip -C sim/resources/h1_2/ rm meshes.zip - -# Stompymicro -echo -echo "Downloading Stompymicro assets..." -gdown --folder https://drive.google.com/drive/folders/1C_v0FKoc6um0tUK2f1e6cWXtfvuc-xsD -O sim/resources/stompymicro/ - diff --git a/sim/scripts/mjcf.py b/sim/scripts/mjcf.py deleted file mode 100755 index e86b070f..00000000 --- a/sim/scripts/mjcf.py +++ /dev/null @@ -1,524 +0,0 @@ -# mypy: disable-error-code="operator,union-attr" -"""Defines common types and functions for exporting MJCF files. - -API reference: -https://github.com/google-deepmind/mujoco/blob/main/src/xml/xml_native_writer.cc#L780 - -Todo: - 1. Clean up the inertia config - 2. Add visual and imu support -""" - -import glob -import os -import shutil -import xml.etree.ElementTree as ET -from dataclasses import dataclass, field -from pathlib import Path -from typing import List, Literal, Optional, Tuple, Union - -import mujoco - - -@dataclass -class Compiler: - coordinate: Optional[Literal["local", "global"]] = None - angle: Literal["radian", "degree"] = "radian" - meshdir: Optional[str] = None - eulerseq: Optional[Literal["xyz", "zxy", "zyx", "yxz", "yzx", "xzy"]] = None - autolimits: Optional[bool] = None - - def to_xml(self, compiler: Optional[ET.Element] = None) -> ET.Element: - if compiler is None: - compiler = ET.Element("compiler") - if self.coordinate is not None: - compiler.set("coordinate", self.coordinate) - compiler.set("angle", self.angle) - if self.meshdir is not None: - compiler.set("meshdir", self.meshdir) - if self.eulerseq is not None: - compiler.set("eulerseq", self.eulerseq) - if self.autolimits is not None: - compiler.set("autolimits", str(self.autolimits).lower()) - return compiler - - -@dataclass -class Mesh: - name: str - file: str - scale: Optional[Tuple[float, float, float]] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - mesh = ET.Element("mesh") if root is None else ET.SubElement(root, "mesh") - mesh.set("name", self.name) - mesh.set("file", self.file) - if self.scale is not None: - mesh.set("scale", " ".join(map(str, self.scale))) - return mesh - - -@dataclass -class Joint: - name: Optional[str] = None - type: Optional[Literal["hinge", "slide", "ball", "free"]] = None - pos: Optional[Tuple[float, float, float]] = None - axis: Optional[Tuple[float, float, float]] = None - limited: Optional[bool] = None - range: Optional[Tuple[float, float]] = None - damping: Optional[float] = None - stiffness: Optional[float] = None - armature: Optional[float] = None - frictionloss: Optional[float] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - joint = ET.Element("joint") if root is None else ET.SubElement(root, "joint") - if self.name is not None: - joint.set("name", self.name) - if self.type is not None: - joint.set("type", self.type) - if self.pos is not None: - joint.set("pos", " ".join(map(str, self.pos))) - if self.axis is not None: - joint.set("axis", " ".join(map(str, self.axis))) - if self.range is not None: - joint.set("range", " ".join(map(str, self.range))) - if self.limited is not None: - joint.set("limited", str(self.limited).lower()) - if self.damping is not None: - joint.set("damping", str(self.damping)) - if self.stiffness is not None: - joint.set("stiffness", str(self.stiffness)) - if self.armature is not None: - joint.set("armature", str(self.armature)) - if self.frictionloss is not None: - joint.set("frictionloss", str(self.frictionloss)) - return joint - - -@dataclass -class Geom: - name: Optional[str] = None - type: Optional[Literal["plane", "sphere", "cylinder", "box", "capsule", "ellipsoid", "mesh"]] = None - plane: Optional[str] = None - rgba: Optional[Tuple[float, float, float, float]] = None - pos: Optional[Tuple[float, float, float]] = None - quat: Optional[Tuple[float, float, float, float]] = None - matplane: Optional[str] = None - material: Optional[str] = None - condim: Optional[int] = None - contype: Optional[int] = None - conaffinity: Optional[int] = None - size: Optional[Tuple[float, float, float]] = None - friction: Optional[Tuple[float, float, float]] = None - solref: Optional[Tuple[float, float]] = None - density: Optional[float] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - geom = ET.Element("geom") if root is None else ET.SubElement(root, "geom") - if self.name is not None: - geom.set("name", self.name) - if self.type is not None: - geom.set("type", self.type) - if self.rgba is not None: - geom.set("rgba", " ".join(map(str, self.rgba))) - if self.pos is not None: - geom.set("pos", " ".join(map(str, self.pos))) - if self.quat is not None: - geom.set("quat", " ".join(map(str, self.quat))) - if self.matplane is not None: - geom.set("matplane", self.matplane) - if self.material is not None: - geom.set("material", self.material) - if self.condim is not None: - geom.set("condim", str(self.condim)) - if self.contype is not None: - geom.set("contype", str(self.contype)) - if self.conaffinity is not None: - geom.set("conaffinity", str(self.conaffinity)) - if self.plane is not None: - geom.set("plane", self.plane) - if self.size is not None: - geom.set("size", " ".join(map(str, self.size))) - if self.friction is not None: - geom.set("friction", " ".join(map(str, self.friction))) - if self.solref is not None: - geom.set("solref", " ".join(map(str, self.solref))) - if self.density is not None: - geom.set("density", str(self.density)) - return geom - - -@dataclass -class Body: - name: str - pos: Optional[Tuple[float, float, float]] = field(default=None) - quat: Optional[Tuple[float, float, float, float]] = field(default=None) - geom: Optional[Geom] = field(default=None) - joint: Optional[Joint] = field(default=None) - # TODO - fix inertia, until then rely on Mujoco's engine - # inertial: Inertial = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - body = ET.Element("body") if root is None else ET.SubElement(root, "body") - body.set("name", self.name) - if self.pos is not None: - body.set("pos", " ".join(map(str, self.pos))) - if self.quat is not None: - body.set("quat", " ".join(f"{x:.8g}" for x in self.quat)) - if self.joint is not None: - self.joint.to_xml(body) - if self.geom is not None: - self.geom.to_xml(body) - return body - - -@dataclass -class Flag: - frictionloss: Optional[str] = None - # removed at 3.1.4 - # sensornoise: str | None = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - flag = ET.Element("flag") if root is None else ET.SubElement(root, "flag") - if self.frictionloss is not None: - flag.set("frictionloss", self.frictionloss) - return flag - - -@dataclass -class Option: - timestep: Optional[float] = None - viscosity: Optional[float] = None - iterations: Optional[int] = None - solver: Optional[Literal["PGS", "CG", "Newton"]] = None - gravity: Optional[Tuple[float, float, float]] = None - flag: Optional[Flag] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - if root is None: - option = ET.Element("option") - else: - option = ET.SubElement(root, "option") - if self.iterations is not None: - option.set("iterations", str(self.iterations)) - if self.timestep is not None: - option.set("timestep", str(self.timestep)) - if self.viscosity is not None: - option.set("viscosity", str(self.viscosity)) - if self.solver is not None: - option.set("solver", self.solver) - if self.gravity is not None: - option.set("gravity", " ".join(map(str, self.gravity))) - if self.flag is not None: - self.flag.to_xml(option) - return option - - -@dataclass -class Motor: - name: Optional[str] = None - joint: Optional[str] = None - ctrlrange: Optional[Tuple[float, float]] = None - ctrllimited: Optional[bool] = None - gear: Optional[float] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - if root is None: - motor = ET.Element("motor") - else: - motor = ET.SubElement(root, "motor") - if self.name is not None: - motor.set("name", self.name) - if self.joint is not None: - motor.set("joint", self.joint) - if self.ctrllimited is not None: - motor.set("ctrllimited", str(self.ctrllimited).lower()) - if self.ctrlrange is not None: - motor.set("ctrlrange", " ".join(map(str, self.ctrlrange))) - if self.gear is not None: - motor.set("gear", str(self.gear)) - return motor - - -@dataclass -class Light: - directional: bool = True - diffuse: Optional[Tuple[float, float, float]] = None - specular: Optional[Tuple[float, float, float]] = None - pos: Optional[Tuple[float, float, float]] = None - dir: Optional[Tuple[float, float, float]] = None - castshadow: Optional[bool] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - if root is None: - light = ET.Element("light") - else: - light = ET.SubElement(root, "light") - if self.directional is not None: - light.set("directional", str(self.directional).lower()) - if self.diffuse is not None: - light.set("diffuse", " ".join(map(str, self.diffuse))) - if self.specular is not None: - light.set("specular", " ".join(map(str, self.specular))) - if self.pos is not None: - light.set("pos", " ".join(map(str, self.pos))) - if self.dir is not None: - light.set("dir", " ".join(map(str, self.dir))) - if self.castshadow is not None: - light.set("castshadow", str(self.castshadow).lower()) - return light - - -@dataclass -class Equality: - solref: Tuple[float, float] - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - equality = ET.Element("equality") if root is None else ET.SubElement(root, "equality") - equality.set("solref", " ".join(map(str, self.solref))) - return equality - - -@dataclass -class Site: - name: Optional[str] = None - size: Optional[float] = None - pos: Optional[Tuple[float, float, float]] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - site = ET.Element("site") if root is None else ET.SubElement(root, "site") - if self.name is not None: - site.set("name", self.name) - if self.size is not None: - site.set("size", str(self.size)) - if self.pos is not None: - site.set("pos", " ".join(map(str, self.pos))) - return site - - -@dataclass -class Default: - joint: Optional[Joint] = None - geom: Optional[Geom] = None - class_: Optional[str] = None - motor: Optional[Motor] = None - equality: Optional[Equality] = None - visual_geom: Optional[ET.Element] = None - - def to_xml(self, default: Optional[ET.Element] = None) -> ET.Element: - default = ET.Element("default") if default is None else ET.SubElement(default, "default") - if self.joint is not None: - self.joint.to_xml(default) - if self.geom is not None: - self.geom.to_xml(default) - if self.class_ is not None: - default.set("class", self.class_) - if self.motor is not None: - self.motor.to_xml(default) - if self.equality is not None: - self.equality.to_xml(default) - if self.visual_geom is not None: - default.append(self.visual_geom) - return default - - -@dataclass -class Actuator: - motors: List[Motor] - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - actuator = ET.Element("actuator") if root is None else ET.SubElement(root, "actuator") - for motor in self.motors: - motor.to_xml(actuator) - return actuator - - -@dataclass -class Actuatorpos: - name: Optional[str] = None - actuator: Optional[str] = None - user: Optional[str] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - actuatorpos = ET.Element("actuatorpos") if root is None else ET.SubElement(root, "actuatorpos") - if self.name is not None: - actuatorpos.set("name", self.name) - if self.actuator is not None: - actuatorpos.set("actuator", self.actuator) - if self.user is not None: - actuatorpos.set("user", self.user) - return actuatorpos - - -@dataclass -class Actuatorvel: - name: Optional[str] = None - actuator: Optional[str] = None - user: Optional[str] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - actuatorvel = ET.Element("actuatorvel") if root is None else ET.SubElement(root, "actuatorvel") - if self.name is not None: - actuatorvel.set("name", self.name) - if self.actuator is not None: - actuatorvel.set("actuator", self.actuator) - if self.user is not None: - actuatorvel.set("user", self.user) - return actuatorvel - - -@dataclass -class Actuatorfrc: - name: Optional[str] = None - actuator: Optional[str] = None - user: Optional[str] = None - noise: Optional[float] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - actuatorfrc = ET.Element("actuatorfrc") if root is None else ET.SubElement(root, "actuatorfrc") - if self.name is not None: - actuatorfrc.set("name", self.name) - if self.actuator is not None: - actuatorfrc.set("actuator", self.actuator) - if self.user is not None: - actuatorfrc.set("user", self.user) - if self.noise is not None: - actuatorfrc.set("noise", str(self.noise)) - return actuatorfrc - - -@dataclass -class Sensor: - actuatorpos: Optional[List[Actuatorpos]] = None - actuatorvel: Optional[List[Actuatorvel]] = None - actuatorfrc: Optional[List[Actuatorfrc]] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - sensor = ET.Element("sensor") if root is None else ET.SubElement(root, "sensor") - if self.actuatorpos is not None: - for actuatorpos in self.actuatorpos: - actuatorpos.to_xml(sensor) - if self.actuatorvel is not None: - for actuatorvel in self.actuatorvel: - actuatorvel.to_xml(sensor) - if self.actuatorfrc is not None: - for actuatorfrc in self.actuatorfrc: - actuatorfrc.to_xml(sensor) - return sensor - - -@dataclass -class Key: - name: Optional[str] = None - qpos: Optional[str] = None - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - key = ET.Element("key") if root is None else ET.SubElement(root, "key") - if self.name is not None: - key.set("name", self.name) - if self.qpos is not None: - key.set("qpos", self.qpos) - return key - - -@dataclass -class Keyframe: - keys: List[Key] - - def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: - keyframe = ET.Element("keyframe") if root is None else ET.SubElement(root, "keyframe") - for key in self.keys: - key.to_xml(keyframe) - return keyframe - - -def _copy_stl_files(source_directory: Union[str, Path], destination_directory: Union[str, Path]) -> None: - # Ensure the destination directory exists, create if not - os.makedirs(destination_directory, exist_ok=True) - - # Use glob to find all .stl files in the source directory - pattern = os.path.join(source_directory, "*.stl") - for file_path in glob.glob(pattern): - destination_path = os.path.join(destination_directory, os.path.basename(file_path)) - shutil.copy(file_path, destination_path) - print(f"Copied {file_path} to {destination_path}") - - -def _remove_stl_files(source_directory: Union[str, Path]) -> None: - for filename in os.listdir(source_directory): - if filename.endswith(".stl"): - file_path = os.path.join(source_directory, filename) - os.remove(file_path) - - -class Robot: - """A class to adapt the world in a Mujoco XML file.""" - - def __init__( - self, - robot_name: str, - output_dir: Union[str, Path], - compiler: Optional[Compiler] = None, - ) -> None: - """Initialize the robot. - - Args: - robot_name (str): The name of the robot. - output_dir (Union[str, Path]): The output directory. - compiler (Compiler, optional): The compiler settings. - """ - self.robot_name = robot_name - self.output_dir = output_dir - self.compiler = compiler - self._set_clean_up() - self.tree = ET.parse(self.output_dir / f"{self.robot_name}.xml") - - def _set_clean_up(self) -> None: - # HACK - # mujoco has a hard time reading meshes - _copy_stl_files(self.output_dir / "meshes", self.output_dir) - - urdf_tree = ET.parse(self.output_dir / f"{self.robot_name}.urdf") - root = urdf_tree.getroot() - - tree = ET.ElementTree(root) - tree.write(self.output_dir / f"{self.robot_name}.urdf", encoding="utf-8") - - model = mujoco.MjModel.from_xml_path((self.output_dir / f"{self.robot_name}.urdf").as_posix()) - mujoco.mj_saveLastXML((self.output_dir / f"{self.robot_name}.xml").as_posix(), model) - # remove all the files - _remove_stl_files(self.output_dir) - - def adapt_world(self) -> None: - root = self.tree.getroot() - - # Turn off internal collisions - for element in root: - if element.tag == "geom": - element.attrib["contype"] = str(1) - element.attrib["conaffinity"] = str(0) - - compiler = root.find("compiler") - if self.compiler is not None: - compiler = self.compiler.to_xml(compiler) - - worldbody = root.find("worldbody") - new_root_body = Body(name="root", pos=(0, 0, 0), quat=(1, 0, 0, 0)).to_xml() - - # List to store items to be moved to the new root body - items_to_move = [] - # Gather all children (geoms and bodies) that need to be moved under the new root body - for element in worldbody: - items_to_move.append(element) - # Move gathered elements to the new root body - for item in items_to_move: - worldbody.remove(item) - new_root_body.append(item) - - # Add the new root body to the worldbody - worldbody.append(new_root_body) - self.tree = ET.ElementTree(root) - - def save(self, path: Union[str, Path]) -> None: - self.tree.write(path, encoding="utf-8", xml_declaration=False) diff --git a/sim/sim2sim.py b/sim/sim2sim.py index 87e40974..b5129523 100755 --- a/sim/sim2sim.py +++ b/sim/sim2sim.py @@ -1,8 +1,9 @@ """Sim2sim deployment test. Run: - python sim/sim2sim.py --load_model examples/walking_policy.onnx --embodiment gpr + python sim/sim2sim.py --load_model examples/gpr_walking.kinfer --embodiment gpr """ + import argparse import math import os @@ -15,14 +16,14 @@ import numpy as np import onnxruntime as ort import pygame -from scipy.spatial.transform import Rotation as R import torch +from kinfer.export.pytorch import export_to_onnx +from kinfer.inference.python import ONNXModel +from model_export import ActorCfg, get_actor_policy +from scipy.spatial.transform import Rotation as R from tqdm import tqdm from sim.h5_logger import HDF5Logger -from model_export import ActorCfg, get_actor_policy -from kinfer.export.pytorch import export_to_onnx -from kinfer.inference.python import ONNXModel @dataclass @@ -34,7 +35,6 @@ class Sim2simCfg: cycle_time: float = 0.25 - def handle_keyboard_input() -> None: global x_vel_cmd, y_vel_cmd, yaw_vel_cmd @@ -186,7 +186,7 @@ def run_mujoco( num_actions=model_info["num_actions"], max_timesteps=stop_state_log, num_observations=model_info["num_observations"], - h5_out_dir=h5_out_dir + h5_out_dir=h5_out_dir, ) # Initialize variables for tracking upright steps and average speed @@ -246,27 +246,29 @@ def run_mujoco( hist_obs = policy_output["x.3"] target_q = positions - + if log_h5: - logger.log_data({ - "t": np.array([count_lowlevel * cfg.dt], dtype=np.float32), - "2D_command": np.array( - [ - np.sin(2 * math.pi * count_lowlevel * cfg.dt / cfg.cycle_time), - np.cos(2 * math.pi * count_lowlevel * cfg.dt / cfg.cycle_time), - ], - dtype=np.float32, - ), - "3D_command": np.array([x_vel_cmd, y_vel_cmd, yaw_vel_cmd], dtype=np.float32), - "joint_pos": cur_pos_obs.astype(np.float32), - "joint_vel": cur_vel_obs.astype(np.float32), - "prev_actions": prev_actions.astype(np.float32), - "curr_actions": curr_actions.astype(np.float32), - "ang_vel": omega.astype(np.float32), - "euler_rotation": eu_ang.astype(np.float32), - "buffer": hist_obs.astype(np.float32) - }) - + logger.log_data( + { + "t": np.array([count_lowlevel * cfg.dt], dtype=np.float32), + "2D_command": np.array( + [ + np.sin(2 * math.pi * count_lowlevel * cfg.dt / cfg.cycle_time), + np.cos(2 * math.pi * count_lowlevel * cfg.dt / cfg.cycle_time), + ], + dtype=np.float32, + ), + "3D_command": np.array([x_vel_cmd, y_vel_cmd, yaw_vel_cmd], dtype=np.float32), + "joint_pos": cur_pos_obs.astype(np.float32), + "joint_vel": cur_vel_obs.astype(np.float32), + "prev_actions": prev_actions.astype(np.float32), + "curr_actions": curr_actions.astype(np.float32), + "ang_vel": omega.astype(np.float32), + "euler_rotation": eu_ang.astype(np.float32), + "buffer": hist_obs.astype(np.float32), + } + ) + prev_actions = curr_actions # Generate PD control @@ -296,6 +298,7 @@ def run_mujoco( if log_h5: logger.close() + if __name__ == "__main__": parser = argparse.ArgumentParser(description="Deployment script.") parser.add_argument("--embodiment", type=str, required=True, help="Embodiment name.") @@ -312,7 +315,7 @@ def run_mujoco( pygame.init() pygame.display.set_caption("Simulation Control") else: - x_vel_cmd, y_vel_cmd, yaw_vel_cmd = np.random.uniform(0.1, 0.5), 0.0, 0.0 + x_vel_cmd, y_vel_cmd, yaw_vel_cmd = 0.2, 0.0, 0.0 policy_cfg = ActorCfg(embodiment=args.embodiment) if args.embodiment == "gpr": @@ -342,12 +345,7 @@ def run_mujoco( # Merge policy_cfg and sim2sim_info into a single config object export_config = {**vars(policy_cfg), **sim2sim_info} print(export_config) - export_to_onnx( - actor_model, - input_tensors=input_tensors, - config=export_config, - save_path="kinfer_test.onnx" - ) + export_to_onnx(actor_model, input_tensors=input_tensors, config=export_config, save_path="kinfer_test.onnx") policy = ONNXModel("kinfer_test.onnx") metadata = policy.get_metadata() diff --git a/sim/sim2sim_old.py b/sim/sim2sim_old.py index fd761aa8..b887e3fd 100755 --- a/sim/sim2sim_old.py +++ b/sim/sim2sim_old.py @@ -1,7 +1,10 @@ +""" Sim2sim script for running a policy in MuJoCo. + +Run: + python sim/sim2sim.py --load_model examples/standing_pro.pt --embodiment gpr + python sim/sim2sim.py --load_model examples/standing_micro.pt --embodiment zeroth """ -python sim/sim2sim.py --load_model examples/standing_pro.pt --embodiment gpr -python sim/sim2sim.py --load_model examples/standing_micro.pt --embodiment zeroth -""" + import argparse import math import os @@ -17,7 +20,6 @@ from scipy.spatial.transform import Rotation as R from tqdm import tqdm -# from sim.h5_logger import HDF5Logger from sim.model_export import ActorCfg, convert_model_to_onnx @@ -76,7 +78,7 @@ def quaternion_to_euler_array(quat: np.ndarray) -> np.ndarray: def get_obs(data: mujoco.MjData) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - """Extracts an observation from the mujoco data structure""" + """Extracts an observation from the mujoco data structure.""" q = data.qpos.astype(np.double) dq = data.qvel.astype(np.double) quat = data.sensor("orientation").data[[1, 2, 3, 0]].astype(np.double) @@ -95,7 +97,7 @@ def pd_control( kd: np.ndarray, default: np.ndarray, ) -> np.ndarray: - """Calculates torques from position commands""" + """Calculates torques from position commands.""" return kp * (target_q + default - q) - kd * dq @@ -109,12 +111,16 @@ def run_mujoco( render: bool = True, h5_out_dir: str = "sim/resources", ) -> None: - """ - Run the Mujoco simulation using the provided policy and configuration. + """Run the Mujoco simulation using the provided policy and configuration. Args: policy: The policy used for controlling the simulation. cfg: The configuration object containing simulation settings. + model_info: The model information dictionary. + keyboard_use: Whether to use keyboard input. + log_h5: Whether to log data to HDF5. + render: Whether to render the simulation. + h5_out_dir: The directory to save HDF5 files. Returns: None @@ -174,16 +180,6 @@ def run_mujoco( "buffer.1": np.zeros(model_info["num_observations"]).astype(np.float32), } - if log_h5: - stop_state_log = int(cfg.sim_duration / cfg.dt) / cfg.decimation - logger = HDF5Logger( - data_name=embodiment, - num_actions=model_info["num_actions"], - max_timesteps=stop_state_log, - num_observations=model_info["num_observations"], - h5_out_dir=h5_out_dir - ) - # Initialize variables for tracking upright steps and average speed upright_steps = 0 total_speed = 0.0 @@ -263,9 +259,6 @@ def run_mujoco( print(f"Number of upright steps: {upright_steps}") print(f"Average speed: {average_speed:.4f} m/s") - if log_h5: - logger.close() - def parse_modelmeta( modelmeta: List[Tuple[str, str]], @@ -333,9 +326,7 @@ def parse_modelmeta( if args.load_model.endswith(".onnx"): policy = ort.InferenceSession(args.load_model) else: - policy = convert_model_to_onnx( - args.load_model, policy_cfg, save_path="policy.onnx" - ) + policy = convert_model_to_onnx(args.load_model, policy_cfg, save_path="policy.onnx") model_info = parse_modelmeta( policy.get_modelmeta().custom_metadata_map.items(),