v0.1.5

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, 依次执行每种维度上的极值

resources/robots/go2/go2.xml

@@ -13,10 +13,12 @@
       </default>
       <default class="hip">
         <default class="front_hip">
-          <joint range="-1.5708 3.4907"/>
+          <!-- <joint range="-1.5708 3.4907"/> -->
+          <joint range="-1.5708 1.5708"/>
         </default>
         <default class="back_hip">
-          <joint range="-0.5236 4.5379"/>
+          <!-- <joint range="-0.5236 4.5379"/> -->
+          <joint range="-0.5236 1.5708"/>
         </default>
       </default>
       <default class="knee">

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")

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()

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)

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

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 = {
+        self.info = {'goal': [], 'metric': []}
-            'goal': [],
+        self.results = {}  # {'goal/sub_goal': {'metric': result}}
-            'metric': [],
-        }
 
-        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,12 +48,18 @@ 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):
             return False
@@ -60,9 +67,19 @@ 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(
         #     goal_type='velocity',
@@ -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_obj = self.goals[self.goal_idx]
-        goal = goal_instance.get_goal(sim_data)
+        goal = goal_obj.get_goal(sim_data)
-        if goal is None:
+
+        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)
+        now_goal_str = str(goal_obj)
-        if now_goal_str != self.goal_str:
+        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)):
+        metrics_results = {}
-            logger.log(sim_data.proprio.joint.force[i], f'dof/force_{i}', step=sim_data.n_step)
+        for metric_name, metric_func in zip(self.info['metric'], self.metrics):
-        for metric_func in self.metrics:
+            val = metric_func(sim_data)
-            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}")

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
+            dof_names = None  # List of DOF names to monitor, None for all
         
-    class commands:
+        class visualization:
-        stance = True
+            enabled = True
-        max_lin_vel = True
+            dof_force = True
-        diagonal_lin_vel = True
+            dof_pos = True
-    

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
+        metric_dt = 0.1  # [s], frequency to compute metrics
-    
+        class dof_limits:
-    class commands:
+            enabled = True
-        stance = True
+            soft_dof_limit_ratio = 0.9
-        max_lin_vel = True
+            dof_names = None  # List of DOF names to monitor, None for all
-        diagonal_lin_vel = True
-    
         
+        class visualization:
+            enabled = True
+            dof_force = True
+            dof_pos = True

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}, "
+                s += f"{field}={getattr(self, field):.1f}_"
-        return s[:-2] if s else "stance"
+        return s[:-1] if s else "stance"
 
 @dataclass
 class PositionGoal:

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
+    name = 'base_goal'
-    total = 0
+
+    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()}

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():
+        for key in ['lin_vel_x', 'lin_vel_y', 'lin_vel_z', 'ang_vel_roll', 'ang_vel_pitch', 'ang_vel_yaw']:
-            if min_max is None: continue
+            if self.max_velocity.get(key) is None: continue
-            for value in min_max:
+            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}"

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_lower_limit = lower_limit + (1 - soft_dof_limit_ratio) * dof_range
-        soft_upper_limit = upper_limit - 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)
+        logger.log(value, f'dof_limits/{dof_name}', step=sim_data.n_step)
-        mean_value += value
+        if dof_names is not None:
-    mean_value /= len(sim_data.proprio.joint.limits)
+            for use_name in dof_names:
-    logger.log(mean_value, f'dof_limits/mean', step=sim_data.n_step)
+                if use_name in dof_name:
-    return mean_value
+                    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

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.gauge_cfg.assets.terrain_spawn_pos,
-            self.robot_cfg.assets.robot_spawn_height,
             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}")

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)
+        model_path = parse_path(cfg.control.model_path)
-        logger.info(f"Loading robot model from '{script_model_path}'")
+        logger.info(f"Loading robot model from '{model_path}'")
-        self.model = torch.jit.load(script_model_path).to(self.device)
+        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:

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
 

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"
+        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_cmd-1,1_38k.pt"
         control_dt = 0.02  # 50 Hz
         control_type = 'P'  # Position control
 

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

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,
+        terrain_spawn_pos: list = None,
-        robot_spawn_height: float = 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:
+        if terrain_spawn_pos is not None:
-            self.terrain_spawn_xy = terrain_spawn_xy
+            self.terrain_spawn_pos = terrain_spawn_pos
-        if robot_spawn_height is not None:
-            self.robot_spawn_height = robot_spawn_height
         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
+        terrain_spawn_pos = self.terrain_spawn_pos
-        robot_spawn_height = self.robot_spawn_height
         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_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,15 +88,35 @@ 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
             )
+
+            # 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,
@@ -112,15 +131,16 @@ class MujocoSimulator:
                 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}")
+            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

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:

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

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))
+        data_path = Path(ROBOGAUGE_LOGS_DIR) / self.experiment_name / 'data' / robot_name / model_name / goal_name / self.tag
-        self.info(f"Tensorboard writer created at: {self.log_dir / run_name}")
+        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)

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

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'):
+    def make_pipeline(self, args=None, sim_cfg=None, gauger_cfg=None, robot_cfg=None):
-        default_cfgs = self.get_cfgs(name)
+        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)
+        pipeline_class = self.get_pipeline_class(args.task_name)
-        return pipeline_class(run_name, sim_cfg, robot_cfg, gauger_cfg)
+        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: