refactor: unify g1 walk env tasks

benchmark/benchmark_env_step.py

@@ -5,7 +5,7 @@
     uv run benchmark/benchmark_env_step.py
 
     # Single task + backend:
-    uv run benchmark/benchmark_env_step.py task=g1_joystick_flat/motrix
+    uv run benchmark/benchmark_env_step.py task=g1_walk_flat/motrix
 
     # Override bench params:
     uv run benchmark/benchmark_env_step.py task=go1_joystick_flat/mujoco num_envs=4096 num_steps=500
@@ -61,7 +61,7 @@ def _load_helper_module(module_name: str, relative_path: str):
 TASK_CONFIGS = {
     "go1": "task=go1_joystick_flat",
     "go2": "task=go2_joystick_flat",
-    "g1": "task=g1_joystick_flat",
+    "g1": "task=g1_walk_flat",
     "g1_mt": "task=g1_motion_tracking",
 }
 

benchmark/benchmark_mujoco_backend_step_detail.py

@@ -19,12 +19,12 @@
     uv run benchmark/benchmark_mujoco_backend_step_detail.py
 
     uv run benchmark/benchmark_mujoco_backend_step_detail.py \
-        --tasks go1_joystick_flat,go2_joystick_flat,g1_joystick_flat \
+        --tasks go1_joystick_flat,go2_joystick_flat,g1_walk_flat \
         --env-nums 256,512,1024,2048,4096,8192 \
         --nsteps 1,2,3,4
 
     uv run benchmark/benchmark_mujoco_backend_step_detail.py \
-        --tasks g1_joystick_flat --env-nums 2048 --nsteps 1,2,3 --iters 20
+        --tasks g1_walk_flat --env-nums 2048 --nsteps 1,2,3 --iters 20
 """
 
 from __future__ import annotations
@@ -89,7 +89,7 @@ def _load_helper_module(module_name: str, relative_path: str):
 TASK_COLORS = {
     "go1_joystick_flat": "#4C78A8",
     "go2_joystick_flat": "#54A24B",
-    "g1_joystick_flat": "#F58518",
+    "g1_walk_flat": "#F58518",
 }
 COMPONENT_COLORS = {
     "set_ctrl_ms": "#9AA0A6",

benchmark/benchmark_physics_step_genesis.py

@@ -3,7 +3,7 @@
 Benchmark Genesis physics execution.
 
 Benchmarks Genesis across current locomotion owner task ids
-(go1_joystick_flat/go2_joystick_flat/g1_joystick_flat) and outputs JSON + plots
+(go1_joystick_flat/go2_joystick_flat/g1_walk_flat) and outputs JSON + plots
 aligned with benchmark/benchmark_physics_step_mj_step.py.
 Legacy env names remain accepted as aliases.
 
@@ -20,32 +20,38 @@
 from dataclasses import asdict, dataclass
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Dict, List
+from typing import Any, Dict, List, cast
 
 import matplotlib
 
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 
 try:
-    import genesis as gs
+    import genesis as _gs
 except ImportError:
     gs = None
+else:
+    gs = _gs
 
 try:
-    from benchmark.core.device_info import get_device_info_dict, get_device_info_line
-    from benchmark.core.task_names import (
-        canonical_locomotion_task_ids,
-        locomotion_task_spec,
-        normalize_locomotion_task_id,
-    )
+    from benchmark.core import device_info as _benchmark_device_info
+    from benchmark.core import task_names as _benchmark_task_names
+
+    _device_info = _benchmark_device_info
+    _task_names = _benchmark_task_names
 except ModuleNotFoundError:
-    from core.device_info import get_device_info_dict, get_device_info_line
-    from core.task_names import (
-        canonical_locomotion_task_ids,
-        locomotion_task_spec,
-        normalize_locomotion_task_id,
-    )
+    from core import device_info as _core_device_info
+    from core import task_names as _core_task_names
+
+    _device_info = _core_device_info
+    _task_names = _core_task_names
+
+get_device_info_dict = _device_info.get_device_info_dict
+get_device_info_line = _device_info.get_device_info_line
+canonical_locomotion_task_ids = _task_names.canonical_locomotion_task_ids
+locomotion_task_spec = _task_names.locomotion_task_spec
+normalize_locomotion_task_id = _task_names.normalize_locomotion_task_id
 
 
 @dataclass
@@ -84,7 +90,8 @@ def _init_genesis() -> None:
     global _GENESIS_INITIALIZED
     if _GENESIS_INITIALIZED:
         return
-    gs.init(backend=gs.gpu)
+    gs_mod = cast(Any, gs)
+    gs_mod.init(backend=gs_mod.gpu)
     _GENESIS_INITIALIZED = True
 
 
@@ -94,14 +101,15 @@ def _load_task_xml(task_name: str) -> str:
 
 
 def _build_scene(xml_path: str, batch_size: int):
-    scene = gs.Scene(
+    gs_mod = cast(Any, gs)
+    scene = gs_mod.Scene(
         show_viewer=False,
-        rigid_options=gs.options.RigidOptions(
+        rigid_options=gs_mod.options.RigidOptions(
             dt=0.01,
-            constraint_solver=gs.constraint_solver.Newton,
+            constraint_solver=gs_mod.constraint_solver.Newton,
         ),
     )
-    scene.add_entity(gs.morphs.MJCF(file=xml_path))
+    scene.add_entity(gs_mod.morphs.MJCF(file=xml_path))
     scene.build(n_envs=batch_size)
     return scene
 

benchmark/benchmark_physics_step_isaacgym.py

@@ -25,7 +25,7 @@
 from dataclasses import asdict, dataclass
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Dict, List
+from typing import Any, Dict, List, cast
 
 import matplotlib
 
@@ -43,13 +43,13 @@
 if str(DEFAULT_ISAACGYM_PYTHON) not in sys.path:
     sys.path.insert(0, str(DEFAULT_ISAACGYM_PYTHON))
 
+_ISAACGYM_IMPORT_ERROR: Exception | None = None
+
 try:
     from isaacgym import gymapi
 except Exception as _isaacgym_error:
     gymapi = None
     _ISAACGYM_IMPORT_ERROR = _isaacgym_error
-else:
-    _ISAACGYM_IMPORT_ERROR = None
 
 
 @dataclass(frozen=True)
@@ -87,9 +87,9 @@ class BenchRecord:
         asset_file="go2_description/urdf/go2_description.urdf",
         initial_height=0.40,
     ),
-    "g1_joystick_flat": TaskSpec(
-        owner_task_id="g1_joystick_flat",
-        display_name="g1_joystick_flat",
+    "g1_walk_flat": TaskSpec(
+        owner_task_id="g1_walk_flat",
+        display_name="g1_walk_flat",
         asset_root=DEFAULT_MODELS_ROOT,
         asset_file="g1_description/g1_29dof_rev_1_0.urdf",
         initial_height=0.78,
@@ -98,16 +98,16 @@ class BenchRecord:
 TASK_ALIASES = {
     "Go1JoystickFlat": "go1_joystick_flat",
     "Go2JoystickFlat": "go2_joystick_flat",
-    "G1JoystickFlat": "g1_joystick_flat",
+    "G1WalkFlat": "g1_walk_flat",
     "task=go1_joystick_flat/isaacgym": "go1_joystick_flat",
     "task=go2_joystick_flat/isaacgym": "go2_joystick_flat",
-    "task=g1_joystick_flat/isaacgym": "g1_joystick_flat",
+    "task=g1_walk_flat/isaacgym": "g1_walk_flat",
     "go1_joystick_flat/isaacgym": "go1_joystick_flat",
     "go2_joystick_flat/isaacgym": "go2_joystick_flat",
-    "g1_joystick_flat/isaacgym": "g1_joystick_flat",
+    "g1_walk_flat/isaacgym": "g1_walk_flat",
     "go1": "go1_joystick_flat",
     "go2": "go2_joystick_flat",
-    "g1": "g1_joystick_flat",
+    "g1": "g1_walk_flat",
 }
 DEFAULT_TASK_IDS = list(TASK_SPECS.keys())
 DEFAULT_BATCH_SIZES = [2**k for k in range(8, 15)]  # 256 .. 16384
@@ -196,19 +196,20 @@ def _default_dof_targets(dof_props) -> np.ndarray:
 
 
 def _create_sim(compute_device_id: int, graphics_device_id: int, nthread: int):
-    gym = gymapi.acquire_gym()
-    sim_params = gymapi.SimParams()
+    gymapi_mod = cast(Any, gymapi)
+    gym = gymapi_mod.acquire_gym()
+    sim_params = gymapi_mod.SimParams()
     sim_params.dt = 0.01
     sim_params.substeps = 1
-    sim_params.up_axis = gymapi.UpAxis.UP_AXIS_Z
-    sim_params.gravity = gymapi.Vec3(0.0, 0.0, -9.81)
+    sim_params.up_axis = gymapi_mod.UpAxis.UP_AXIS_Z
+    sim_params.gravity = gymapi_mod.Vec3(0.0, 0.0, -9.81)
     sim_params.physx.solver_type = 1
     sim_params.physx.num_position_iterations = 4
     sim_params.physx.num_velocity_iterations = 1
     sim_params.physx.num_threads = nthread
     sim_params.physx.use_gpu = True
     sim_params.use_gpu_pipeline = True
-    sim = gym.create_sim(compute_device_id, graphics_device_id, gymapi.SIM_PHYSX, sim_params)
+    sim = gym.create_sim(compute_device_id, graphics_device_id, gymapi_mod.SIM_PHYSX, sim_params)
     if sim is None:
         raise RuntimeError("Failed to create Isaac Gym simulation.")
     return gym, sim
@@ -223,31 +224,32 @@ def _build_task_sim(
 ):
     spec = _task_spec(task_name)
     gym, sim = _create_sim(compute_device_id, graphics_device_id, nthread)
+    gymapi_mod = cast(Any, gymapi)
 
-    plane_params = gymapi.PlaneParams()
-    plane_params.normal = gymapi.Vec3(0.0, 0.0, 1.0)
+    plane_params = gymapi_mod.PlaneParams()
+    plane_params.normal = gymapi_mod.Vec3(0.0, 0.0, 1.0)
     gym.add_ground(sim, plane_params)
 
-    asset_options = gymapi.AssetOptions()
+    asset_options = gymapi_mod.AssetOptions()
     asset_options.flip_visual_attachments = True
     asset_options.armature = 0.01
-    asset_options.default_dof_drive_mode = int(gymapi.DOF_MODE_POS)
+    asset_options.default_dof_drive_mode = int(gymapi_mod.DOF_MODE_POS)
 
     asset = gym.load_asset(sim, str(spec.asset_root), spec.asset_file, asset_options)
     dof_props = gym.get_asset_dof_properties(asset)
-    dof_props["driveMode"][:].fill(gymapi.DOF_MODE_POS)
+    dof_props["driveMode"][:].fill(gymapi_mod.DOF_MODE_POS)
     dof_props["stiffness"][:].fill(1000.0)
     dof_props["damping"][:].fill(10.0)
     dof_targets = _default_dof_targets(dof_props)
 
     spacing = 1.0
-    env_lower = gymapi.Vec3(-spacing, 0.0, -spacing)
-    env_upper = gymapi.Vec3(spacing, spacing, spacing)
+    env_lower = gymapi_mod.Vec3(-spacing, 0.0, -spacing)
+    env_upper = gymapi_mod.Vec3(spacing, spacing, spacing)
     num_per_row = max(1, int(math.sqrt(batch_size)))
 
-    pose = gymapi.Transform()
-    pose.r = gymapi.Quat(0.0, 0.0, 0.0, 1.0)
-    pose.p = gymapi.Vec3(0.0, 0.0, spec.initial_height)
+    pose = gymapi_mod.Transform()
+    pose.r = gymapi_mod.Quat(0.0, 0.0, 0.0, 1.0)
+    pose.p = gymapi_mod.Vec3(0.0, 0.0, spec.initial_height)
 
     for env_idx in range(batch_size):
         env = gym.create_env(sim, env_lower, env_upper, num_per_row)