From ad907089a8a555db813f8e70d624ed4dbb7d668b Mon Sep 17 00:00:00 2001
From: wty-yy <993660140@qq.com>
Date: Mon, 1 Dec 2025 18:01:12 +0800
Subject: [PATCH] v0.1.5
---
UPDATE.md | 6 +
resources/robots/go2/go2.xml | 6 +-
robogauge/__init__.py | 1 +
robogauge/scripts/run.py | 4 +-
robogauge/tasks/__init__.py | 4 +-
robogauge/tasks/custom/go2_flat_task.py | 14 +++
robogauge/tasks/gauge/base_gauge.py | 67 ++++++++---
robogauge/tasks/gauge/base_gauge_config.py | 17 +--
.../gauge/gauge_configs/flat_gauge_config.py | 26 +++--
robogauge/tasks/gauge/goal_data.py | 4 +-
robogauge/tasks/gauge/goals/base_goal.py | 42 ++++++-
robogauge/tasks/gauge/goals/velocity_goals.py | 17 ++-
robogauge/tasks/gauge/metrics/__init__.py | 43 +++++--
robogauge/tasks/pipeline/base_pipeline.py | 19 +++-
robogauge/tasks/robots/base_robot.py | 6 +-
robogauge/tasks/robots/base_robot_config.py | 5 +-
robogauge/tasks/robots/go2/go2_config.py | 5 +-
robogauge/tasks/simulator/mujoco_config.py | 3 +
robogauge/tasks/simulator/mujoco_simulator.py | 106 +++++++++++-------
robogauge/tasks/simulator/sim_data.py | 1 +
robogauge/utils/helpers.py | 5 +-
robogauge/utils/logger.py | 16 ++-
robogauge/utils/measure.py | 13 +++
robogauge/utils/task_register.py | 14 ++-
24 files changed, 324 insertions(+), 120 deletions(-)
create mode 100644 robogauge/tasks/custom/go2_flat_task.py
create mode 100644 robogauge/utils/measure.py
diff --git a/UPDATE.md b/UPDATE.md
index ee51819..4e0f625 100644
--- a/UPDATE.md
+++ b/UPDATE.md
@@ -1,4 +1,10 @@
# UPDATE
+## 20251201
+### v0.1.5
+1. 加入goals, metrics结果存储
+2. 优化路径存储: `logs/experiment_name/`下分别有两个文件`{time_tag}_{run_name}`存储实验启动的参数, 保存视频; `data/robot/model/goal/{time_tag}_{run_name}`下存储tensorboard
+3. 优化视频存储, 优先记录可视化界面, 否则使用跟随base的相机
+4. 修改`go2.xml`大腿的电机范围不超过base的高度
## 20251130
### v0.1.4
1. 加入velocity_goals中的MaxVelocityGoal, 依次执行每种维度上的极值
diff --git a/resources/robots/go2/go2.xml b/resources/robots/go2/go2.xml
index d558681..77b209a 100644
--- a/resources/robots/go2/go2.xml
+++ b/resources/robots/go2/go2.xml
@@ -13,10 +13,12 @@
-
+
+
-
+
+
diff --git a/robogauge/__init__.py b/robogauge/__init__.py
index 6d75003..469db74 100644
--- a/robogauge/__init__.py
+++ b/robogauge/__init__.py
@@ -2,3 +2,4 @@ from pathlib import Path
__version__ = "0.1.0"
ROBOGAUGE_ROOT_DIR = str(Path(__file__).parents[1])
+ROBOGAUGE_LOGS_DIR = str(Path(ROBOGAUGE_ROOT_DIR) / "logs")
diff --git a/robogauge/scripts/run.py b/robogauge/scripts/run.py
index eb5dc82..599a8f5 100644
--- a/robogauge/scripts/run.py
+++ b/robogauge/scripts/run.py
@@ -17,7 +17,7 @@ from robogauge.utils.logger import logger
if __name__ == '__main__':
args = parse_args()
- logger.create(args.experiment_name)
+ logger.create(args.experiment_name, args.run_name)
logger.info(f"Starting experiment: {args.experiment_name}")
- pipeline: BasePipeline = task_register.make_pipeline(args.task_name, args=args)
+ pipeline: BasePipeline = task_register.make_pipeline(args=args)
pipeline.run()
diff --git a/robogauge/tasks/__init__.py b/robogauge/tasks/__init__.py
index ce85cdb..f732871 100644
--- a/robogauge/tasks/__init__.py
+++ b/robogauge/tasks/__init__.py
@@ -4,5 +4,7 @@ from robogauge.tasks.robots import RobotConfig, Go2Config
from robogauge.tasks.pipeline import BasePipeline
from robogauge.tasks.gauge import BaseGaugeConfig
+from robogauge.tasks.custom.go2_flat_task import Go2FlatGaugeConfig
+
task_register.register('base', BasePipeline, MujocoConfig, BaseGaugeConfig, RobotConfig)
-task_register.register('go2', BasePipeline, MujocoConfig, BaseGaugeConfig, Go2Config)
+task_register.register('go2_flat', BasePipeline, MujocoConfig, Go2FlatGaugeConfig, Go2Config)
diff --git a/robogauge/tasks/custom/go2_flat_task.py b/robogauge/tasks/custom/go2_flat_task.py
new file mode 100644
index 0000000..1bfd07e
--- /dev/null
+++ b/robogauge/tasks/custom/go2_flat_task.py
@@ -0,0 +1,14 @@
+from robogauge.tasks.robots import Go2Config
+from robogauge.tasks.gauge import FlatGaugeConfig
+
+class Go2FlatGaugeConfig(FlatGaugeConfig):
+ class metrics(FlatGaugeConfig.metrics):
+ class dof_limits(FlatGaugeConfig.metrics.dof_limits):
+ enabled = True
+ soft_dof_limit_ratio = 0.7
+ dof_names = ['hip', 'thigh'] # List of DOF names to monitor, None for all
+
+ class goals(FlatGaugeConfig.goals):
+ class max_velocity(FlatGaugeConfig.goals.max_velocity):
+ enabled = True
+ cmd_duration = 5.0
diff --git a/robogauge/tasks/gauge/base_gauge.py b/robogauge/tasks/gauge/base_gauge.py
index c8994ad..7cf6ed6 100644
--- a/robogauge/tasks/gauge/base_gauge.py
+++ b/robogauge/tasks/gauge/base_gauge.py
@@ -7,7 +7,9 @@
@Blog : https://wty-yy.github.io/
@Desc : Base Gauge for Robogauge
'''
+import yaml
from typing import List
+from pathlib import Path
from functools import partial
from robogauge.utils.logger import logger
@@ -19,11 +21,12 @@ from robogauge.tasks.gauge.goal_data import GoalData, VelocityGoal, PositionGoal
from robogauge.tasks.simulator.sim_data import SimData
from robogauge.tasks.gauge.goals import BaseGoal, MaxVelocityGoal
-from robogauge.tasks.gauge.metrics import dof_limits_metric
+from robogauge.tasks.gauge.metrics import *
class BaseGauge:
def __init__(self, cfg: BaseGaugeConfig, robot_cfg: RobotConfig):
self.cfg = cfg
+ self.robot_cfg = robot_cfg
self.goals_cfg = class_to_dict(self.cfg.goals)
self.metrics_cfg = class_to_dict(self.cfg.metrics)
@@ -31,12 +34,10 @@ class BaseGauge:
self.goal_idx = 0
self.goals: List[BaseGoal] = []
self.metrics: List[function] = []
- self.info = {
- 'goal': [],
- 'metric': [],
- }
+ self.info = {'goal': [], 'metric': []}
+ self.results = {} # {'goal/sub_goal': {'metric': result}}
- log_str = "Initialized Gauge with Goals:\n"
+ log_str = "Initialized Gauge with Goals and Metrics:\n"
for name, kwargs in self.goals_cfg.items():
if not kwargs['enabled']: continue
if name == 'max_velocity':
@@ -47,11 +48,17 @@ class BaseGauge:
self.info['goal'].append(name)
for name, enabled in self.metrics_cfg.items():
if not enabled: continue
+ if name in ['metric_dt']: continue
metric_func = eval(f"{name}_metric")
self.metrics.append(partial(metric_func, robot_cfg=robot_cfg, **self.metrics_cfg[name]))
log_str += f" - Metric: {name}\n"
self.info['metric'].append(name)
logger.info(log_str.strip())
+
+ if len(self.goals) == 0:
+ logger.warning("No goals have been configured for the Gauge. Exiting.")
+ else:
+ self.create_new_goal_logger()
def is_reset(self, sim_data: SimData) -> bool:
if self.goal_idx >= len(self.goals):
@@ -60,8 +67,18 @@ class BaseGauge:
def is_done(self) -> bool:
if self.goal_idx >= len(self.goals):
+ self.save_results()
return True
return False
+
+ def create_new_goal_logger(self):
+ """ Create a new logger for new goal to metrics. """
+ if self.goal_idx >= len(self.goals): return
+ logger.create_tensorboard(
+ self.robot_cfg.robot_name,
+ Path(self.robot_cfg.control.model_path).stem,
+ self.goals[self.goal_idx].name
+ )
def get_goal(self, sim_data: SimData) -> GoalData:
# goal = GoalData(
@@ -73,23 +90,37 @@ class BaseGauge:
if self.goal_idx >= len(self.goals):
logger.error("All goals have been exhausted.")
return None
- goal_instance = self.goals[self.goal_idx]
- goal = goal_instance.get_goal(sim_data)
- if goal is None:
+ goal_obj = self.goals[self.goal_idx]
+ goal = goal_obj.get_goal(sim_data)
+
+ if goal is None: # goal obj finished
+ self.results[str(goal_obj)] = goal_obj.sub_goal_mean_metrics
+ self.results[goal_obj.name] = goal_obj.goal_mean_metrics
self.goal_idx += 1
+ self.create_new_goal_logger()
return None
- now_goal_str = str(goal_instance)
- if now_goal_str != self.goal_str:
+ now_goal_str = str(goal_obj)
+ if now_goal_str != self.goal_str: # sub goal changed
+ if self.goal_str != "":
+ self.results[self.goal_str] = goal_obj.sub_goal_mean_metrics
self.goal_str = now_goal_str
- logger.info(f"New Goal [{self.goal_idx+1}/{len(self.goals)}] [{goal_instance.count+1}/{goal_instance.total}]: {self.goal_str}")
+ 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):
- if sim_data.n_step % int(0.1 / sim_data.sim_dt) != 0:
+ if sim_data.n_step % int(self.cfg.metrics.metric_dt / sim_data.sim_dt) != 0:
return
- for i in range(len(sim_data.proprio.joint.force)):
- logger.log(sim_data.proprio.joint.force[i], f'dof/force_{i}', step=sim_data.n_step)
- for metric_func in self.metrics:
- metric_func(sim_data)
-
\ No newline at end of file
+ metrics_results = {}
+ for metric_name, metric_func in zip(self.info['metric'], self.metrics):
+ val = metric_func(sim_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. """
+ save_path = Path(logger.log_dir) / "results.yaml"
+ with open(save_path, 'w') as file:
+ yaml.dump(self.results, file)
+ logger.info(f"Saved metric results to {save_path}")
diff --git a/robogauge/tasks/gauge/base_gauge_config.py b/robogauge/tasks/gauge/base_gauge_config.py
index e3a42b6..b3172eb 100644
--- a/robogauge/tasks/gauge/base_gauge_config.py
+++ b/robogauge/tasks/gauge/base_gauge_config.py
@@ -14,20 +14,21 @@ class BaseGaugeConfig(Config):
class assets:
terrain_xml = '{ROBOGAUGE_ROOT_DIR}/resources/terrains/flat.xml'
- terrain_spawn_xy = [0, 0] # x y [m]
+ terrain_spawn_pos = [0, 0, 0] # x y z [m], robot freejoint spawn position on the terrain
class goals:
class max_velocity: # goal with maximum velocity
enabled = True
- cmd_duration = 3.0 # duration for each velocity command [s]
+ cmd_duration = 5.0 # duration for each velocity command [s]
class metrics:
+ metric_dt = 0.1 # [s], frequency to compute metrics
class dof_limits:
enabled = True
soft_dof_limit_ratio = 0.9
-
- class commands:
- stance = True
- max_lin_vel = True
- diagonal_lin_vel = True
-
+ dof_names = None # List of DOF names to monitor, None for all
+
+ class visualization:
+ enabled = True
+ dof_force = True
+ dof_pos = True
diff --git a/robogauge/tasks/gauge/gauge_configs/flat_gauge_config.py b/robogauge/tasks/gauge/gauge_configs/flat_gauge_config.py
index 567e118..fdc11d5 100644
--- a/robogauge/tasks/gauge/gauge_configs/flat_gauge_config.py
+++ b/robogauge/tasks/gauge/gauge_configs/flat_gauge_config.py
@@ -14,17 +14,21 @@ class FlatGaugeConfig(BaseGaugeConfig):
class assets:
terrain_xml = '{ROBOGAUGE_ROOT_DIR}/resources/terrains/flat.xml'
- terrain_spawn_xy = [0, 0] # x y [m]
-
+ terrain_spawn_pos = [0, 0, 0] # x y z [m], robot freejoint spawn position on the terrain
+
class goals:
- max_velocity = True # goal with maximum velocity
+ class max_velocity: # goal with maximum velocity
+ enabled = True
+ cmd_duration = 5.0 # duration for each velocity command [s]
class metrics:
- dof_limits = True
-
- class commands:
- stance = True
- max_lin_vel = True
- diagonal_lin_vel = True
-
-
+ metric_dt = 0.1 # [s], frequency to compute metrics
+ class dof_limits:
+ enabled = True
+ soft_dof_limit_ratio = 0.9
+ dof_names = None # List of DOF names to monitor, None for all
+
+ class visualization:
+ enabled = True
+ dof_force = True
+ dof_pos = True
diff --git a/robogauge/tasks/gauge/goal_data.py b/robogauge/tasks/gauge/goal_data.py
index 767052e..b06ea67 100644
--- a/robogauge/tasks/gauge/goal_data.py
+++ b/robogauge/tasks/gauge/goal_data.py
@@ -15,8 +15,8 @@ class VelocityGoal:
s = ""
for field in self.__dataclass_fields__:
if getattr(self, field) != 0.0:
- s += f"{field}={getattr(self, field):.1f}, "
- return s[:-2] if s else "stance"
+ s += f"{field}={getattr(self, field):.1f}_"
+ return s[:-1] if s else "stance"
@dataclass
class PositionGoal:
diff --git a/robogauge/tasks/gauge/goals/base_goal.py b/robogauge/tasks/gauge/goals/base_goal.py
index 98178da..89e525b 100644
--- a/robogauge/tasks/gauge/goals/base_goal.py
+++ b/robogauge/tasks/gauge/goals/base_goal.py
@@ -8,12 +8,22 @@
@Desc : Base Goal Class
'''
-from robogauge.tasks.simulator.sim_data import SimData
+from collections import defaultdict
+from robogauge.utils.measure import Average
from robogauge.tasks.gauge.goal_data import GoalData
+from robogauge.tasks.simulator.sim_data import SimData
class BaseGoal:
- count = 0
- total = 0
+ name = 'base_goal'
+
+ def __init__(self):
+ self.count = 0
+ self.total = 0
+ self.sub_name = None
+
+ self._goal_mean_metrics = defaultdict(Average)
+ self.last_sub_name = None
+ self._sub_goal_mean_metrics = defaultdict(Average)
def is_done(self) -> bool:
raise NotImplementedError
@@ -23,3 +33,29 @@ class BaseGoal:
def get_goal(self, sim_data: SimData) -> GoalData:
raise NotImplementedError
+
+ def __repr__(self):
+ if self.sub_name is None:
+ return f"{self.name}"
+ return f"{self.name}/{self.sub_name}"
+
+ def update_metrics(self, metrics: dict):
+ """ Update step metrics for the current goal and sub-goal."""
+ if self.last_sub_name is None or self.last_sub_name != self.sub_name:
+ self.last_sub_name = self.sub_name
+ self._sub_goal_mean_metrics = defaultdict(Average)
+ for metric_name, value in metrics.items():
+ self._goal_mean_metrics[metric_name].update(value)
+ self._sub_goal_mean_metrics[metric_name].update(value)
+
+ @property
+ def goal_mean_metrics(self):
+ """ Get the mean metrics for the current goal. """
+ return {k: float(v.mean) for k, v in self._goal_mean_metrics.items()}
+
+ @property
+ def sub_goal_mean_metrics(self):
+ """ Get the mean metrics for the current sub-goal. """
+ if len(self._sub_goal_mean_metrics) == 0:
+ return {}
+ return {k: float(v.mean) for k, v in self._sub_goal_mean_metrics.items()}
diff --git a/robogauge/tasks/gauge/goals/velocity_goals.py b/robogauge/tasks/gauge/goals/velocity_goals.py
index 276fe25..96e17e9 100644
--- a/robogauge/tasks/gauge/goals/velocity_goals.py
+++ b/robogauge/tasks/gauge/goals/velocity_goals.py
@@ -17,8 +17,11 @@ from robogauge.utils.helpers import class_to_dict
from robogauge.utils.logger import logger
class MaxVelocityGoal(BaseGoal):
+ name = "max_velocity"
+
""" Goal class for maximizing velocity commands. """
def __init__(self, max_velocity: RobotConfig.commands, cmd_duration: float = 5, **kwargs):
+ super().__init__()
kwargs.pop('enabled', None)
if kwargs:
logger.warning(f"Unused kwargs in MaxVelocityGoal: {kwargs}")
@@ -28,14 +31,18 @@ class MaxVelocityGoal(BaseGoal):
self.last_reset_time = 0.0
self.goals = []
- for key, min_max in self.max_velocity.items():
- if min_max is None: continue
- for value in min_max:
+ for key in ['lin_vel_x', 'lin_vel_y', 'lin_vel_z', 'ang_vel_roll', 'ang_vel_pitch', 'ang_vel_yaw']:
+ if self.max_velocity.get(key) is None: continue
+ for value in self.max_velocity[key]:
if value != 0:
self.goals.append(VelocityGoal(**{key: value}))
self.count = 0
self.total = len(self.goals)
+ if len(self.goals) == 0:
+ logger.warning("MaxVelocityGoal initialized with no valid velocity commands.")
+ return
+ self.sub_name = str(self.goals[0])
def is_reset(self, sim_data: SimData) -> bool:
if sim_data.sim_time - self.last_reset_time >= self.cmd_duration:
@@ -48,10 +55,8 @@ class MaxVelocityGoal(BaseGoal):
if self.count >= self.total:
return None
self.current_goal = self.goals[self.count]
+ self.sub_name = str(self.current_goal)
return GoalData(
goal_type='velocity',
velocity_goal=self.current_goal
)
-
- def __repr__(self):
- return f"{self.current_goal}"
diff --git a/robogauge/tasks/gauge/metrics/__init__.py b/robogauge/tasks/gauge/metrics/__init__.py
index 5c89953..16db731 100644
--- a/robogauge/tasks/gauge/metrics/__init__.py
+++ b/robogauge/tasks/gauge/metrics/__init__.py
@@ -1,3 +1,5 @@
+import numpy as np
+
from robogauge.tasks.robots import RobotConfig
from robogauge.tasks.simulator.sim_data import SimData
@@ -18,26 +20,51 @@ def dof_limits_metric(
sim_data: SimData,
robot_cfg: RobotConfig,
soft_dof_limit_ratio: float = 0.9,
+ dof_names: list = None,
**kwargs
) -> float:
""" Metric to log DOF limit violations. """
- mean_value = 0.0
+ values = []
for i in range(len(sim_data.proprio.joint.limits)):
lower_limit = sim_data.proprio.joint.limits[i, 0]
upper_limit = sim_data.proprio.joint.limits[i, 1]
dof_range = upper_limit - lower_limit
- soft_lower_limit = lower_limit + soft_dof_limit_ratio * dof_range
- soft_upper_limit = upper_limit - soft_dof_limit_ratio * dof_range
+ soft_lower_limit = lower_limit + (1 - soft_dof_limit_ratio) * dof_range
+ soft_upper_limit = upper_limit - (1 - soft_dof_limit_ratio) * dof_range
pos = sim_data.proprio.joint.pos[i]
+ dof_name = sim_data.proprio.joint.names[i]
value = 0
if pos < soft_lower_limit:
value = soft_lower_limit - pos
elif pos > soft_upper_limit:
value = pos - soft_upper_limit
value /= dof_range # Normalize by DOF range
- logger.log(value, f'dof_limits/{i}', step=sim_data.n_step)
- mean_value += value
- mean_value /= len(sim_data.proprio.joint.limits)
- logger.log(mean_value, f'dof_limits/mean', step=sim_data.n_step)
- return mean_value
+ logger.log(value, f'dof_limits/{dof_name}', step=sim_data.n_step)
+ if dof_names is not None:
+ for use_name in dof_names:
+ if use_name in dof_name:
+ values.append(value)
+ else:
+ values.append(value)
+ rms_value = 1 - np.sqrt(np.mean(np.square(values)))
+ logger.log(1 - rms_value, f'dof_limits/rms', step=sim_data.n_step)
+ return rms_value
+
+def visualization_metric(
+ sim_data: SimData,
+ robot_cfg: RobotConfig,
+ dof_force: bool = False,
+ dof_pos: bool = False,
+ **kwargs,
+):
+ """ Metric to visualize various robot states in the simulator. """
+ for i in range(len(sim_data.proprio.joint.force)):
+ name = sim_data.proprio.joint.names[i]
+ if dof_force:
+ force = sim_data.proprio.joint.force[i]
+ logger.log(force, f'dof_force/{name}', step=sim_data.n_step)
+ if dof_pos:
+ pos = sim_data.proprio.joint.pos[i]
+ logger.log(pos, f'dof_pos/{name}', step=sim_data.n_step)
+ return 0.0
diff --git a/robogauge/tasks/pipeline/base_pipeline.py b/robogauge/tasks/pipeline/base_pipeline.py
index efb10e7..df41c88 100644
--- a/robogauge/tasks/pipeline/base_pipeline.py
+++ b/robogauge/tasks/pipeline/base_pipeline.py
@@ -7,11 +7,14 @@
@Blog : https://wty-yy.github.io/
@Desc : Base Pipeline for Robogauge
'''
-import traceback
+import yaml
+from pathlib import Path
+
from robogauge.utils.logger import logger
from robogauge.tasks.simulator import MujocoSimulator, MujocoConfig
from robogauge.tasks.robots import BaseRobot, RobotConfig, Go2Config, Go2
from robogauge.tasks.gauge import BaseGauge, BaseGaugeConfig
+from robogauge.utils.helpers import class_to_dict
class BasePipeline:
def __init__(self,
@@ -29,13 +32,20 @@ class BasePipeline:
self.robot: BaseRobot = eval(robot_cfg.robot_class)(robot_cfg)
self.gauge: BaseGauge = eval(gauge_cfg.gauge_class)(gauge_cfg, robot_cfg)
+ # save configs
+ cfg = {}
+ for name in ['simulator_cfg', 'robot_cfg', 'gauge_cfg']:
+ obj = getattr(self, name)
+ obj_dict = class_to_dict(obj)
+ cfg.update({name: obj_dict})
+ with open(Path(logger.log_dir) / "configs.yaml", 'w') as file:
+ yaml.dump(cfg, file)
+
def load(self):
- logger.create_tensorboard(self.run_name)
self.sim.load(
self.gauge_cfg.assets.terrain_xml,
self.robot_cfg.assets.robot_xml,
- self.gauge_cfg.assets.terrain_spawn_xy,
- self.robot_cfg.assets.robot_spawn_height,
+ self.gauge_cfg.assets.terrain_spawn_pos,
self.robot_cfg.control.default_dof_pos
)
@@ -61,5 +71,6 @@ class BasePipeline:
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}")
diff --git a/robogauge/tasks/robots/base_robot.py b/robogauge/tasks/robots/base_robot.py
index 9eaf936..f8abd34 100644
--- a/robogauge/tasks/robots/base_robot.py
+++ b/robogauge/tasks/robots/base_robot.py
@@ -25,9 +25,9 @@ class BaseRobot:
self.control_type = cfg.control.control_type
self.p_gains = np.array(cfg.control.p_gains)
self.d_gains = np.array(cfg.control.d_gains)
- script_model_path = parse_path(cfg.control.torch_script_model_path)
- logger.info(f"Loading robot model from '{script_model_path}'")
- self.model = torch.jit.load(script_model_path).to(self.device)
+ model_path = parse_path(cfg.control.model_path)
+ logger.info(f"Loading robot model from '{model_path}'")
+ self.model = torch.jit.load(model_path).to(self.device)
self.model.eval()
def build_observation(self, sim_data: SimData, goal_data: GoalData) -> np.ndarray:
diff --git a/robogauge/tasks/robots/base_robot_config.py b/robogauge/tasks/robots/base_robot_config.py
index adf52d5..ccf4901 100644
--- a/robogauge/tasks/robots/base_robot_config.py
+++ b/robogauge/tasks/robots/base_robot_config.py
@@ -12,15 +12,16 @@ from typing import Optional, List
from robogauge.utils.config import Config
class RobotConfig(Config):
+ robot_name = 'base_robot'
robot_class = 'BaseRobot'
class assets:
robot_xml = "{ROBOGAUGE_ROOT_DIR}/resources/robots/go2/go2.xml"
- robot_spawn_height = 0.1 # z [m]
class control:
device = 'cpu'
- torch_script_model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_cts_83501.pt"
+ # torch script model path
+ model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_cts_83501.pt"
control_dt = 0.02 # 50 Hz
control_type = 'P' # Position control
diff --git a/robogauge/tasks/robots/go2/go2_config.py b/robogauge/tasks/robots/go2/go2_config.py
index 210601b..84c7a97 100644
--- a/robogauge/tasks/robots/go2/go2_config.py
+++ b/robogauge/tasks/robots/go2/go2_config.py
@@ -11,6 +11,7 @@ from typing_extensions import Literal
from robogauge.tasks.robots import RobotConfig
class Go2Config(RobotConfig):
+ robot_name = 'go2'
robot_class = 'Go2'
class assets:
@@ -19,8 +20,8 @@ class Go2Config(RobotConfig):
class control(RobotConfig.control):
device = 'cpu'
- torch_script_model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_cts_83501.pt"
- # torch_script_model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_cts_cmd-1,1_38k.pt"
+ model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_cts_83501.pt"
+ # model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_cts_cmd-1,1_38k.pt"
control_dt = 0.02 # 50 Hz
control_type = 'P' # Position control
diff --git a/robogauge/tasks/simulator/mujoco_config.py b/robogauge/tasks/simulator/mujoco_config.py
index f5bc7c3..a96fbb5 100644
--- a/robogauge/tasks/simulator/mujoco_config.py
+++ b/robogauge/tasks/simulator/mujoco_config.py
@@ -18,6 +18,9 @@ class MujocoConfig(Config):
class viewer:
headless = False
block_rendering = True # Whether to block rendering in the viewer loop.
+ camera_distance = 2.0
+ camera_elevation = -20.0
+ camera_azimuth = 60.0
class render:
save_video = False
diff --git a/robogauge/tasks/simulator/mujoco_simulator.py b/robogauge/tasks/simulator/mujoco_simulator.py
index a27018b..f9ba215 100644
--- a/robogauge/tasks/simulator/mujoco_simulator.py
+++ b/robogauge/tasks/simulator/mujoco_simulator.py
@@ -15,6 +15,7 @@ import re
import time
import imageio
import numpy as np
+from pathlib import Path
from typing import Literal
from robogauge.utils.logger import logger
@@ -30,10 +31,11 @@ class MujocoSimulator:
self.cfg = sim_cfg
self.terrain_xml = None
self.robot_xml = None
- self.terrain_spawn_xy = None
+ self.terrain_spawn_pos = None
self.robot_spawn_height = None
self.default_dof_pos = None
self.viewer = None
+ self.offscreen_cam = mujoco.MjvCamera()
self.renderer = None
self.vid_writer = None
self.vid_count = 0
@@ -45,8 +47,7 @@ class MujocoSimulator:
self,
terrain_xml: str = None,
robot_xml: str = None,
- terrain_spawn_xy: list = None,
- robot_spawn_height: float = None,
+ terrain_spawn_pos: list = None,
default_dof_pos: list = None,
):
""" Load terrain and robot into the simulator, support re-loading. """
@@ -54,22 +55,20 @@ class MujocoSimulator:
self.terrain_xml = parse_path(terrain_xml)
if robot_xml is not None:
self.robot_xml = parse_path(robot_xml)
- if terrain_spawn_xy is not None:
- self.terrain_spawn_xy = terrain_spawn_xy
- if robot_spawn_height is not None:
- self.robot_spawn_height = robot_spawn_height
+ if terrain_spawn_pos is not None:
+ self.terrain_spawn_pos = terrain_spawn_pos
if default_dof_pos is not None:
self.default_dof_pos = default_dof_pos
terrain_xml = self.terrain_xml
robot_xml = self.robot_xml
- terrain_spawn_xy = self.terrain_spawn_xy
- robot_spawn_height = self.robot_spawn_height
+ terrain_spawn_pos = self.terrain_spawn_pos
if terrain_xml is None or robot_xml is None:
raise ValueError("Terrain and robot XML paths must be provided.")
if default_dof_pos is None:
raise ValueError("Default DOF positions must be provided.")
+ # Create MJCF models
robot_mjcf = mjcf.from_path(robot_xml)
terrain_mjcf = mjcf.from_path(terrain_xml)
for j in robot_mjcf.find_all('joint'):
@@ -77,10 +76,10 @@ class MujocoSimulator:
j.remove()
attachment_frame = terrain_mjcf.attach(robot_mjcf)
attachment_frame.add('freejoint')
- attachment_frame.pos = [*terrain_spawn_xy, robot_spawn_height]
+ attachment_frame.pos = terrain_spawn_pos
- if self.viewer is not None:
- self.close_viewer()
+ self.close_viewer()
+ self.close_video_writer()
self.mj_physics = mjcf.Physics.from_mjcf_model(terrain_mjcf)
self.mj_model = self.mj_physics.model.ptr
self.mj_data = self.mj_physics.data.ptr
@@ -89,38 +88,59 @@ class MujocoSimulator:
self.mj_data.qpos[7:] = default_dof_pos
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
+ self.offscreen_cam.elevation = self.cfg.viewer.camera_elevation
+ self.offscreen_cam.azimuth = self.cfg.viewer.camera_azimuth
+ self.offscreen_cam.lookat = np.array([0.0, 0.0, 0.0])
+
+ # Setup viewer
self.headless = self.cfg.viewer.headless
- if self.cfg.render.save_video and self.headless:
- logger.warning("Cannot save video in headless mode, disabling video saving.")
- self.cfg.render.save_video = False
if not self.headless:
self.viewer = mujoco.viewer.launch_passive(
self.mj_model, self.mj_data, key_callback=self.key_callback
)
- self.last_render_time = time.time()
- if self.cfg.render.save_video:
- self.renderer = mujoco.Renderer(
- self.mj_model, height=self.cfg.render.height,
- width=self.cfg.render.width
- )
- vid_dir = logger.log_dir / "videos"
- vid_dir.mkdir(parents=True, exist_ok=True)
- vid_path = str(vid_dir / f"sim_video_{self.vid_count:03d}.mp4")
- self.vid_writer = imageio.get_writer(
- vid_path,
- fps=self.cfg.render.video_fps,
- )
- self.vid_frame_skip = int(1 / (self.cfg.render.video_fps * self.sim_dt * 2))
- logger.info(f"Saving simulation video to: {vid_path}")
- self.vid_count += 1
+ # set viewer.camera to follow robot
+ self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
+ self.viewer.cam.trackbodyid = body_id
+ self.viewer.cam.distance = self.cfg.viewer.camera_distance
+ self.viewer.cam.elevation = self.cfg.viewer.camera_elevation
+ self.viewer.cam.azimuth = self.cfg.viewer.camera_azimuth
+ self.last_render_time = time.time()
+
+ # Setup video writer
+ if self.cfg.render.save_video:
+ self.renderer = mujoco.Renderer(
+ self.mj_model, height=self.cfg.render.height,
+ width=self.cfg.render.width
+ )
+
+ vid_dir = logger.log_dir / "videos"
+ vid_dir.mkdir(parents=True, exist_ok=True)
+ vid_path = str(vid_dir / f"sim_video_{self.vid_count:03d}.mp4")
+ self.vid_writer = imageio.get_writer(
+ vid_path,
+ fps=self.cfg.render.video_fps,
+ )
+ self.vid_frame_skip = int(1 / (self.cfg.render.video_fps * self.sim_dt * 2))
+ logger.info(f"Simulation video saved at: {vid_path}")
+ self.vid_count += 1
+
+ # Initialize simulation state
self._pause = False
self.n_step = 0
self.sim_time = 0.0
self.load_dof_limits()
self.preload_sensors()
-
# Robot controller placeholders
self.action = None
self.p_gains = None
@@ -138,29 +158,36 @@ class MujocoSimulator:
time.sleep(0.1)
self.update_torque()
self.mj_physics.step()
+
+ # Viewer sync
if self.viewer is not None:
if self.viewer.is_running():
+ self.viewer.sync()
time_untile_next_render = self.cfg.physics.simulation_dt - (
time.time() - self.last_render_time
)
if time_untile_next_render > 0:
time.sleep(time_untile_next_render)
- self.viewer.sync()
- if self.vid_writer is not None and self.n_step % self.vid_frame_skip == 0:
- self.renderer.update_scene(self.mj_data, camera=self.viewer.cam)
- frame = self.renderer.render()
- self.vid_writer.append_data(frame)
self.last_render_time = time.time()
else:
- logger.warning("Viewer closed by user, stop video recording.")
+ logger.warning("Viewer closed by user.")
self.close_viewer()
+ # Video recording
+ if self.vid_writer is not None and self.n_step % self.vid_frame_skip == 0:
+ render_cam = self.viewer.cam if self.viewer is not None else self.offscreen_cam
+ # mujoco.mjv_updateCamera(render_cam)
+ self.renderer.update_scene(self.mj_data, camera=render_cam)
+ frame = self.renderer.render()
+ self.vid_writer.append_data(frame)
+
self.proprio = proprio = RobotProprioception(
joint=JointState(
pos=self.get_sensor_data('joint_pos'),
vel=self.get_sensor_data('joint_vel'),
force=self.get_sensor_data('joint_eff'),
limits=self.dof_limits,
+ names=self.dof_names,
),
imu=IMUState(
pos=self.get_sensor_data('imu_pos'),
@@ -233,6 +260,9 @@ class MujocoSimulator:
self.viewer.close()
self.viewer = None
logger.info("Closing viewer.")
+
+ def close_video_writer(self):
+ """ Close the video writer if exists. """
if self.vid_writer is not None:
self.vid_writer.close()
self.vid_writer = None
diff --git a/robogauge/tasks/simulator/sim_data.py b/robogauge/tasks/simulator/sim_data.py
index 964fce3..d3b0893 100644
--- a/robogauge/tasks/simulator/sim_data.py
+++ b/robogauge/tasks/simulator/sim_data.py
@@ -7,6 +7,7 @@ class JointState:
vel: np.ndarray # [rad/s] shape (n_dof,)
force: np.ndarray # [N*m] shape (n_dof,)
limits: np.ndarray # [rad] shape (n_dof, 2), lower and upper limits
+ names: list # list of joint names
@dataclass
class BaseState:
diff --git a/robogauge/utils/helpers.py b/robogauge/utils/helpers.py
index a096795..dacb78b 100644
--- a/robogauge/utils/helpers.py
+++ b/robogauge/utils/helpers.py
@@ -10,6 +10,7 @@
- Class to dict conversion
- Path parsing
'''
+import yaml
from argparse import ArgumentParser
from pathlib import Path
from robogauge import ROBOGAUGE_ROOT_DIR
@@ -49,6 +50,8 @@ def parse_args():
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": "--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."},
]
@@ -56,5 +59,5 @@ def parse_args():
parser.add_argument(param['name'], **{k: v for k, v in param.items() if k != 'name'})
args = parser.parse_args()
if args.experiment_name is None:
- args.experiment_name = f"{args.task_name}_exp"
+ args.experiment_name = f"exp"
return args
diff --git a/robogauge/utils/logger.py b/robogauge/utils/logger.py
index d46e57f..07a561c 100644
--- a/robogauge/utils/logger.py
+++ b/robogauge/utils/logger.py
@@ -13,7 +13,7 @@
import time
import logging
from pathlib import Path
-from robogauge import ROBOGAUGE_ROOT_DIR
+from robogauge import ROBOGAUGE_LOGS_DIR
from torch.utils.tensorboard import SummaryWriter
class LogColor:
@@ -56,6 +56,7 @@ class Logger:
def create(self,
experiment_name,
+ run_name,
console_output=True, color_output=True,
log_level=logging.DEBUG, save_file_mode='a'
):
@@ -75,6 +76,9 @@ class Logger:
self.logger = logging.getLogger(experiment_name + "_logger")
self.logger.setLevel(log_level)
self.logger.propagate = False
+ self.time_tag = time.strftime("%Y%m%d-%H-%M-%S")
+ self.tag = f"{self.time_tag}_{run_name}"
+ self.experiment_name = experiment_name
console_formatter = ColorFormatter( # console output format
fmt="%(asctime)s - %(color_level)s - %(filename)s:%(lineno)d - %(message)s",
@@ -91,19 +95,21 @@ class Logger:
sh.setFormatter(console_formatter)
self.logger.addHandler(sh)
- self.log_dir = Path(ROBOGAUGE_ROOT_DIR) / "logs" / experiment_name / time.strftime("%Y%m%d-%H-%M-%S")
+ self.log_dir = Path(ROBOGAUGE_LOGS_DIR) / experiment_name / self.tag
self.log_dir.mkdir(parents=True, exist_ok=True)
path_log_file = self.log_dir / "stdout.log"
if path_log_file:
fh = logging.FileHandler(path_log_file, mode=save_file_mode, encoding='utf-8')
fh.setFormatter(file_formatter)
self.logger.addHandler(fh)
+ self.info(f"Logs saved at: {path_log_file}")
- def create_tensorboard(self, run_name: str):
+ def create_tensorboard(self, robot_name: str, model_name: str, goal_name: str):
if self.writer is not None:
self.writer.close()
- self.writer = SummaryWriter(str(self.log_dir / run_name))
- self.info(f"Tensorboard writer created at: {self.log_dir / run_name}")
+ data_path = Path(ROBOGAUGE_LOGS_DIR) / self.experiment_name / 'data' / robot_name / model_name / goal_name / self.tag
+ self.writer = SummaryWriter(str(data_path))
+ self.info(f"Tensorboard writer created at: {data_path}")
def debug(self, msg, *args, **kwargs):
self.logger.debug(msg, *args, **kwargs, stacklevel=2)
diff --git a/robogauge/utils/measure.py b/robogauge/utils/measure.py
new file mode 100644
index 0000000..3cd8905
--- /dev/null
+++ b/robogauge/utils/measure.py
@@ -0,0 +1,13 @@
+class Average:
+ def __init__(self):
+ self.avg = 0.0
+ self.count = 0
+
+ def update(self, value: float):
+ self.avg += (value - self.avg) / (self.count + 1)
+ self.count += 1
+ return self.avg
+
+ @property
+ def mean(self):
+ return self.avg
diff --git a/robogauge/utils/task_register.py b/robogauge/utils/task_register.py
index 60dc0c7..78f0652 100644
--- a/robogauge/utils/task_register.py
+++ b/robogauge/utils/task_register.py
@@ -7,6 +7,8 @@
@Blog : https://wty-yy.github.io/
@Desc : Task Registration Utility
'''
+from robogauge.utils.helpers import parse_args
+
class TaskRegister():
def __init__(self):
self.pipeline_classes = {}
@@ -33,8 +35,10 @@ class TaskRegister():
robot_cfg = self.robot_cfgs[name]
return sim_cfg, gauger_cfg, robot_cfg
- def make_pipeline(self, name, args=None, sim_cfg=None, gauger_cfg=None, robot_cfg=None, run_name='0'):
- default_cfgs = self.get_cfgs(name)
+ def make_pipeline(self, args=None, sim_cfg=None, gauger_cfg=None, robot_cfg=None):
+ if args is None:
+ args = parse_args()
+ default_cfgs = self.get_cfgs(args.task_name)
if sim_cfg is None:
sim_cfg = default_cfgs[0]
if gauger_cfg is None:
@@ -43,10 +47,12 @@ class TaskRegister():
robot_cfg = default_cfgs[2]
if args is not None:
self.update_args_to_cfg(sim_cfg, gauger_cfg, robot_cfg, args)
- pipeline_class = self.get_pipeline_class(name)
- return pipeline_class(run_name, sim_cfg, robot_cfg, gauger_cfg)
+ pipeline_class = self.get_pipeline_class(args.task_name)
+ return pipeline_class(args.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:
+ robot_cfg.control.model_path = args.model_path
if args.headless is not None:
sim_cfg.viewer.headless = args.headless
if args.save_video is not None: