cyy/RoboGauge
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

v0.1.8

Browse Source
This commit is contained in:
wty-yy
2025-12-06 17:11:40 +08:00
parent a05d047b39
commit 9d509829a5
19 changed files with 364 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@@ -27,7 +27,7 @@
| 参数名称 | 变量名 | 范围 |
| - | - | - |
| 电机动作执行随机延迟 | `action delay` | `<= RL控制间隔` |
| base负重 | `base mass` | `(-1, 5) kg` |
| base负重 | `base mass` | `-1, 0, 1, 2, 3 kg` |
#### 地形
1. 支持legged_gym中的部分地形, 包括: `wave, slope, rough_slope, stairs up, stairs down, obstacles, flat`, 除`flat`地形外其他地形可进行难度系数提升

5
UPDATE.md
View File

@@ -1,4 +1,9 @@
# UPDATE
## 20251206
### v0.1.8
1. 加入run_eval_models.sh多模型评估bash脚本
2. 加入域随机化, `action delay`, 基于配置修改的`base mass`, `friction`
3. 加入MultiPipeline支持多种seeds, frictions, base masses并行评估
## 20251205
### v0.1.7
1. 加入moe模型的测试, 及moe模型, 平地的指令最大范围开到2

2
resources/terrains/flat.xml
View File

@@ -17,6 +17,6 @@
<worldbody>
<light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
<geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
<geom name="floor" size="0 0 0.05" type="plane" material="groundplane" priority="1"/>
</worldbody>
</mujoco>

9
robogauge/scripts/evaluate_models.txt Normal file
View File

@@ -0,0 +1,9 @@
### Support using '#' as the starting comments ###
### Uncommented lines provide the path to evaluated model ###
/path/to/your/eval_model1.pt
/path/to/your/eval_model2.pt
/path/to/your/eval_model3.pt
# /path/to/your/uneval_model1.pt
# /path/to/your/uneval_model2.pt

13
robogauge/scripts/run.py
View File

@@ -11,13 +11,18 @@ import os
os.environ['MUJOCO_GL'] = 'glfw' # avoid mujoco.Renderer EGL context error
from robogauge.tasks import *
from robogauge.tasks.pipeline.multi_pipeline import MultiPipeline
from robogauge.utils.task_register import task_register
from robogauge.utils.helpers import parse_args
from robogauge.utils.logger import logger
if __name__ == '__main__':
args = parse_args()
logger.create(args.experiment_name, args.run_name)
logger.info(f"Starting experiment: {args.experiment_name}")
pipeline: BasePipeline = task_register.make_pipeline(args=args)
pipeline.run()
if not args.multi:
pipeline: BasePipeline = task_register.make_pipeline(args=args)
pipeline.run()
else:
multi_pipeline = MultiPipeline(args)
multi_pipeline.run()

74
robogauge/scripts/run_eval_models.sh Executable file
View File

@@ -0,0 +1,74 @@
#!/bin/bash
# Description: This script automates the evaluation of multiple models
# specified in default ./evaluate_models.txt file.
#
# Usage: ./run_scripts.sh [-n EXP_NAME] [-t TASK_NAME] [-s]
# -n EXP_NAME Experiment name (default: go2_moe_flat)
# -t TASK_NAME Task name (default: go2_moe_flat)
# -s Save video (default: false)
# -h Show this help message
### Find the directory of the script ###
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODELS_FILE="$SCRIPT_DIR/evaluate_models.txt"
RUN_PY="$SCRIPT_DIR/run.py"
### Default Configure ###
EXP_NAME="go2_moe_flat" # Experiment name [-n]
TASK_NAME="go2_moe_flat" # Task name [-t]
SAVE_VIDEO=false # Whether to save video [-s]
### Parse Arguments ###
usage() {
echo "Usage: $0 [-n EXP_NAME] [-t TASK_NAME] [-s]"
echo " -n EXP_NAME Experiment name (default: ${EXP_NAME})"
echo " -t TASK_NAME Task name (default: ${TASK_NAME})"
echo " -s Save video (default: ${SAVE_VIDEO})"
echo " -h Show this help message"
exit 1
}
while getopts "n:t:sh" opt; do
case "${opt}" in
n) EXP_NAME="$OPTARG" ;; # Experiment name [-n]
t) TASK_NAME="$OPTARG" ;; # Task name [-t]
s) SAVE_VIDEO=true ;; # Whether to save video [-s]
h) usage ;; # Print usage [-h]
*) usage ;; # Print usage for invalid options
esac
done
### Activate Conda Environment ###
eval "$(conda shell.bash hook)"
conda activate kaiwu
### Read Models from File ###
echo "Reading models from: $MODELS_FILE"
mapfile -t models_paths < <(grep -v -e '^[[:space:]]*$' -e '^[[:space:]]*#' "$MODELS_FILE")
if [ ${#models_paths[@]} -eq 0 ]; then
echo "Error: $MODELS_FILE is empty or contains only blank lines."
exit 1
fi
### Run Evaluation Scripts ###
base_args="--task $TASK_NAME --headless --experiment-name $EXP_NAME"
if [ "$SAVE_VIDEO" = true ]; then
base_args="$base_args --save-video"
fi
echo "================ Run Settings ================"
echo "Script Dir: $SCRIPT_DIR"
echo "Runner: $RUN_PY"
echo "Task: $TASK_NAME"
echo "Exp Name: $EXP_NAME"
echo "Save Video: $SAVE_VIDEO"
echo "Models Qty: ${#models_paths[@]}"
echo "=============================================="
for model_path in "${models_paths[@]}"; do
echo "🚀 Evaluating model: $model_path 🚀"
python "$RUN_PY" $base_args --model-path "$model_path"
done

2
robogauge/tasks/custom/go2_flat_task.py
View File

@@ -14,7 +14,7 @@ class Go2FlatGaugeConfig(FlatGaugeConfig):
cmd_duration = 5.0
class diagonal_velocity(FlatGaugeConfig.goals.diagonal_velocity):
enabled = True
enabled = False
cmd_duration = 6.0
class Go2FlatConfig(Go2Config):

26
robogauge/tasks/gauge/base_gauge.py
View File

@@ -11,7 +11,9 @@ import yaml
from typing import List
from pathlib import Path
from functools import partial
from collections import defaultdict
import numpy as np
from robogauge.utils.logger import logger
from robogauge.utils.helpers import class_to_dict, snake_to_pascal
@@ -56,7 +58,7 @@ class BaseGauge:
metric_class = eval(metric_class_name)
self.metrics.append(metric_class(robot_cfg=robot_cfg, **self.metrics_cfg[name]))
log_str += f" - Metric: {name}\n"
self.info['metric'].append(metric_class_name)
self.info['metric'].append(name)
logger.info(log_str.strip())
if len(self.goals) == 0:
@@ -109,22 +111,36 @@ class BaseGauge:
logger.info(f"New Goal [{self.goal_idx+1}/{len(self.goals)}] [{goal_obj.count+1}/{goal_obj.total}]: {self.goal_str}")
return goal
def update_metrics(self, sim_data: SimData):
def update_metrics(self, sim_data: SimData, goal_data: GoalData):
if sim_data.n_step % int(self.cfg.metrics.metric_dt / sim_data.sim_dt) != 0:
return
metrics_results = {}
for metric_name, metric_obj in zip(self.info['metric'], self.metrics):
val = metric_obj(sim_data)
val = metric_obj(sim_data, goal_data)
if metric_name not in ['visualization']:
metrics_results[metric_name] = val
self.goals[self.goal_idx].update_metrics(metrics_results)
def save_results(self):
""" Save the results to a yaml file. """
metrics = defaultdict(lambda: defaultdict(list))
for goal in self.results:
for metric_name, quantiles in self.results[goal].items():
for quantile, val in quantiles.items():
metrics[metric_name][quantile].append(val)
self.results['summary'] = {}
for metric_name, quantiles in metrics.items():
if metric_name not in self.results['summary']:
self.results['summary'][metric_name] = {}
for quantile, vals in quantiles.items():
mean = float(np.mean(vals))
std = float(np.std(vals))
self.results['summary'][metric_name][quantile] = f"{mean:.4f} ± {std:.4f}"
save_path = Path(logger.log_dir) / "results.yaml"
with open(save_path, 'w') as file:
yaml.dump(self.results, file)
yaml_str = yaml.dump(self.results)
yaml.dump(self.results, file, allow_unicode=True)
yaml_str = yaml.dump(self.results, allow_unicode=True)
logger.info(
f"""\n{'='*20} Goals and Metrics results {'='*20}\n"""
f"""{yaml_str}"""

3
robogauge/tasks/gauge/metrics/base_metric.py
View File

@@ -1,6 +1,7 @@
from robogauge.utils.logger import logger
from robogauge.tasks.robots import RobotConfig
from robogauge.tasks.simulator.sim_data import SimData
from robogauge.tasks.gauge.goals.base_goal import GoalData
class BaseMetric:
""" Base class for all metric functions. """
@@ -9,7 +10,7 @@ class BaseMetric:
def __init__(self, robot_cfg: RobotConfig, **kwargs):
self.robot_cfg = robot_cfg
def __call__(self, sim_data: SimData) -> float:
def __call__(self, sim_data: SimData, goal_data: GoalData) -> float:
value = 0.0
logger.log(value, self.name, step=sim_data.n_step)
return value

5
robogauge/tasks/gauge/metrics/dof_metrics.py
View File

@@ -1,8 +1,7 @@
import numpy as np
from robogauge.tasks.robots import RobotConfig
from robogauge.tasks.simulator.sim_data import SimData
from robogauge.tasks.gauge.metrics.base_metric import BaseMetric
from robogauge.tasks.gauge.metrics.base_metric import BaseMetric, GoalData, SimData
from robogauge.utils.logger import logger
@@ -21,7 +20,7 @@ class DofLimitsMetric(BaseMetric):
self.soft_dof_limit_ratio = soft_dof_limit_ratio
self.calc_dof_names = dof_names
def __call__(self, sim_data: SimData) -> float:
def __call__(self, sim_data: SimData, goal_data: GoalData) -> float:
values = []
for i in range(len(sim_data.proprio.joint.limits)):
lower_limit = sim_data.proprio.joint.limits[i, 0]

28
robogauge/tasks/gauge/metrics/vel_metrics.py
View File

@@ -1,11 +1,10 @@
import numpy as np
from robogauge.tasks.robots import RobotConfig
from robogauge.tasks.simulator.sim_data import SimData
from robogauge.tasks.gauge.metrics.base_metric import BaseMetric
from robogauge.tasks.gauge.goals.base_goal import GoalData
from robogauge.tasks.gauge.metrics.base_metric import BaseMetric, SimData, GoalData
from robogauge.utils.logger import logger
from robogauge.utils.helpers import class_to_dict
class LinVelErrMetric(BaseMetric):
@@ -15,13 +14,11 @@ class LinVelErrMetric(BaseMetric):
def __init__(self, robot_cfg: RobotConfig, **kwargs):
super().__init__(robot_cfg)
max_ranges = []
cfg_commands = class_to_dict(robot_cfg.commands)
for name in ['lin_vel_x', 'lin_vel_y', 'lin_vel_z']:
max_ranges.append(
max(
abs(getattr(robot_cfg.commands, name)[0]),
abs(getattr(robot_cfg.commands, name)[1])
)
)
cmds = cfg_commands.get(name)
if cmds is not None:
max_ranges.append(max(abs(cmds[0]), abs(cmds[1])))
self.norm_vel = np.linalg.norm(max_ranges)
def __call__(self, sim_data: SimData, goal_data: GoalData) -> float:
@@ -41,16 +38,11 @@ class AngVelErrMetric(BaseMetric):
def __init__(self, robot_cfg: RobotConfig, **kwargs):
super().__init__(robot_cfg)
max_ranges = []
cfg_commands = class_to_dict(robot_cfg.commands)
for name in ['ang_vel_roll', 'ang_vel_pitch', 'ang_vel_yaw']:
cmd_range = getattr(robot_cfg.commands, name)
if cmd_range is None:
cmd_range = [0, 0]
max_ranges.append(
max(
abs(cmd_range[0]),
abs(cmd_range[1])
)
)
cmds = cfg_commands.get(name)
if cmds is not None:
max_ranges.append(max(abs(cmds[0]), abs(cmds[1])))
self.norm_vel = np.linalg.norm(max_ranges)
def __call__(self, sim_data: SimData, goal_data: GoalData) -> float:

5
robogauge/tasks/gauge/metrics/visualization.py
View File

@@ -1,6 +1,5 @@
from robogauge.tasks.robots import RobotConfig
from robogauge.tasks.simulator.sim_data import SimData
from robogauge.tasks.gauge.metrics.base_metric import BaseMetric
from robogauge.tasks.gauge.metrics.base_metric import BaseMetric, SimData, GoalData
from robogauge.utils.logger import logger
@@ -18,7 +17,7 @@ class VisualizationMetric(BaseMetric):
self.dof_force = dof_force
self.dof_pos = dof_pos
def __call__(self, sim_data: SimData) -> float:
def __call__(self, sim_data: SimData, goal_data: GoalData) -> float:
for i in range(len(sim_data.proprio.joint.force)):
name = sim_data.proprio.joint.names[i]
if self.dof_force:

55
robogauge/tasks/pipeline/base_pipeline.py
View File

@@ -8,10 +8,12 @@
@Desc : Base Pipeline for Robogauge
'''
import yaml
import random
from copy import deepcopy
from pathlib import Path
from robogauge.utils.logger import logger
from robogauge.tasks.simulator import MujocoSimulator, MujocoConfig
from robogauge.tasks.simulator import MujocoSimulator, MujocoConfig, SimData
from robogauge.tasks.robots import (
BaseRobot, RobotConfig, Go2Config, Go2, Go2MoEConfig, Go2MoE
)
@@ -26,7 +28,7 @@ class BasePipeline:
gauge_cfg: BaseGaugeConfig
):
self.run_name = run_name
self.simulator_cfg = simulator_cfg
self.sim_cfg = simulator_cfg
self.robot_cfg = robot_cfg
self.gauge_cfg = gauge_cfg
@@ -36,7 +38,7 @@ class BasePipeline:
# save configs
cfg = {}
for name in ['simulator_cfg', 'robot_cfg', 'gauge_cfg']:
for name in ['sim_cfg', 'robot_cfg', 'gauge_cfg']:
obj = getattr(self, name)
obj_dict = class_to_dict(obj)
cfg.update({name: obj_dict})
@@ -52,27 +54,54 @@ class BasePipeline:
)
def run(self):
logger.info(f"🚀 Starting single run: {self.run_name}")
try:
self.load()
sim_data = self.sim.step()
frame_skip = int(self.robot_cfg.control.control_dt / self.simulator_cfg.physics.simulation_dt)
logger.info(f"Sim FPS: {1.0 / self.simulator_cfg.physics.simulation_dt:.2f}, Control FPS: {1.0 / self.robot_cfg.control.control_dt:.2f}, Frame Skip: {frame_skip:d}")
frame_skip = int(self.robot_cfg.control.control_dt / self.sim_cfg.physics.simulation_dt)
logger.info(f"Sim FPS: {1.0 / self.sim_cfg.physics.simulation_dt:.2f}, Control FPS: {1.0 / self.robot_cfg.control.control_dt:.2f}, Frame Skip: {frame_skip:d}")
logger.info("Running pipeline...")
while not self.gauge.is_done():
goal = self.gauge.get_goal(sim_data)
if goal is None:
goal_data = self.gauge.get_goal(sim_data)
if goal_data is None:
continue
obs = self.robot.build_observation(sim_data, goal)
obs = self.robot.build_observation(self.add_noise(sim_data), goal_data)
action, p_gains, d_gains, control_type = self.robot.get_action(obs)
self.sim.setup_action(action, p_gains, d_gains, control_type)
for _ in range(frame_skip):
if self.sim_cfg.domain_rand.action_delay:
actions_start_decimation = random.randint(0, frame_skip)
else:
self.sim.setup_action(action, p_gains, d_gains, control_type)
for i in range(frame_skip):
if self.sim_cfg.domain_rand.action_delay and i == actions_start_decimation:
self.sim.setup_action(action, p_gains, d_gains, control_type)
sim_data = self.sim.step()
self.gauge.update_metrics(sim_data)
self.gauge.update_metrics(sim_data, goal_data)
if self.gauge.is_reset(sim_data):
self.sim.reset()
sim_data = self.sim.step()
finally:
self.sim.close_viewer()
self.sim.close_video_writer()
logger.info("Pipeline execution finished.")
logger.info(f"Logging saved at: {logger.log_dir}")
logger.info("✅ Pipeline execution finished.")
logger.info(f"📁 Logging saved at: {logger.log_dir}")
return logger.log_dir
def add_noise(self, sim_data: SimData):
sim_data = deepcopy(sim_data)
noise_cfg = self.sim_cfg.noise
if not noise_cfg.enabled:
return sim_data
proprio = sim_data.proprio
def add_uniform_noise(data, noise_level):
for i in range(len(data)):
noise = random.uniform(-noise_level, noise_level)
data[i] += noise
add_uniform_noise(proprio.joint.pos, noise_cfg.joint_pos)
add_uniform_noise(proprio.joint.vel, noise_cfg.joint_vel)
add_uniform_noise(proprio.base.lin_vel, noise_cfg.lin_vel)
add_uniform_noise(proprio.base.ang_vel, noise_cfg.ang_vel)
add_uniform_noise(proprio.imu.lin_vel, noise_cfg.lin_vel)
add_uniform_noise(proprio.imu.ang_vel, noise_cfg.ang_vel)
return sim_data

113
robogauge/tasks/pipeline/multi_pipeline.py Normal file
View File

@@ -0,0 +1,113 @@
# -*- coding: utf-8 -*-
'''
@File : multi_pipeline.py
@Time : 2025/12/06 15:32:44
@Author : wty-yy
@Version : 1.0
@Blog : https://wty-yy.github.io/
@Desc : Multiprocessing Pipeline for Robogauge
'''
import yaml
import functools
import numpy as np
from tqdm import tqdm
import multiprocessing
from pathlib import Path
from copy import deepcopy
from itertools import product
from collections import defaultdict
from robogauge.tasks.pipeline.base_pipeline import BasePipeline
from robogauge.utils.task_register import task_register
from robogauge.utils.logger import logger
def run_single_process(args, data):
seed, base_mass, friction = data
local_args = deepcopy(args)
local_args.seed = seed
run_name = f"{local_args.run_name}_{seed}_baseMass{base_mass}_friction{friction}"
logger.create(
experiment_name=local_args.experiment_name,
run_name=run_name,
console_output=False
)
pipeline = task_register.make_pipeline(args=local_args, create_logger=False)
log_dir = pipeline.run()
return log_dir
class MultiPipeline:
def __init__(self, args):
self.args = args
self.seeds = args.seeds
self.frictions = args.frictions
self.base_masses = args.base_masses
self.num_processes = args.num_processes
logger.create(args.experiment_name+'_multi', args.run_name+'_multi')
def run(self):
logger.info(f"🚀 Starting Multi-Process Evaluation with {self.num_processes} processes.")
logger.info(f"🔢 Seeds: {self.seeds}, Frictions: {self.frictions}, Base masses: {self.base_masses}")
process_args = list(product(self.seeds, self.base_masses, self.frictions))
ctx = multiprocessing.get_context('spawn')
worker_func = functools.partial(run_single_process, self.args)
result_log_dirs = []
with ctx.Pool(processes=self.num_processes) as pool:
iterator = pool.imap_unordered(worker_func, process_args)
for log_dir in tqdm(iterator, total=len(process_args), desc="Evaluation"):
result_log_dirs.append(log_dir)
logger.info("✅ Multi-Process Evaluation Completed.")
self.aggregate_results(result_log_dirs)
def aggregate_results(self, log_dirs):
""" Process results.yaml from each log_dir """
all_results = []
all_yaml_paths = []
for path in log_dirs:
yaml_path = Path(path) / "results.yaml"
if not yaml_path.exists():
logger.warning(f"Results file not found: {yaml_path}, skipping.")
continue
with open(yaml_path, 'r') as file:
data = yaml.safe_load(file)
if data:
all_results.append(data)
all_yaml_paths.append(yaml_path)
if not all_results:
logger.error("No results to aggregate.")
return
yaml_paths_str = '\n'.join([str(p) for p in all_yaml_paths])
logger.info(
f"""\n{'='*20} Results Files {'='*20}\n"""
f"""{yaml_paths_str}\n"""
f"""{'='*56}"""
)
value_collections = defaultdict(lambda: defaultdict(list))
for result in all_results:
for goal, metrics in result.items():
if goal != 'summary':
continue
for metric, means in metrics.items():
for mean_name, mean_value in means.items():
value_collections[metric][mean_name].append(float(mean_value.split(' ')[0]))
summary = {}
for metric, means in value_collections.items():
summary[metric] = {}
for mean_name, values in means.items():
summary[metric][mean_name] = f"{float(np.mean(values)):.4f} ± {float(np.std(values)):.4f}"
save_path = logger.log_dir / "aggregated_results.yaml"
with open(save_path, 'w') as file:
yaml.dump(summary, file, allow_unicode=True)
logger.info("✅ Aggregated execution finished.")
logger.info(f"📁 Aggregated results saved to: {save_path}")
logger.info(
f"""\n{'='*20} Multi-Run Summary {'='*20}\n"""
f"""{yaml.dump(summary, allow_unicode=True)}"""
f"""{'='*60}"""
)

1
robogauge/tasks/simulator/__init__.py
View File

@@ -1,2 +1,3 @@
from .mujoco_simulator import MujocoSimulator
from .mujoco_config import MujocoConfig
from .sim_data import SimData

16
robogauge/tasks/simulator/mujoco_config.py
View File

@@ -27,3 +27,19 @@ class MujocoConfig(Config):
video_fps = 30
height = 480
width = 640
class domain_rand:
# With randomization
action_delay = True # [0, control_dt]
# Setup by config file, ensure evaluation coverage
base_mass = 0.0 # [kg], {-1, 0, 1, 2, 3}
friction = 1.0 # [N.s/m], {0.4, 0.7, 1.0, 1.3, 1.6}
class noise:
# Uniform noise
enabled = True
lin_vel = 0.05 # [m/s]
ang_vel = 0.8 # [rad/s]
joint_pos = 0.01 # [rad]
joint_vel = 3.0 # [rad/s]

25
robogauge/tasks/simulator/mujoco_simulator.py
View File

@@ -86,14 +86,29 @@ class MujocoSimulator:
self.mj_model.opt.timestep = self.cfg.physics.simulation_dt
self.sim_dt = self.cfg.physics.simulation_dt
self.mj_data.qpos[7:] = default_dof_pos
# Domain randomization: base mass
base_body_name = f'{Path(self.robot_xml).stem}/base_link'
body_id = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, base_body_name)
assert body_id != -1, f"Body '{base_body_name}' not found in the model."
if self.cfg.domain_rand.base_mass != 0.0:
original_mass = self.mj_model.body_mass[body_id]
new_mass = max(0.01, original_mass + self.cfg.domain_rand.base_mass)
self.mj_model.body_mass[body_id] = new_mass
logger.info(f"Randomized base mass: {original_mass:.3f} -> {new_mass:.3f} kg")
# Domain randomization: friction
if self.cfg.domain_rand.friction != 1.0:
for i in range(self.mj_model.ngeom):
# Both change robot friction and terrain friction, usually robot friction < 1.0
# Mujoco friction calculation takes the *max* between two contacting geoms
geom_friction = self.mj_model.geom_friction[i]
geom_friction[0] *= self.cfg.domain_rand.friction
self.mj_model.geom_friction[i] = geom_friction
logger.info(f"Scaled geom friction by factor: {self.cfg.domain_rand.friction:.3f}")
mujoco.mj_forward(self.mj_model, self.mj_data)
# Setup offscreen camera
base_body_name = f'{Path(self.robot_xml).stem}/base_link'
body_id = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, base_body_name)
if body_id == -1:
body_id = 1
logger.warning(f"Body '{base_body_name}' not found, tracking body ID 1 instead.")
self.offscreen_cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
self.offscreen_cam.trackbodyid = body_id
self.offscreen_cam.distance = self.cfg.viewer.camera_distance

26
robogauge/utils/helpers.py
View File

@@ -38,6 +38,21 @@ def class_to_dict(obj) -> dict:
result[key] = element
return result
def set_seed(seed: int):
import os
import torch
import random
import numpy as np
assert seed >= 0, "Seed must be non-negative."
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def str2bool(v):
if v.lower() in ('yes', 'true', 't', 'y', '1'):
return True
@@ -48,12 +63,21 @@ def str2bool(v):
def parse_args():
parser = ArgumentParser()
parameters = [
# Single run parameters
{"name": "--task-name", "type": str, "default": "base", "help": "Name of the task to run."},
{"name": "--experiment-name", "type": str, "help": "Name of the experiment to run."},
{"name": "--run-name", "type": str, "default": "run1", "help": "Name of the run."},
{"name": "--run-name", "type": str, "default": "run", "help": "Name of the run."},
{"name": "--model-path", "type": str, "help": "Path to the model file."},
{"name": "--headless", "action": "store_true", "default": False, "help": "Run in headless mode."},
{"name": "--save-video", "action": "store_true", "default": False, "help": "Save video output."},
{"name": "--seed", "type": int, "default": 42, "help": "Random seed."},
# Multiprocessing parameters, with different seeds
{"name": "--multi", "action": "store_true", "default": False, "help": "Enable multiprocessing."},
{"name": "--num-processes", "type": int, "default": 2, "help": "Number of parallel processes."},
{"name": "--seeds", "type": int, "nargs": "+", "default": [0], "help": "List of random seeds for multiple runs."},
{"name": "--base-masses", "type": float, "nargs": "+", "default": [-1, 0, 1], "help": "List of base masses for the model."},
{"name": "--frictions", "type": float, "nargs": "+", "default": [0.5, 1.0, 1.5], "help": "List of friction coefficients for the model."}
]
for param in parameters:
parser.add_argument(param['name'], **{k: v for k, v in param.items() if k != 'name'})

14
robogauge/utils/task_register.py
View File

@@ -7,7 +7,9 @@
@Blog : https://wty-yy.github.io/
@Desc : Task Registration Utility
'''
from robogauge.utils.helpers import parse_args
from robogauge import ROBOGAUGE_ROOT_DIR
from robogauge.utils.logger import logger
from robogauge.utils.helpers import parse_args, set_seed
class TaskRegister():
def __init__(self):
@@ -29,13 +31,13 @@ class TaskRegister():
def get_cfgs(self, name):
if name not in self.sim_cfgs:
raise ValueError(f"Task '{name}' is not registered.")
raise ValueError(f"Task '{name}' is not registered, checkout '{ROBOGAUGE_ROOT_DIR}/robogauge/tasks/__init__.py'.")
sim_cfg = self.sim_cfgs[name]
gauger_cfg = self.gauger_cfgs[name]
robot_cfg = self.robot_cfgs[name]
return sim_cfg, gauger_cfg, robot_cfg
def make_pipeline(self, args=None, sim_cfg=None, gauger_cfg=None, robot_cfg=None):
def make_pipeline(self, args=None, sim_cfg=None, gauger_cfg=None, robot_cfg=None, create_logger=True):
if args is None:
args = parse_args()
default_cfgs = self.get_cfgs(args.task_name)
@@ -48,7 +50,11 @@ class TaskRegister():
if args is not None:
self.update_args_to_cfg(sim_cfg, gauger_cfg, robot_cfg, args)
pipeline_class = self.get_pipeline_class(args.task_name)
return pipeline_class(args.run_name, sim_cfg, robot_cfg, gauger_cfg)
set_seed(args.seed)
run_name = args.run_name + f'_{args.seed}'
if create_logger:
logger.create(args.experiment_name, run_name)
return pipeline_class(run_name, sim_cfg, robot_cfg, gauger_cfg)
def update_args_to_cfg(self, sim_cfg, gauger_cfg, robot_cfg, args):
if args.model_path is not None: