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

v0.1.5

Browse Source
This commit is contained in:
wty-yy
2025-12-01 18:01:12 +08:00
parent a91ea5b62e
commit ad907089a8
24 changed files with 324 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
robogauge/__init__.py
View File

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

4
robogauge/scripts/run.py
View File

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

4
robogauge/tasks/__init__.py
View File

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

14
robogauge/tasks/custom/go2_flat_task.py Normal file
View File

@@ -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

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

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

17
robogauge/tasks/gauge/base_gauge_config.py
View File

@@ -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

26
robogauge/tasks/gauge/gauge_configs/flat_gauge_config.py
View File

@@ -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

4
robogauge/tasks/gauge/goal_data.py
View File

@@ -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:

42
robogauge/tasks/gauge/goals/base_goal.py
View File

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

17
robogauge/tasks/gauge/goals/velocity_goals.py
View File

@@ -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}"

43
robogauge/tasks/gauge/metrics/__init__.py
View File

@@ -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

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

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

6
robogauge/tasks/robots/base_robot.py
View File

@@ -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:

5
robogauge/tasks/robots/base_robot_config.py
View File

@@ -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

5
robogauge/tasks/robots/go2/go2_config.py
View File

@@ -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

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

@@ -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

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

@@ -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

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

@@ -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:

5
robogauge/utils/helpers.py
View File

@@ -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

16
robogauge/utils/logger.py
View File

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

13
robogauge/utils/measure.py Normal file
View File

@@ -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

14
robogauge/utils/task_register.py
View File

@@ -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: