From 4546be76bf973dc98ea32d5ec852a5e74a253b0d Mon Sep 17 00:00:00 2001 From: wty-yy Date: Sun, 21 Dec 2025 00:16:07 +0800 Subject: [PATCH] v0.1.11 --- .vscode/launch.json | 20 +++ UPDATE.md | 9 ++ robogauge/scripts/run.py | 2 + robogauge/tasks/custom/go2_flat_task.py | 17 ++- robogauge/tasks/gauge/base_gauge.py | 31 +++-- robogauge/tasks/gauge/base_gauge_config.py | 9 ++ .../gauge/gauge_configs/flat_gauge_config.py | 13 +- robogauge/tasks/gauge/goal_data.py | 1 + robogauge/tasks/gauge/goals/__init__.py | 2 +- robogauge/tasks/gauge/goals/base_goal.py | 6 +- robogauge/tasks/gauge/goals/velocity_goals.py | 114 ++++++++++++++++-- robogauge/tasks/pipeline/base_pipeline.py | 30 +++-- robogauge/tasks/simulator/mujoco_simulator.py | 27 +++++ robogauge/utils/math_utils.py | 10 ++ 14 files changed, 246 insertions(+), 45 deletions(-) create mode 100644 .vscode/launch.json diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000..abad633 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,20 @@ +{ + // 使用 IntelliSense 了解相关属性。 + // 悬停以查看现有属性的描述。 + // 欲了解更多信息，请访问: https://go.microsoft.com/fwlink/?linkid=830387 + "version": "0.2.0", + "configurations": [ + { + "name": "Python 调试程序: run.py", + "type": "debugpy", + "request": "launch", + "program": "${workspaceFolder}/robogauge/scripts/run.py", + "args": [ + "--task", "go2_moe_flat", + "--experiment-name", "debug", + "--headless" + ], + "console": "integratedTerminal" + } + ] +} diff --git a/UPDATE.md b/UPDATE.md index 7452410..9456232 100644 --- a/UPDATE.md +++ b/UPDATE.md @@ -1,5 +1,14 @@ # UPDATE TODO: 在模型崩溃时也记录下最后的gauge信息 +## 20251220 +### v0.1.11 +1. 加入`os.environ["OMP_NUM_THREADS"] = "2"; os.environ["MKL_NUM_THREADS"] = "2"`避免并行时cpu线程爆炸, `--multi`模式能稳定提高速度了 +2. 完成target_position_goal, 超参数包含: 目标点位置, 最大线速度角速度, 最长追踪时间, 追踪到达阈值范围; 并在mujoco中绘制红色目标点, result.yaml中记录当前terrain和success +3. 修改goal的reset逻辑 + 旧版: 在sub_goal开始时通过实现类中的get_goal异常返回None判断当前goals全部结束, 并且goals全部结束也不reset环境; + 新版: 加入goal.pre_get_goal, 判断当前系列goals是否全部结束, 并根据sim_data更新当前的sub_goal索引, 实现类中无需考虑异常处理, 并且goals全部结束时判断为change goal执行一次reset环境, 保证新的goal可以直接无缝衔接上 +4. 添加vscode python debug启动配置文件, 支持参数输入调试 +Fix Bugs: 修复sub_goal重名, 日志info不显示的问题 ## 20251218 ### v0.1.10 1. 加入`--write-tensorboard`参数, 默认为`False`即不记录`gauge`的日志信息 diff --git a/robogauge/scripts/run.py b/robogauge/scripts/run.py index 970d0fa..b91d713 100644 --- a/robogauge/scripts/run.py +++ b/robogauge/scripts/run.py @@ -9,6 +9,8 @@ ''' import os os.environ['MUJOCO_GL'] = 'glfw' # avoid mujoco.Renderer EGL context error +os.environ["OMP_NUM_THREADS"] = "1" +os.environ["MKL_NUM_THREADS"] = "1" from robogauge.tasks import * from robogauge.tasks.pipeline.multi_pipeline import MultiPipeline diff --git a/robogauge/tasks/custom/go2_flat_task.py b/robogauge/tasks/custom/go2_flat_task.py index 1047c8a..af53a8e 100644 --- a/robogauge/tasks/custom/go2_flat_task.py +++ b/robogauge/tasks/custom/go2_flat_task.py @@ -29,21 +29,30 @@ class Go2FlatGaugeConfig(FlatGaugeConfig): enabled = True cmd_duration = 6.0 + class target_pos_velocity(FlatGaugeConfig.goals.target_pos_velocity): # goal to reach a target position by velocity command, config target at assets.target_pos + enabled = True + target_pos = [2, 2, 0] # x y z [m], target position in the environment, used for target position goal + lin_vel_x = 1.0 # +/- m/s + lin_vel_y = 1.0 # +/- m/s + ang_vel_yaw = 1.5 # +/- rad/s + max_cmd_duration = 10.0 # [s] maximum duration to reach the target position + reach_threshold = 0.1 # [m] distance threshold to consider the target reached + class Go2FlatConfig(Go2Config): class commands(Go2Config.commands): lin_vel_x = [-1.8, 1.8] # min max [m/s] lin_vel_y = [-1.8, 1.8] # min max [m/s] - ang_vel_yaw = [-1.8, 1.8] # min max [rad/s] + ang_vel_yaw = [-2.0, 2.0] # min max [rad/s] class Go2MoEFlatConfig(Go2MoEConfig): class commands(Go2Config.commands): lin_vel_x = [-1.8, 1.8] # min max [m/s] lin_vel_y = [-1.8, 1.8] # min max [m/s] - ang_vel_yaw = [-1.8, 1.8] # min max [rad/s] + ang_vel_yaw = [-2.0, 2.0] # min max [rad/s] class control(Go2Config.control): - # model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_moe_cts_124k.pt" - model_path = "/home/xfy/Coding/kaiwu2025/rob_finals/sim2real/models/v6-2_106503/kaiwu_script_v6-2_106503.pt" + model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_moe_cts_124k.pt" + # model_path = "/home/xfy/Coding/kaiwu2025/rob_finals/sim2real/models/v6-2_106503/kaiwu_script_v6-2_106503.pt" class Go2MoEFlatMujocoConfig(MujocoConfig): class domain_rand(MujocoConfig.domain_rand): diff --git a/robogauge/tasks/gauge/base_gauge.py b/robogauge/tasks/gauge/base_gauge.py index d95a843..364d007 100644 --- a/robogauge/tasks/gauge/base_gauge.py +++ b/robogauge/tasks/gauge/base_gauge.py @@ -22,7 +22,7 @@ from robogauge.tasks.gauge.base_gauge_config import BaseGaugeConfig 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, DiagonalVelocityGoal +from robogauge.tasks.gauge.goals import * from robogauge.tasks.gauge.metrics import * class BaseGauge: @@ -40,7 +40,8 @@ class BaseGauge: self.goals: List[BaseGoal] = [] self.metrics: List[function] = [] self.info = {'goal': [], 'metric': []} - self.results = {} # {'goal/sub_goal': {'metric': result}} + self.results = { + } # {'goal/sub_goal': {'metric': result}} log_str = "Initialized Gauge with Goals and Metrics:\n" for name, kwargs in self.goals_cfg.items(): @@ -51,6 +52,9 @@ class BaseGauge: elif name == 'diagonal_velocity': self.goals.append(DiagonalVelocityGoal(robot_cfg.control.control_dt, robot_cfg.commands, **kwargs)) log_str += f" - Diagonal Velocity Goal: {kwargs}\n" + elif name == 'target_pos_velocity': + self.goals.append(TargetPosVelocityGoal(robot_cfg.control.control_dt,**kwargs)) + log_str += f" - Target Position Velocity Goal: {kwargs}\n" else: raise NotImplementedError(f"Goal '{name}' is not implemented in BaseGauge.") self.info['goal'].append(name) @@ -97,21 +101,23 @@ class BaseGauge: # ang_vel_yaw=-5.0, # ) # ) - if self.goal_idx >= len(self.goals): - logger.error("All goals have been exhausted.") - return None goal_obj = self.goals[self.goal_idx] - goal = goal_obj.get_goal(sim_data) + if goal_obj.pre_get_goal(sim_data): + metrics = goal_obj.goal_mean_metrics + if hasattr(goal_obj, 'success'): # target position goal + metrics['success'] = {'mean': float(goal_obj.success)} + + key = f"{goal_obj.name}" + self.results[key] = metrics - if goal is None: # goal obj finished - self.results[goal_obj.name] = goal_obj.goal_mean_metrics self.goal_idx += 1 self.create_new_goal_logger() return None + goal = goal_obj.get_goal(sim_data) now_goal_str = str(goal_obj) - if now_goal_str != self.goal_str: # sub goal changed - self.goal_str = now_goal_str + if f"{goal_obj.count+1}_{now_goal_str}" != self.goal_str: # sub goal changed + self.goal_str = f"{goal_obj.count+1}_{now_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 @@ -142,9 +148,10 @@ class BaseGauge: self.results['summary'][metric_name][quantile] = f"{mean:.4f} ± {std:.4f}" save_path = Path(logger.log_dir) / "results.yaml" + self.results["terrain"] = f"{self.cfg.assets.terrain_name}" with open(save_path, 'w') as file: - yaml.dump(self.results, file, allow_unicode=True) - yaml_str = yaml.dump(self.results, allow_unicode=True) + yaml.dump(self.results, file, allow_unicode=True, sort_keys=False) + yaml_str = yaml.dump(self.results, allow_unicode=True, sort_keys=False) logger.info( f"""\n{'='*20} Goals and Metrics results {'='*20}\n""" f"""{yaml_str}""" diff --git a/robogauge/tasks/gauge/base_gauge_config.py b/robogauge/tasks/gauge/base_gauge_config.py index 4a2f435..e3c39fa 100644 --- a/robogauge/tasks/gauge/base_gauge_config.py +++ b/robogauge/tasks/gauge/base_gauge_config.py @@ -14,6 +14,7 @@ class BaseGaugeConfig(Config): write_tensorboard = False # Whether to write tensorboard logs class assets: + terrain_name = "flat_0" # {type}_{level} terrain_xml = '{ROBOGAUGE_ROOT_DIR}/resources/terrains/flat.xml' terrain_spawn_pos = [0, 0, 0] # x y z [m], robot freejoint spawn position on the terrain @@ -26,6 +27,14 @@ class BaseGaugeConfig(Config): enabled = True cmd_duration = 6.0 # [s] duration for a pair of diagonal velocity commands + class target_pos_velocity: # goal to reach a target position by velocity command + enabled = True + target_pos = [5, 0, 0] # x y z [m], target position in the environment, used for target position goal + lin_vel_x = 1.0 # +/- m/s + ang_vel_yaw = 1.0 # +/- rad/s + max_cmd_duration = 10.0 # [s] maximum duration to reach the target position + reach_threshold = 0.1 + class metrics: metric_dt = 0.1 # [s] frequency to compute metrics class dof_limits: diff --git a/robogauge/tasks/gauge/gauge_configs/flat_gauge_config.py b/robogauge/tasks/gauge/gauge_configs/flat_gauge_config.py index 55cc2a3..f1aa193 100644 --- a/robogauge/tasks/gauge/gauge_configs/flat_gauge_config.py +++ b/robogauge/tasks/gauge/gauge_configs/flat_gauge_config.py @@ -12,11 +12,12 @@ from robogauge.tasks.gauge.base_gauge_config import BaseGaugeConfig class FlatGaugeConfig(BaseGaugeConfig): gauge_class = 'BaseGauge' - class assets: + class assets(BaseGaugeConfig.assets): + terrain_name = "flat_0" # {type}_{level} terrain_xml = '{ROBOGAUGE_ROOT_DIR}/resources/terrains/flat.xml' terrain_spawn_pos = [0, 0, 0] # x y z [m], robot freejoint spawn position on the terrain - class goals: + class goals(BaseGaugeConfig.goals): class max_velocity: # goal with maximum velocity enabled = True move_duration = 5.0 # [s] duration for each velocity command @@ -27,6 +28,14 @@ class FlatGaugeConfig(BaseGaugeConfig): enabled = True cmd_duration = 6.0 # [s] duration for a pair of diagonal velocity commands + class target_pos_velocity: # goal to reach a target position by velocity command, config target at assets.target_pos + enabled = True + target_pos = [5, 0, 0] # x y z [m], target position in the environment, used for target position goal + lin_vel_x = 1.0 # +/- m/s + ang_vel_yaw = 1.0 # +/- rad/s + max_cmd_duration = 10.0 # [s] maximum duration to reach the target position + reach_threshold = 0.1 + class metrics(BaseGaugeConfig.metrics): metric_dt = 0.1 # [s] frequency to compute metrics class dof_limits: diff --git a/robogauge/tasks/gauge/goal_data.py b/robogauge/tasks/gauge/goal_data.py index 86c516d..7a1b1b3 100644 --- a/robogauge/tasks/gauge/goal_data.py +++ b/robogauge/tasks/gauge/goal_data.py @@ -41,3 +41,4 @@ class GoalData: goal_type: Literal['velocity', 'position'] velocity_goal: Optional[VelocityGoal] = None position_goal: Optional[PositionGoal] = None + visualization_pos: Optional[Tuple[float, float, float]] = None diff --git a/robogauge/tasks/gauge/goals/__init__.py b/robogauge/tasks/gauge/goals/__init__.py index 2e6d1c3..32b6f0b 100644 --- a/robogauge/tasks/gauge/goals/__init__.py +++ b/robogauge/tasks/gauge/goals/__init__.py @@ -1,2 +1,2 @@ from robogauge.tasks.gauge.goals.base_goal import BaseGoal -from robogauge.tasks.gauge.goals.velocity_goals import MaxVelocityGoal, DiagonalVelocityGoal +from robogauge.tasks.gauge.goals.velocity_goals import MaxVelocityGoal, DiagonalVelocityGoal, TargetPosVelocityGoal diff --git a/robogauge/tasks/gauge/goals/base_goal.py b/robogauge/tasks/gauge/goals/base_goal.py index 7a9a940..cdf93d2 100644 --- a/robogauge/tasks/gauge/goals/base_goal.py +++ b/robogauge/tasks/gauge/goals/base_goal.py @@ -17,13 +17,13 @@ class BaseGoal: name = 'base_goal' def __init__(self): - self.count = 0 - self.total = 0 + self.count = 0 # current task index + self.total = 0 # total tasks self.sub_name = None self._goal_mean_metrics = defaultdict(list) - def is_done(self) -> bool: + def pre_get_goal(self) -> bool: raise NotImplementedError def is_reset(self, sim_data: SimData) -> bool: diff --git a/robogauge/tasks/gauge/goals/velocity_goals.py b/robogauge/tasks/gauge/goals/velocity_goals.py index 04b4f94..248ba89 100644 --- a/robogauge/tasks/gauge/goals/velocity_goals.py +++ b/robogauge/tasks/gauge/goals/velocity_goals.py @@ -7,8 +7,9 @@ @Blog : https://wty-yy.github.io/ @Desc : Velocity Goals Implementation ''' +import numpy as np from copy import deepcopy -from typing import Optional +from typing import Optional, List from robogauge.tasks.gauge.goals import BaseGoal from robogauge.tasks.robots import RobotConfig @@ -16,27 +17,48 @@ from robogauge.tasks.simulator.sim_data import SimData from robogauge.tasks.gauge.goal_data import GoalData, VelocityGoal from robogauge.utils.helpers import class_to_dict from robogauge.utils.logger import logger +from robogauge.utils.math_utils import quat_rotate_inverse + +PI = np.pi class BaseVelocityGoal(BaseGoal): name = "base_velocity_goal" def __init__(self, control_dt: float, cmd_duration: float = 5, **kwargs): + """ + Args: + control_dt (float): Control timestep. + cmd_duration (float): Duration for each sub-task (e.g. different velocity commands). + """ super().__init__() self.control_dt = control_dt self.cmd_duration = cmd_duration self.goal_runtime = 0.0 self.first_goal_after_reset = True - self.last_reset_time = 0.0 + self.last_reset_time = -1 self.goals = [] def is_reset(self, sim_data: SimData) -> bool: + if self.last_reset_time == -1: + self.last_reset_time = sim_data.sim_time if sim_data.sim_time - self.last_reset_time >= self.cmd_duration: self.last_reset_time = sim_data.sim_time self.first_goal_after_reset = True return True return False + + def pre_get_goal(self, sim_data: SimData) -> bool: + """ Run before getting the goal + Returns: + bool: whether the goal sequence is done + """ + self.update_runtime_count(sim_data) + return self.count >= self.total def update_runtime_count(self, sim_data: SimData): + """ + Update the goal runtime and count based on the simulation time. + """ if self.first_goal_after_reset: self.last_reset_time = sim_data.sim_time # update last reset time (stance after reset) self.first_goal_after_reset = False @@ -46,7 +68,7 @@ class BaseVelocityGoal(BaseGoal): class MaxVelocityGoal(BaseVelocityGoal): name = "max_velocity" - """ Goal class for maximizing velocity commands. """ + """ Goal class for maximizing velocity commands, maximum 6 sub-tasks. """ def __init__(self, control_dt: float, max_velocity: RobotConfig.commands, @@ -82,9 +104,6 @@ class MaxVelocityGoal(BaseVelocityGoal): self.sub_name = str(self.goals[0]) def get_goal(self, sim_data: SimData) -> Optional[GoalData]: - self.update_runtime_count(sim_data) - if self.count >= self.total: - return None self.current_goal = self.goals[self.count] if self.end_stance and self.goal_runtime - self.count * self.cmd_duration >= self.move_duration: self.current_goal = VelocityGoal() # zero velocity @@ -103,7 +122,7 @@ class DiagonalVelocityGoal(BaseVelocityGoal): cmd_duration: float = 6, **kwargs ): - """ Goal class for diagonal velocity changes. + """ Goal class for diagonal velocity changes, maximum 8 sub-tasks. Args: control_dt (float): Control timestep. max_velocity (RobotConfig.commands): Maximum velocity commands. @@ -114,9 +133,6 @@ class DiagonalVelocityGoal(BaseVelocityGoal): if kwargs: logger.warning(f"Unused kwargs in DiagonalVelocityGoal: {kwargs}") - self.cmd_duration = cmd_duration - self.last_reset_time = 0.0 - self.goals = [] lin_vel_x_vals = max_velocity.lin_vel_x + [0] lin_vel_y_vals = max_velocity.lin_vel_y + [0] @@ -136,9 +152,6 @@ class DiagonalVelocityGoal(BaseVelocityGoal): self.sub_name = str(self.goals[0]) def get_goal(self, sim_data: SimData) -> Optional[GoalData]: - self.update_runtime_count(sim_data) - if self.count >= self.total: - return None self.current_goal = self.goals[self.count] if self.goal_runtime - self.count * self.cmd_duration >= self.cmd_duration / 2: self.current_goal = self.current_goal.invert() @@ -147,3 +160,78 @@ class DiagonalVelocityGoal(BaseVelocityGoal): goal_type='velocity', velocity_goal=self.current_goal ) + +class TargetPosVelocityGoal(BaseVelocityGoal): + name = "target_pos_velocity" + + def __init__(self, + control_dt: float, + target_pos: List[float], + lin_vel_x: float, + lin_vel_y: float, + ang_vel_yaw: float, + max_cmd_duration: float, + reach_threshold: float, + **kwargs + ): + super().__init__(control_dt=control_dt, cmd_duration=max_cmd_duration) + + self.target_pos = target_pos + self.lin_vel_x = lin_vel_x + self.lin_vel_y = lin_vel_y + self.ang_vel_yaw = ang_vel_yaw + self.reach_threshold = reach_threshold + + kwargs.pop('enabled', None) + if kwargs: + logger.warning(f"Unused kwargs in TargetPosVelocity: {kwargs}") + self.sub_name = None + self.count = 0 + self.total = 1 # only one task + self.done = False + self.success = False + + def get_goal(self, sim_data: SimData) -> Optional[GoalData]: + self.update_runtime_count(sim_data) + delta_pos, _, delta_ang = self.get_delta_info(sim_data) + ang_vel_yaw = np.sign(delta_ang) * min(abs(delta_ang) * 2, self.ang_vel_yaw) + self.current_goal = VelocityGoal( + lin_vel_x=self.lin_vel_x if abs(delta_ang) < PI / 4 else 0.0, + lin_vel_y=np.sign(delta_pos[1]) * min(abs(delta_pos[1]), self.lin_vel_y) if abs(delta_ang) >= PI / 4 else 0.0, + ang_vel_yaw=ang_vel_yaw + ) + return GoalData( + goal_type='velocity', + velocity_goal=self.current_goal, + visualization_pos=self.target_pos + ) + + def is_reset(self, sim_data: SimData) -> bool: + _, norm, _ = self.get_delta_info(sim_data) + time_out = super().is_reset(sim_data) + reached = norm < self.reach_threshold + if time_out or reached: + self.done = True + if reached: + self.success = True + return True + return False + + def pre_get_goal(self, sim_data: SimData) -> bool: + """ Run before getting the goal + Returns: + bool: whether the goal is done + """ + return self.count >= self.total or self.done + + def get_delta_info(self, sim_data: SimData) -> List[float]: + current_pos = sim_data.proprio.base.pos + delta_pos = [ + self.target_pos[0] - current_pos[0], + self.target_pos[1] - current_pos[1], + self.target_pos[2] - current_pos[2], + ] + delta_pos = quat_rotate_inverse(sim_data.proprio.base.quat, delta_pos) + norm = np.linalg.norm(delta_pos[:2]) + delta_ang = np.arctan2(delta_pos[1], delta_pos[0]) + return delta_pos, norm, delta_ang diff --git a/robogauge/tasks/pipeline/base_pipeline.py b/robogauge/tasks/pipeline/base_pipeline.py index 9d16472..89aed8b 100644 --- a/robogauge/tasks/pipeline/base_pipeline.py +++ b/robogauge/tasks/pipeline/base_pipeline.py @@ -38,6 +38,7 @@ class BasePipeline: self.robot: BaseRobot = eval(robot_cfg.robot_class)(robot_cfg) self.gauge: BaseGauge = eval(gauge_cfg.gauge_class)(gauge_cfg, robot_cfg) + self.first_reset = True self.last_reset_time = 0.0 # save configs @@ -61,7 +62,6 @@ class BasePipeline: logger.info(f"🚀 Starting single run: {self.run_name}") try: self.load() - first_reset = True sim_data = self.sim.step() frame_skip = int(self.robot_cfg.control.control_dt / self.sim_cfg.physics.simulation_dt) assert frame_skip * self.sim_cfg.physics.simulation_dt == self.robot_cfg.control.control_dt, \ @@ -69,20 +69,26 @@ class BasePipeline: logger.info(f"Sim FPS: {1.0 / self.sim_cfg.physics.simulation_dt:.2f}, Control FPS: {1.0 / self.robot_cfg.control.control_dt:.2f}, Frame Skip: {frame_skip:d}") logger.info("Running pipeline...") while not self.gauge.is_done(): - if first_reset: # wait for robot to be still + if self.first_reset: # wait for robot to be still goal_data = GoalData( goal_type=self.robot_cfg.control.support_goal, velocity_goal=VelocityGoal(), # zero velocity position_goal=PositionGoal(), # current position ) - lin_vel = np.linalg.norm(sim_data.proprio.base.lin_vel) - # print(lin_vel, sim_data.sim_time - self.last_reset_time) if np.linalg.norm(sim_data.proprio.base.lin_vel) < 0.05 and sim_data.sim_time - self.last_reset_time > 0.1: - first_reset = False + self.first_reset = False else: goal_data = self.gauge.get_goal(sim_data) - if goal_data is None: + + if goal_data is None: # Change goal + sim_data = self.reset_sim(sim_data) continue + + if goal_data.visualization_pos is not None: + self.sim.set_target_pos(goal_data.visualization_pos) + else: + self.sim.set_target_pos(None) + obs = self.robot.build_observation(self.add_noise(sim_data), goal_data) action, p_gains, d_gains, control_type = self.robot.get_action(obs) @@ -96,10 +102,7 @@ class BasePipeline: sim_data = self.sim.step() self.gauge.update_metrics(sim_data, goal_data) if self.gauge.is_reset(sim_data): - self.sim.reset() - first_reset = True - self.last_reset_time = sim_data.sim_time - sim_data = self.sim.step() + sim_data = self.reset_sim(sim_data) except Exception as e: logger.error(f"❌ Pipeline execution failed with error: {e}") raise e @@ -110,6 +113,13 @@ class BasePipeline: logger.info(f"📁 Logging saved at: {logger.log_dir}") return logger.log_dir + + def reset_sim(self, sim_data: SimData): + self.sim.reset() + self.last_reset_time = sim_data.sim_time + self.first_reset = True + sim_data = self.sim.step() + return sim_data def add_noise(self, sim_data: SimData): sim_data = deepcopy(sim_data) diff --git a/robogauge/tasks/simulator/mujoco_simulator.py b/robogauge/tasks/simulator/mujoco_simulator.py index ec0f5c1..1342f71 100644 --- a/robogauge/tasks/simulator/mujoco_simulator.py +++ b/robogauge/tasks/simulator/mujoco_simulator.py @@ -42,6 +42,7 @@ class MujocoSimulator: self._pause = True self.n_step = 0 self.sim_time = 0.0 + self.target_pos = None def load( self, @@ -167,6 +168,9 @@ class MujocoSimulator: self._pause = not self._pause logger.info(f"Pause toggled: {self._pause}") + def set_target_pos(self, pos): + self.target_pos = pos + def step(self) -> SimData: """ Simulation step, pause will block thread. """ while self._pause: @@ -177,6 +181,17 @@ class MujocoSimulator: # Viewer sync if self.viewer is not None: if self.viewer.is_running(): + if self.target_pos is not None: + self.viewer.user_scn.ngeom = 0 + mujoco.mjv_initGeom( + self.viewer.user_scn.geoms[0], + type=mujoco.mjtGeom.mjGEOM_SPHERE, + size=[0.1, 0, 0], + pos=self.target_pos, + mat=np.eye(3).flatten(), + rgba=[1, 0, 0, 1] + ) + self.viewer.user_scn.ngeom = 1 self.viewer.sync() time_untile_next_render = self.cfg.physics.simulation_dt - ( time.time() - self.last_render_time @@ -193,6 +208,18 @@ class MujocoSimulator: 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) + + if self.target_pos is not None: + self.renderer.scene.ngeom += 1 + mujoco.mjv_initGeom( + self.renderer.scene.geoms[self.renderer.scene.ngeom - 1], + type=mujoco.mjtGeom.mjGEOM_SPHERE, + size=[0.1, 0, 0], + pos=self.target_pos, + mat=np.eye(3).flatten(), + rgba=[1, 0, 0, 1] + ) + frame = self.renderer.render() self.vid_writer.append_data(frame) diff --git a/robogauge/utils/math_utils.py b/robogauge/utils/math_utils.py index 48e349e..e8ad178 100644 --- a/robogauge/utils/math_utils.py +++ b/robogauge/utils/math_utils.py @@ -16,3 +16,13 @@ def get_projected_gravity(quat): gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz) return gravity_orientation + +def quat_rotate_inverse(q, v): + q = np.array(q, np.float32) + v = np.array(v, np.float32) + q_w = q[0] + q_vec = q[1:] + a = v * (2.0 * q_w ** 2 - 1.0) + b = np.cross(q_vec, v) * q_w * 2.0 + c = q_vec * np.dot(q_vec, v) * 2.0 + return a - b + c