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

v1.0.5; add num_processes and support dual moe

Browse Source
This commit is contained in:
wty-yy
2026-01-07 00:20:22 +08:00
parent 9c3184a43d
commit 0f2dc29dfc
3 changed files with 24 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
robogauge/scripts/server.py
View File

@@ -31,7 +31,7 @@ default_args_list = [
'--stress-benchmark',
'--stress-terrain-names', 'flat', 'wave', 'slope_fd', 'slope_bd', 'stairs_fd', 'stairs_bd', 'obstacle',
# '--stress-terrain-names', 'flat', 'wave',
'--num-processes', '30',
# '--num-processes', '30', # Set in CLI
'--seeds', '0', '1', '2',
'--search-seeds', '0', '1', '2', '3', '4',
'--frictions', '0.5', '0.75', '1.0', '1.25', '1.5', '1.75', '2.0', '2.25', '2.5',
@@ -95,6 +95,7 @@ def run_api_server(input_queue: multiprocessing.Queue, result_dict: dict, port=9
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--port', type=int, default=9973, help='API server port')
parser.add_argument('--num-processes', type=int, default=30, help='Number of parallel processes for StressPipeline')
args_cli = parser.parse_args()
print("🤖 RoboGauge Evaluation Server Starting...")
ctx = multiprocessing.get_context('spawn')
@@ -122,7 +123,12 @@ def main():
results_store[task_id] = {"status": ResponseStatus.PROCESSING}
args_list = default_args_list.copy()
args_list += ['--model-path', task_data.model_path, '--task-name', task_data.task_name, '--experiment-name', task_data.experiment_name]
args_list += [
'--model-path', task_data.model_path,
'--task-name', task_data.task_name,
'--experiment-name', task_data.experiment_name,
'--num-processes', str(args_cli.num_processes),
]
args = parse_args(args_list)
print(f"📋 Running with args:")

18
robogauge/tasks/robots/go2/go2_moe.py
View File

@@ -23,13 +23,23 @@ class Go2MoE(Go2):
def get_action(self, obs: np.ndarray):
obs_tensor = torch.tensor(obs, dtype=torch.float32).unsqueeze(0).to(self.device)
action, results = self.model(obs_tensor)
if isinstance(results, tuple) and len(results) == 2:
weights, latent = results
if isinstance(results, tuple):
latent = results[-1]
latent = latent.detach().cpu().numpy().squeeze(0) if latent is not None else None
weights = weights.detach().cpu().numpy().squeeze(0) if weights is not None else None
student_weights = None
actor_weights = None
if len(results) >= 2:
student_weights = results[0]
student_weights = student_weights.detach().cpu().numpy().squeeze(0) if student_weights is not None else None
if len(results) >= 3:
actor_weights = results[1]
actor_weights = actor_weights.detach().cpu().numpy().squeeze(0) if actor_weights is not None else None
if self.cfg.control.save_additional_output:
self.save_info['latent'].append(latent)
self.save_info['weights'].append(weights)
if student_weights is not None:
self.save_info['weights'].append(student_weights)
if actor_weights is not None:
self.save_info['actor_weights'].append(actor_weights)
action = action.detach().cpu().numpy().squeeze(0)[self.model2mj_idx]
self.last_action = action
target_dof_pos = action * self.action_scale + self.default_dof_pos