[software] 添加16DOF早期训练仿真与Sim2Real闭环
This commit is contained in:
@@ -0,0 +1,194 @@
|
||||
"""Body impulse demo with force visualization.
|
||||
|
||||
A cylinder hangs from a ball joint like a punching bag. Random impulses
|
||||
swing it around while magenta arrows show the applied forces. Both native
|
||||
and Viser viewers render the arrows via ``apply_body_impulse``'s built-in
|
||||
debug visualization.
|
||||
|
||||
Run with:
|
||||
uv run mjpython scripts/demos/body_impulse.py # macOS
|
||||
uv run python scripts/demos/body_impulse.py # Linux
|
||||
uv run python scripts/demos/body_impulse.py --viewer viser # Viser
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
|
||||
import mujoco
|
||||
import torch
|
||||
import tyro
|
||||
|
||||
import mjlab
|
||||
from mjlab.entity import EntityCfg
|
||||
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
||||
from mjlab.envs.mdp.events import apply_body_impulse, reset_scene_to_default
|
||||
from mjlab.managers.event_manager import EventTermCfg
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.rl import RslRlVecEnvWrapper
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.utils.torch import configure_torch_backends
|
||||
from mjlab.viewer import NativeMujocoViewer, ViserPlayViewer
|
||||
|
||||
BAG_RADIUS = 0.12 # punching bag radius
|
||||
BAG_HALF_HEIGHT = 0.25 # punching bag half-height
|
||||
BAG_DENSITY = 500.0 # ~5.7 kg
|
||||
ROPE_LEN = 0.3 # rope length in meters
|
||||
CEILING_Z = 1.0 # pivot height
|
||||
ROPE_RADIUS = 0.008 # visual rope thickness
|
||||
|
||||
|
||||
def create_pendulum_spec() -> mujoco.MjSpec:
|
||||
spec = mujoco.MjSpec()
|
||||
spec.modelname = "impulse_punching_bag"
|
||||
|
||||
# Ground plane.
|
||||
ground = spec.worldbody.add_geom()
|
||||
ground.type = mujoco.mjtGeom.mjGEOM_PLANE
|
||||
ground.size[:] = (5.0, 5.0, 0.1)
|
||||
ground.rgba[:] = (0.4, 0.5, 0.6, 1.0)
|
||||
|
||||
# Light.
|
||||
light = spec.worldbody.add_light()
|
||||
light.pos[:] = (0, 0, 4)
|
||||
light.dir[:] = (0, 0, -1)
|
||||
light.diffuse[:] = (0.8, 0.8, 0.8)
|
||||
|
||||
# Ceiling anchor (visual only).
|
||||
anchor = spec.worldbody.add_geom()
|
||||
anchor.type = mujoco.mjtGeom.mjGEOM_BOX
|
||||
anchor.size[:] = (0.04, 0.04, 0.02)
|
||||
anchor.pos[:] = (0, 0, CEILING_Z)
|
||||
anchor.rgba[:] = (0.3, 0.3, 0.3, 1.0)
|
||||
anchor.contype = 0
|
||||
anchor.conaffinity = 0
|
||||
|
||||
# Pendulum body. Ball joint pivot is at (0, 0, CEILING_Z).
|
||||
bag_body = spec.worldbody.add_body()
|
||||
bag_body.name = "bag"
|
||||
bag_body.pos[:] = (0, 0, CEILING_Z)
|
||||
|
||||
joint = bag_body.add_joint()
|
||||
joint.name = "bag_joint"
|
||||
joint.type = mujoco.mjtJoint.mjJNT_BALL
|
||||
joint.damping[:] = 3.0
|
||||
joint.frictionloss = 0.5
|
||||
|
||||
# Rope (visual only capsule from pivot to bag center).
|
||||
rope = bag_body.add_geom()
|
||||
rope.type = mujoco.mjtGeom.mjGEOM_CAPSULE
|
||||
rope.size[:2] = (ROPE_RADIUS, ROPE_LEN / 2)
|
||||
rope.pos[:] = (0, 0, -ROPE_LEN / 2)
|
||||
rope.rgba[:] = (0.5, 0.4, 0.3, 1.0)
|
||||
rope.contype = 0
|
||||
rope.conaffinity = 0
|
||||
rope.mass = 0.001 # negligible mass
|
||||
|
||||
# Punching bag cylinder hanging at the end of the rope.
|
||||
geom = bag_body.add_geom()
|
||||
geom.name = "bag_geom"
|
||||
geom.type = mujoco.mjtGeom.mjGEOM_CYLINDER
|
||||
geom.size[:2] = (BAG_RADIUS, BAG_HALF_HEIGHT)
|
||||
geom.pos[:] = (0, 0, -ROPE_LEN - BAG_HALF_HEIGHT)
|
||||
geom.density = BAG_DENSITY
|
||||
geom.rgba[:] = (0.55, 0.15, 0.1, 0.35)
|
||||
|
||||
return spec
|
||||
|
||||
|
||||
def create_env_cfg() -> ManagerBasedRlEnvCfg:
|
||||
bag_cfg = EntityCfg(
|
||||
spec_fn=create_pendulum_spec,
|
||||
init_state=EntityCfg.InitialStateCfg(
|
||||
pos=(0.0, 0.0, 0.0),
|
||||
),
|
||||
)
|
||||
|
||||
bag_mass = BAG_DENSITY * math.pi * BAG_RADIUS**2 * (2 * BAG_HALF_HEIGHT)
|
||||
weight = bag_mass * 9.81
|
||||
force_mag = weight * 0.8 # 0.8x body weight
|
||||
|
||||
cfg = ManagerBasedRlEnvCfg(
|
||||
decimation=10,
|
||||
scene=SceneCfg(
|
||||
num_envs=1,
|
||||
env_spacing=0.0,
|
||||
extent=2.0,
|
||||
entities={"bag": bag_cfg},
|
||||
),
|
||||
events={
|
||||
"reset_scene_to_default": EventTermCfg(
|
||||
func=reset_scene_to_default,
|
||||
mode="reset",
|
||||
),
|
||||
"impulse": EventTermCfg(
|
||||
func=apply_body_impulse,
|
||||
mode="step",
|
||||
params={
|
||||
"force_range": (-force_mag, force_mag),
|
||||
"torque_range": (-force_mag * 0.3, force_mag * 0.3),
|
||||
"duration_s": (0.05, 0.1),
|
||||
"cooldown_s": (1.0, 2.5),
|
||||
"asset_cfg": SceneEntityCfg("bag", body_names=("bag",)),
|
||||
},
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
cfg.viewer.distance = 1.8
|
||||
cfg.viewer.elevation = -10.0
|
||||
|
||||
return cfg
|
||||
|
||||
|
||||
class ZeroPolicy:
|
||||
def __call__(self, obs: object) -> torch.Tensor:
|
||||
del obs
|
||||
return torch.zeros(1, 0)
|
||||
|
||||
|
||||
def main(device: str = "cpu", viewer: str = "auto") -> None:
|
||||
configure_torch_backends()
|
||||
|
||||
env_cfg = create_env_cfg()
|
||||
env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
|
||||
env = RslRlVecEnvWrapper(env)
|
||||
|
||||
# Print force scaling info.
|
||||
mjm = env.unwrapped.sim.mj_model
|
||||
bag_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_BODY, "bag")
|
||||
subtree_mass = mjm.body_subtreemass[bag_id]
|
||||
weight = subtree_mass * 9.81
|
||||
bag_mass = BAG_DENSITY * math.pi * BAG_RADIUS**2 * (2 * BAG_HALF_HEIGHT)
|
||||
force_mag = bag_mass * 9.81 * 0.8
|
||||
print("=" * 50)
|
||||
print("Body Impulse Demo (punching bag)")
|
||||
print(f" Bag mass : {bag_mass:.2f} kg")
|
||||
print(f" Bag weight : {weight:.2f} N")
|
||||
print(f" Rope length : {ROPE_LEN} m")
|
||||
print(f" Force range : +/-{force_mag:.0f} N")
|
||||
print(f" Force/weight : {force_mag / weight:.1f}x")
|
||||
print("=" * 50)
|
||||
|
||||
if viewer == "auto":
|
||||
has_display = bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"))
|
||||
resolved = "native" if has_display else "viser"
|
||||
else:
|
||||
resolved = viewer
|
||||
|
||||
policy = ZeroPolicy()
|
||||
if resolved == "native":
|
||||
print("Launching native viewer...")
|
||||
NativeMujocoViewer(env, policy).run()
|
||||
elif resolved == "viser":
|
||||
print("Launching Viser viewer...")
|
||||
ViserPlayViewer(env, policy).run()
|
||||
else:
|
||||
raise ValueError(f"Unknown viewer: {viewer}")
|
||||
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tyro.cli(main, config=mjlab.TYRO_FLAGS)
|
||||
@@ -0,0 +1,271 @@
|
||||
"""Demo: contact sensor history catches collisions missed by decimation.
|
||||
|
||||
A bouncing ball with high restitution contacts the ground briefly during each
|
||||
bounce. With large decimation, the contact may start and end within
|
||||
intermediate substeps, so by the final substep there is no active contact and
|
||||
instantaneous sensor reads miss it. Setting ``history_length = decimation``
|
||||
captures every substep.
|
||||
|
||||
Run with:
|
||||
uv run python scripts/demos/contact_sensor_decimation.py
|
||||
uv run python scripts/demos/contact_sensor_decimation.py --viewer
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from mjlab.entity import EntityCfg
|
||||
from mjlab.scene import Scene, SceneCfg
|
||||
from mjlab.sensor.contact_sensor import ContactMatch, ContactSensorCfg
|
||||
from mjlab.sim.sim import Simulation, SimulationCfg
|
||||
|
||||
BOUNCING_BALL_XML = """
|
||||
<mujoco>
|
||||
<option timestep="0.001"/>
|
||||
<worldbody>
|
||||
<body name="ground" pos="0 0 0">
|
||||
<geom name="ground_geom" type="plane" size="5 5 0.1"/>
|
||||
</body>
|
||||
<body name="ball" pos="0 0 1">
|
||||
<freejoint/>
|
||||
<geom name="ball_geom" type="sphere" size="0.05" mass="0.1"
|
||||
solref="-1000 0"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
DECIMATION = 20
|
||||
NUM_ENVS = 1
|
||||
NUM_POLICY_STEPS = 200
|
||||
PHYSICS_DT = 0.001
|
||||
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
|
||||
def build(history_length: int) -> tuple[Scene, Simulation]:
|
||||
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(BOUNCING_BALL_XML))
|
||||
sensor_cfg = ContactSensorCfg(
|
||||
name="ball_contact",
|
||||
primary=ContactMatch(mode="geom", pattern="ball_geom", entity="ball"),
|
||||
secondary=None,
|
||||
fields=("found", "force"),
|
||||
history_length=history_length,
|
||||
)
|
||||
scene_cfg = SceneCfg(
|
||||
num_envs=NUM_ENVS,
|
||||
env_spacing=3.0,
|
||||
entities={"ball": entity_cfg},
|
||||
sensors=(sensor_cfg,),
|
||||
)
|
||||
scene = Scene(scene_cfg, DEVICE)
|
||||
model = scene.compile()
|
||||
sim = Simulation(
|
||||
num_envs=NUM_ENVS,
|
||||
cfg=SimulationCfg(njmax=50),
|
||||
model=model,
|
||||
device=DEVICE,
|
||||
)
|
||||
scene.initialize(sim.mj_model, sim.model, sim.data)
|
||||
return scene, sim
|
||||
|
||||
|
||||
def run_no_history():
|
||||
"""Read instantaneous contact at the end of each policy step."""
|
||||
scene, sim = build(history_length=0)
|
||||
sensor = scene["ball_contact"]
|
||||
|
||||
contact_detected = []
|
||||
for _ in range(NUM_POLICY_STEPS):
|
||||
for _ in range(DECIMATION):
|
||||
sim.step()
|
||||
scene.update(dt=PHYSICS_DT)
|
||||
found = sensor.data.found[0, 0].item() > 0
|
||||
contact_detected.append(found)
|
||||
|
||||
return contact_detected
|
||||
|
||||
|
||||
def run_with_history():
|
||||
"""Read full substep history to catch mid-decimation contacts."""
|
||||
scene, sim = build(history_length=DECIMATION)
|
||||
sensor = scene["ball_contact"]
|
||||
|
||||
contact_detected_instant = []
|
||||
contact_detected_history = []
|
||||
ball_height_substep = []
|
||||
for _ in range(NUM_POLICY_STEPS):
|
||||
for _ in range(DECIMATION):
|
||||
sim.step()
|
||||
scene.update(dt=PHYSICS_DT)
|
||||
# qpos is always current after step; qpos[2] is z for a freejoint.
|
||||
ball_height_substep.append(sim.data.qpos[0, 2].item())
|
||||
data = sensor.data
|
||||
found_instant = data.found[0, 0].item() > 0
|
||||
# Check whether any substep in the decimation window had contact.
|
||||
force_hist = data.force_history # [B, N, H, 3]
|
||||
found_history = (force_hist[0, 0].norm(dim=-1) > 1e-6).any().item()
|
||||
contact_detected_instant.append(found_instant)
|
||||
contact_detected_history.append(found_history)
|
||||
|
||||
return contact_detected_instant, contact_detected_history, ball_height_substep
|
||||
|
||||
|
||||
def run_viewer():
|
||||
"""Launch a Viser viewer showing the bouncing ball with contact forces."""
|
||||
import viser
|
||||
|
||||
from mjlab.viewer.viser import ViserMujocoScene
|
||||
|
||||
scene, sim = build(history_length=0)
|
||||
|
||||
server = viser.ViserServer(label="Bouncing Ball")
|
||||
viz = ViserMujocoScene(server, sim.mj_model, num_envs=NUM_ENVS)
|
||||
viz.show_contact_forces = True
|
||||
viz.show_contact_points = True
|
||||
viz.create_scene_gui(
|
||||
camera_distance=2.0,
|
||||
camera_azimuth=90.0,
|
||||
camera_elevation=20.0,
|
||||
)
|
||||
|
||||
print("Open the Viser URL above to watch the bouncing ball.")
|
||||
print("Contact forces and points are enabled by default.")
|
||||
print("Press Ctrl+C to stop.\n")
|
||||
|
||||
try:
|
||||
while True:
|
||||
for _ in range(DECIMATION):
|
||||
sim.step()
|
||||
scene.update(dt=PHYSICS_DT)
|
||||
viz.update(sim.data)
|
||||
if viz.needs_update:
|
||||
viz.refresh_visualization()
|
||||
time.sleep(DECIMATION * PHYSICS_DT)
|
||||
except KeyboardInterrupt:
|
||||
print("\nShutting down...")
|
||||
server.stop()
|
||||
|
||||
|
||||
def run_analysis():
|
||||
"""Run the analysis comparing instantaneous vs history contact detection."""
|
||||
print("=" * 70)
|
||||
print("Contact Sensor Decimation Demo")
|
||||
print(f" Ball dropped from 1m, restitution ~ 1, decimation = {DECIMATION}")
|
||||
print(f" Physics dt = {PHYSICS_DT}s, policy dt = {DECIMATION * PHYSICS_DT}s")
|
||||
print("=" * 70)
|
||||
|
||||
no_hist = run_no_history()
|
||||
instant, history, ball_height = run_with_history()
|
||||
|
||||
# Policy steps where history caught a contact that instant missed.
|
||||
missed = []
|
||||
for i in range(NUM_POLICY_STEPS):
|
||||
if history[i] and not instant[i]:
|
||||
missed.append(i)
|
||||
|
||||
total_contacts_instant = sum(instant)
|
||||
total_contacts_history = sum(history)
|
||||
total_contacts_no_hist = sum(no_hist)
|
||||
|
||||
print()
|
||||
print(f"Total policy steps with contact (no history): {total_contacts_no_hist}")
|
||||
print(f"Total policy steps with contact (instant only): {total_contacts_instant}")
|
||||
print(f"Total policy steps with contact (with history): {total_contacts_history}")
|
||||
print()
|
||||
|
||||
if missed:
|
||||
print(f"Contacts MISSED by instantaneous read but CAUGHT by history: {len(missed)}")
|
||||
print(f" Policy steps: {missed}")
|
||||
else:
|
||||
print("No missed contacts (try increasing decimation or adjusting drop height)")
|
||||
|
||||
print()
|
||||
print("Step-by-step (showing first 60 policy steps):")
|
||||
print(f"{'step':>6} {'no_hist':>8} {'instant':>8} {'history':>8} {'missed':>8}")
|
||||
print("-" * 50)
|
||||
for i in range(min(60, NUM_POLICY_STEPS)):
|
||||
flag = " <<<" if (history[i] and not instant[i]) else ""
|
||||
print(f"{i:>6} {no_hist[i]!s:>8} {instant[i]!s:>8} {history[i]!s:>8} {flag}")
|
||||
|
||||
# --- Plot ---
|
||||
total_substeps = NUM_POLICY_STEPS * DECIMATION
|
||||
t_substep = np.arange(total_substeps) * PHYSICS_DT
|
||||
|
||||
# Place markers at the minimum height within each policy step window.
|
||||
min_height = []
|
||||
min_time = []
|
||||
for i in range(NUM_POLICY_STEPS):
|
||||
start = i * DECIMATION
|
||||
end = (i + 1) * DECIMATION
|
||||
window = ball_height[start:end]
|
||||
j = int(np.argmin(window))
|
||||
min_height.append(window[j])
|
||||
min_time.append(t_substep[start + j])
|
||||
|
||||
# Separate history detections into: caught by both, caught only by history.
|
||||
idx_both = [i for i in range(NUM_POLICY_STEPS) if instant[i] and history[i]]
|
||||
idx_history_only = missed # history=True, instant=False
|
||||
|
||||
fig, ax = plt.subplots(figsize=(12, 4))
|
||||
ax.plot(t_substep, ball_height, color="0.4", linewidth=0.8, label="Ball height")
|
||||
|
||||
if idx_both:
|
||||
ax.scatter(
|
||||
[min_time[i] for i in idx_both],
|
||||
[min_height[i] for i in idx_both],
|
||||
color="tab:green",
|
||||
s=40,
|
||||
zorder=3,
|
||||
label="Detected by both",
|
||||
)
|
||||
|
||||
if idx_history_only:
|
||||
ax.scatter(
|
||||
[min_time[i] for i in idx_history_only],
|
||||
[min_height[i] for i in idx_history_only],
|
||||
color="tab:red",
|
||||
s=60,
|
||||
marker="x",
|
||||
linewidths=2,
|
||||
zorder=4,
|
||||
label="Caught only by history",
|
||||
)
|
||||
|
||||
ax.set_xlabel("Time (s)")
|
||||
ax.set_ylabel("Ball height (m)")
|
||||
ax.set_title(
|
||||
f"Contact sensor with decimation = {DECIMATION}: "
|
||||
f"{len(missed)} collisions missed without history"
|
||||
)
|
||||
ax.legend(loc="upper right")
|
||||
ax.set_ylim(bottom=-0.05)
|
||||
fig.tight_layout()
|
||||
fig.savefig("scripts/demos/contact_sensor_decimation.png", dpi=150)
|
||||
print("\nPlot saved to scripts/demos/contact_sensor_decimation.png")
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--viewer",
|
||||
action="store_true",
|
||||
help="Launch a Viser viewer instead of running the analysis.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.viewer:
|
||||
run_viewer()
|
||||
else:
|
||||
run_analysis()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,194 @@
|
||||
"""Interactive IK control demo.
|
||||
|
||||
Drag the 3D transform control in the viser viewer to move the YAM end-effector.
|
||||
|
||||
Run with:
|
||||
MJLAB_WARP_QUIET=1 uv run scripts/demos/differential_ik.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
import viser
|
||||
|
||||
from mjlab.asset_zoo.robots.i2rt_yam.yam_constants import get_yam_robot_cfg
|
||||
from mjlab.entity import Entity, EntityCfg
|
||||
from mjlab.envs.mdp.actions import DifferentialIKAction, DifferentialIKActionCfg
|
||||
from mjlab.sim.sim import MujocoCfg, Simulation, SimulationCfg
|
||||
from mjlab.utils.lab_api.math import quat_from_matrix
|
||||
from mjlab.viewer.viser import ViserMujocoScene
|
||||
|
||||
DEMO_INIT_STATE = EntityCfg.InitialStateCfg(
|
||||
pos=(0.0, 0.0, 0.01),
|
||||
joint_pos={
|
||||
"joint2": 0.6,
|
||||
"joint3": 0.6,
|
||||
"joint4": 0.0,
|
||||
"left_finger": 0.037,
|
||||
"right_finger": -0.037,
|
||||
},
|
||||
joint_vel={".*": 0.0},
|
||||
)
|
||||
|
||||
IK_ITERATIONS = 10
|
||||
|
||||
|
||||
def main() -> None:
|
||||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
robot_cfg = get_yam_robot_cfg()
|
||||
robot_cfg.init_state = DEMO_INIT_STATE
|
||||
entity = Entity(robot_cfg)
|
||||
model = entity.compile()
|
||||
sim_cfg = SimulationCfg(mujoco=MujocoCfg(gravity=(0, 0, -9.81)))
|
||||
sim = Simulation(num_envs=1, cfg=sim_cfg, model=model, device=device)
|
||||
entity.initialize(model, sim.model, sim.data, device)
|
||||
entity.write_joint_position_to_sim(entity.data.default_joint_pos, joint_ids=None)
|
||||
sim.forward()
|
||||
|
||||
env = SimpleNamespace(num_envs=1, device=device, scene={"robot": entity}, sim=sim)
|
||||
ik_cfg = DifferentialIKActionCfg(
|
||||
entity_name="robot",
|
||||
actuator_names=("joint.*",),
|
||||
frame_name="grasp_site",
|
||||
frame_type="site",
|
||||
posture_weight=0.02,
|
||||
joint_limit_weight=1e-1,
|
||||
damping=1e-1,
|
||||
use_relative_mode=False,
|
||||
)
|
||||
ik_action: DifferentialIKAction = ik_cfg.build(env) # type: ignore[arg-type]
|
||||
joint_ids = ik_action._joint_ids
|
||||
|
||||
grip_ids, _ = entity.find_joints("left_finger")
|
||||
grip_joint_ids = torch.tensor(grip_ids, device=device, dtype=torch.long)
|
||||
grip_open = torch.tensor([[0.037]], device=device)
|
||||
|
||||
server = viser.ViserServer(label="IK Control Demo")
|
||||
scene = ViserMujocoScene(server, sim.mj_model, num_envs=1)
|
||||
scene.create_scene_gui(
|
||||
camera_distance=0.1,
|
||||
camera_azimuth=135.0,
|
||||
camera_elevation=30.0,
|
||||
)
|
||||
|
||||
site_id = ik_action._frame_id
|
||||
pos = sim.data.site_xpos[0, site_id].cpu().numpy()
|
||||
xmat = sim.data.site_xmat[0, site_id]
|
||||
quat = quat_from_matrix(xmat).cpu().numpy()
|
||||
|
||||
transform_ctrl = server.scene.add_transform_controls(
|
||||
"/ik_target",
|
||||
position=(float(pos[0]), float(pos[1]), float(pos[2])),
|
||||
wxyz=(float(quat[0]), float(quat[1]), float(quat[2]), float(quat[3])),
|
||||
scale=0.12,
|
||||
)
|
||||
|
||||
needs_reset = [False]
|
||||
|
||||
with server.gui.add_folder("IK Control"):
|
||||
reset_button = server.gui.add_button("Reset")
|
||||
reset_button.on_click(lambda _: needs_reset.__setitem__(0, True))
|
||||
iterations_slider = server.gui.add_slider(
|
||||
"IK Iterations",
|
||||
min=1,
|
||||
max=50,
|
||||
step=1,
|
||||
initial_value=IK_ITERATIONS,
|
||||
)
|
||||
|
||||
with server.gui.add_folder("IK Weights"):
|
||||
damping_slider = server.gui.add_slider(
|
||||
"Damping (λ)",
|
||||
min=1e-2,
|
||||
max=1.0,
|
||||
step=1e-3,
|
||||
initial_value=ik_cfg.damping,
|
||||
)
|
||||
pos_w_slider = server.gui.add_slider(
|
||||
"Position Weight",
|
||||
min=0.0,
|
||||
max=10.0,
|
||||
step=0.1,
|
||||
initial_value=ik_cfg.position_weight,
|
||||
)
|
||||
ori_w_slider = server.gui.add_slider(
|
||||
"Orientation Weight",
|
||||
min=0.0,
|
||||
max=10.0,
|
||||
step=0.1,
|
||||
initial_value=ik_cfg.orientation_weight,
|
||||
)
|
||||
jlim_w_slider = server.gui.add_slider(
|
||||
"Joint Limit Weight",
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
initial_value=ik_cfg.joint_limit_weight,
|
||||
)
|
||||
posture_w_slider = server.gui.add_slider(
|
||||
"Posture Weight",
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
initial_value=ik_cfg.posture_weight,
|
||||
)
|
||||
|
||||
print("=" * 60)
|
||||
print("IK Control Demo")
|
||||
print(" Open the viser URL printed above")
|
||||
print(" Drag the 3D transform control to move the end-effector")
|
||||
print("=" * 60)
|
||||
|
||||
target_action = torch.zeros(1, 7, device=device)
|
||||
|
||||
def _reset() -> None:
|
||||
entity.write_joint_position_to_sim(entity.data.default_joint_pos, joint_ids=None)
|
||||
sim.forward()
|
||||
ik_action.reset()
|
||||
p = sim.data.site_xpos[0, site_id].cpu().numpy()
|
||||
q = quat_from_matrix(sim.data.site_xmat[0, site_id]).cpu().numpy()
|
||||
transform_ctrl.position = (float(p[0]), float(p[1]), float(p[2]))
|
||||
transform_ctrl.wxyz = (float(q[0]), float(q[1]), float(q[2]), float(q[3]))
|
||||
|
||||
try:
|
||||
while True:
|
||||
if needs_reset[0]:
|
||||
needs_reset[0] = False
|
||||
_reset()
|
||||
|
||||
ik_cfg.damping = max(damping_slider.value, 1e-2)
|
||||
ik_cfg.position_weight = max(pos_w_slider.value, 0.0)
|
||||
ik_cfg.orientation_weight = max(ori_w_slider.value, 0.0)
|
||||
ik_cfg.joint_limit_weight = max(jlim_w_slider.value, 0.0)
|
||||
ik_cfg.posture_weight = max(posture_w_slider.value, 0.0)
|
||||
|
||||
p = transform_ctrl.position
|
||||
w = transform_ctrl.wxyz
|
||||
target_action[0, :3] = torch.tensor([p[0], p[1], p[2]], device=device)
|
||||
target_action[0, 3:] = torch.tensor([w[0], w[1], w[2], w[3]], device=device)
|
||||
ik_action.process_actions(target_action)
|
||||
|
||||
n_iter = int(iterations_slider.value)
|
||||
for _ in range(n_iter):
|
||||
dq = ik_action.compute_dq()
|
||||
q = entity.data.joint_pos[:, joint_ids] + dq
|
||||
entity.write_joint_position_to_sim(q, joint_ids=joint_ids)
|
||||
entity.write_joint_position_to_sim(grip_open, joint_ids=grip_joint_ids)
|
||||
sim.forward()
|
||||
|
||||
scene.update(sim.data)
|
||||
if scene.needs_update:
|
||||
scene.refresh_visualization()
|
||||
|
||||
time.sleep(1.0 / 30.0)
|
||||
except KeyboardInterrupt:
|
||||
print("\nShutting down...")
|
||||
server.stop()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,125 @@
|
||||
"""Flat patch terrain demo.
|
||||
|
||||
Spawns a Go1 on rough terrain with flat-patch sampling.
|
||||
On each reset, the robot lands on a flat patch.
|
||||
|
||||
Run with:
|
||||
uv run python scripts/demos/flat_patch_terrain.py [--viewer native|viser]
|
||||
|
||||
Toggle visualization group 3 to see flat patch locations visualized as box sites.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
import tyro
|
||||
|
||||
import mjlab
|
||||
import mjlab.terrains as terrain_gen
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
from mjlab.envs.mdp import events as mdp
|
||||
from mjlab.managers.event_manager import EventTermCfg
|
||||
from mjlab.rl import RslRlVecEnvWrapper
|
||||
from mjlab.tasks.velocity.config.go1.env_cfgs import unitree_go1_rough_env_cfg
|
||||
from mjlab.terrains import FlatPatchSamplingCfg
|
||||
from mjlab.terrains.terrain_generator import TerrainGeneratorCfg
|
||||
from mjlab.utils.torch import configure_torch_backends
|
||||
from mjlab.viewer import NativeMujocoViewer, ViserPlayViewer
|
||||
|
||||
|
||||
def main(viewer: str = "auto") -> None:
|
||||
configure_torch_backends()
|
||||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
cfg = unitree_go1_rough_env_cfg(play=True)
|
||||
|
||||
spawn_patch_cfg = FlatPatchSamplingCfg(
|
||||
num_patches=100,
|
||||
patch_radius=0.3,
|
||||
max_height_diff=0.05,
|
||||
)
|
||||
|
||||
# Override terrain: 1 row x 2 cols, curriculum mode so each column is deterministic.
|
||||
# Column 0 = discrete obstacles, Column 1 = pyramid slope.
|
||||
assert cfg.scene.terrain is not None
|
||||
cfg.scene.terrain.terrain_generator = TerrainGeneratorCfg(
|
||||
size=(4.0, 4.0),
|
||||
num_rows=1,
|
||||
num_cols=2,
|
||||
border_width=1.0,
|
||||
curriculum=True,
|
||||
add_lights=True,
|
||||
sub_terrains={
|
||||
"discrete_obstacles": terrain_gen.HfDiscreteObstaclesTerrainCfg(
|
||||
proportion=0.5,
|
||||
obstacle_height_range=(0.05, 0.5),
|
||||
obstacle_width_range=(0.4, 1.2),
|
||||
num_obstacles=30,
|
||||
platform_width=1.5,
|
||||
border_width=0.25,
|
||||
flat_patch_sampling={"spawn": spawn_patch_cfg},
|
||||
),
|
||||
"pyramid_slope": terrain_gen.HfPyramidSlopedTerrainCfg(
|
||||
proportion=0.5,
|
||||
slope_range=(0.3, 0.8),
|
||||
platform_width=1.5,
|
||||
border_width=0.25,
|
||||
flat_patch_sampling={"spawn": spawn_patch_cfg},
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
# Remove all termination conditions except time limit.
|
||||
for key in list(cfg.terminations):
|
||||
if key != "time_out":
|
||||
del cfg.terminations[key]
|
||||
|
||||
# Reset every 2 seconds to better showcase flat patch spawning.
|
||||
cfg.episode_length_s = 2.0
|
||||
|
||||
# Replace reset_base event with flat-patch spawning.
|
||||
cfg.events["reset_base"] = EventTermCfg(
|
||||
func=mdp.reset_root_state_from_flat_patches,
|
||||
mode="reset",
|
||||
params={
|
||||
"patch_name": "spawn",
|
||||
"pose_range": {"z": (0.01, 0.05), "yaw": (-3.14, 3.14)},
|
||||
},
|
||||
)
|
||||
|
||||
print("=" * 60)
|
||||
print("Flat Patch Terrain Demo")
|
||||
print(" Toggle group 3 to see flat patch markers (orange spheres)")
|
||||
print(" Press Enter in terminal to reset robot onto a flat patch")
|
||||
print("=" * 60)
|
||||
|
||||
env = ManagerBasedRlEnv(cfg=cfg, device=device)
|
||||
env = RslRlVecEnvWrapper(env)
|
||||
|
||||
class ZeroPolicy:
|
||||
def __call__(self, obs) -> torch.Tensor:
|
||||
del obs
|
||||
return torch.zeros(env.unwrapped.action_space.shape, device=device)
|
||||
|
||||
policy = ZeroPolicy()
|
||||
|
||||
if viewer == "auto":
|
||||
has_display = bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"))
|
||||
resolved_viewer = "native" if has_display else "viser"
|
||||
else:
|
||||
resolved_viewer = viewer
|
||||
|
||||
if resolved_viewer == "native":
|
||||
NativeMujocoViewer(env, policy).run()
|
||||
elif resolved_viewer == "viser":
|
||||
ViserPlayViewer(env, policy).run()
|
||||
else:
|
||||
raise ValueError(f"Unknown viewer: {viewer}")
|
||||
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tyro.cli(main, config=mjlab.TYRO_FLAGS)
|
||||
@@ -0,0 +1,252 @@
|
||||
"""Raycast sensor demo.
|
||||
|
||||
Run with:
|
||||
uv run mjpython scripts/demos/raycast_sensor.py [--viewer native|viser] # macOS
|
||||
uv run python scripts/demos/raycast_sensor.py [--viewer native|viser] # Linux
|
||||
|
||||
Examples:
|
||||
# Grid pattern (default)
|
||||
uv run python scripts/demos/raycast_sensor.py --pattern grid
|
||||
|
||||
# Pinhole camera pattern
|
||||
uv run python scripts/demos/raycast_sensor.py --pattern pinhole
|
||||
|
||||
# With yaw alignment (ignores pitch/roll)
|
||||
uv run python scripts/demos/raycast_sensor.py --alignment yaw
|
||||
|
||||
If using the native viewer, you can launch in interactive mode with:
|
||||
uv run mjpython scripts/demos/raycast_sensor.py --viewer native --interactive
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import torch
|
||||
import tyro
|
||||
|
||||
import mjlab
|
||||
import mjlab.terrains as terrain_gen
|
||||
from mjlab.entity import EntityCfg
|
||||
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
||||
from mjlab.rl import RslRlVecEnvWrapper
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.sensor import (
|
||||
GridPatternCfg,
|
||||
ObjRef,
|
||||
PinholeCameraPatternCfg,
|
||||
RayCastSensorCfg,
|
||||
)
|
||||
from mjlab.terrains.terrain_entity import TerrainEntityCfg
|
||||
from mjlab.terrains.terrain_generator import TerrainGeneratorCfg
|
||||
from mjlab.utils.torch import configure_torch_backends
|
||||
from mjlab.viewer import NativeMujocoViewer, ViserPlayViewer
|
||||
|
||||
|
||||
def create_scanner_spec() -> mujoco.MjSpec:
|
||||
spec = mujoco.MjSpec()
|
||||
spec.modelname = "scanner"
|
||||
|
||||
mat = spec.add_material()
|
||||
mat.name = "scanner_mat"
|
||||
mat.rgba[:] = (1.0, 0.5, 0.0, 0.9)
|
||||
|
||||
scanner = spec.worldbody.add_body(mocap=True)
|
||||
scanner.name = "scanner"
|
||||
scanner.pos[:] = (0, 0, 2.0)
|
||||
|
||||
geom = scanner.add_geom()
|
||||
geom.name = "scanner_geom"
|
||||
geom.type = mujoco.mjtGeom.mjGEOM_BOX
|
||||
geom.size[:] = (0.15, 0.15, 0.05)
|
||||
geom.mass = 1.0
|
||||
geom.material = "scanner_mat"
|
||||
|
||||
scanner.add_camera(name="scanner", fovy=58.0, resolution=(16, 12))
|
||||
|
||||
record_cam = scanner.add_camera(name="record_cam")
|
||||
record_cam.pos[:] = (2, 0, 2)
|
||||
record_cam.fovy = 40.0
|
||||
record_cam.mode = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODY
|
||||
record_cam.targetbody = "scanner"
|
||||
|
||||
return spec
|
||||
|
||||
|
||||
def create_env_cfg(
|
||||
pattern: Literal["grid", "pinhole"],
|
||||
alignment: Literal["base", "yaw", "world"],
|
||||
) -> ManagerBasedRlEnvCfg:
|
||||
custom_terrain_cfg = TerrainGeneratorCfg(
|
||||
size=(4.0, 4.0),
|
||||
border_width=0.5,
|
||||
num_rows=1,
|
||||
num_cols=4,
|
||||
curriculum=True,
|
||||
sub_terrains={
|
||||
"pyramid_stairs_inv": terrain_gen.BoxInvertedPyramidStairsTerrainCfg(
|
||||
proportion=0.25,
|
||||
step_height_range=(0.1, 0.25),
|
||||
step_width=0.3,
|
||||
platform_width=1.5,
|
||||
border_width=0.25,
|
||||
),
|
||||
"hf_pyramid_slope_inv": terrain_gen.HfPyramidSlopedTerrainCfg(
|
||||
proportion=0.25,
|
||||
slope_range=(0.6, 1.5),
|
||||
platform_width=1.5,
|
||||
border_width=0.25,
|
||||
inverted=True,
|
||||
),
|
||||
"random_rough": terrain_gen.HfRandomUniformTerrainCfg(
|
||||
proportion=0.25,
|
||||
noise_range=(0.05, 0.15),
|
||||
noise_step=0.02,
|
||||
border_width=0.25,
|
||||
),
|
||||
"wave_terrain": terrain_gen.HfWaveTerrainCfg(
|
||||
proportion=0.25,
|
||||
amplitude_range=(0.15, 0.25),
|
||||
num_waves=3,
|
||||
border_width=0.25,
|
||||
),
|
||||
},
|
||||
add_lights=True,
|
||||
)
|
||||
|
||||
terrain_cfg = TerrainEntityCfg(
|
||||
terrain_type="generator",
|
||||
terrain_generator=custom_terrain_cfg,
|
||||
num_envs=1,
|
||||
)
|
||||
|
||||
scanner_entity_cfg = EntityCfg(
|
||||
spec_fn=create_scanner_spec,
|
||||
init_state=EntityCfg.InitialStateCfg(pos=(0.65, -0.4, 0.5)),
|
||||
)
|
||||
|
||||
if pattern == "grid":
|
||||
pattern_cfg = GridPatternCfg(
|
||||
size=(0.6, 0.6),
|
||||
resolution=0.1,
|
||||
direction=(0.0, 0.0, -1.0),
|
||||
)
|
||||
else:
|
||||
assert pattern == "pinhole"
|
||||
pattern_cfg = PinholeCameraPatternCfg.from_mujoco_camera("scanner/scanner")
|
||||
|
||||
raycast_cfg = RayCastSensorCfg(
|
||||
name="terrain_scan",
|
||||
frame=ObjRef(type="body", name="scanner", entity="scanner"),
|
||||
pattern=pattern_cfg,
|
||||
ray_alignment=alignment,
|
||||
max_distance=5.0,
|
||||
exclude_parent_body=True,
|
||||
debug_vis=True,
|
||||
viz=RayCastSensorCfg.VizCfg(
|
||||
hit_color=(0.0, 1.0, 0.0, 0.9),
|
||||
miss_color=(1.0, 0.0, 0.0, 0.5),
|
||||
show_rays=False,
|
||||
show_normals=True,
|
||||
),
|
||||
)
|
||||
|
||||
cfg = ManagerBasedRlEnvCfg(
|
||||
decimation=10,
|
||||
scene=SceneCfg(
|
||||
num_envs=1,
|
||||
env_spacing=0.0,
|
||||
extent=2.0,
|
||||
terrain=terrain_cfg,
|
||||
entities={"scanner": scanner_entity_cfg},
|
||||
sensors=(raycast_cfg,),
|
||||
),
|
||||
)
|
||||
|
||||
cfg.viewer.body_name = "scanner"
|
||||
cfg.viewer.distance = 12.0
|
||||
cfg.viewer.elevation = -25.0
|
||||
cfg.viewer.azimuth = 135.0
|
||||
|
||||
return cfg
|
||||
|
||||
|
||||
def main(
|
||||
viewer: str = "auto",
|
||||
interactive: bool = False,
|
||||
pattern: Literal["grid", "pinhole"] = "grid",
|
||||
alignment: Literal["base", "yaw", "world"] = "base",
|
||||
) -> None:
|
||||
configure_torch_backends()
|
||||
|
||||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
print("=" * 60)
|
||||
print("Raycast Sensor Demo - 4 Terrain Types")
|
||||
print(f" Pattern: {pattern}")
|
||||
print(f" Alignment: {alignment}")
|
||||
print("=" * 60)
|
||||
print()
|
||||
|
||||
env_cfg = create_env_cfg(pattern, alignment)
|
||||
env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
|
||||
env = RslRlVecEnvWrapper(env)
|
||||
|
||||
if viewer == "auto":
|
||||
has_display = bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"))
|
||||
resolved_viewer = "native" if has_display else "viser"
|
||||
else:
|
||||
resolved_viewer = viewer
|
||||
|
||||
use_auto_scan = (resolved_viewer == "viser") or (not interactive)
|
||||
|
||||
if use_auto_scan:
|
||||
|
||||
class AutoScanPolicy:
|
||||
def __init__(self):
|
||||
self.step_count = 0
|
||||
|
||||
def __call__(self, obs) -> torch.Tensor:
|
||||
del obs
|
||||
t = self.step_count * 0.005
|
||||
y_period = 1000
|
||||
y_normalized = (self.step_count % y_period) / y_period
|
||||
y = -8.0 + 16.0 * y_normalized
|
||||
x = 1.5 * np.sin(2 * np.pi * t * 0.3)
|
||||
z = 1.0
|
||||
env.unwrapped.sim.data.mocap_pos[0, 0, :] = torch.tensor(
|
||||
[x, y, z], device=device, dtype=torch.float32
|
||||
)
|
||||
env.unwrapped.sim.data.mocap_quat[0, 0, :] = torch.tensor(
|
||||
[1, 0, 0, 0], device=device, dtype=torch.float32
|
||||
)
|
||||
self.step_count += 1
|
||||
return torch.zeros(env.unwrapped.action_space.shape, device=device)
|
||||
|
||||
policy = AutoScanPolicy()
|
||||
else:
|
||||
|
||||
class PolicyZero:
|
||||
def __call__(self, obs) -> torch.Tensor:
|
||||
del obs
|
||||
return torch.zeros(env.unwrapped.action_space.shape, device=device)
|
||||
|
||||
policy = PolicyZero()
|
||||
|
||||
if resolved_viewer == "native":
|
||||
print("Launching native viewer...")
|
||||
NativeMujocoViewer(env, policy).run()
|
||||
elif resolved_viewer == "viser":
|
||||
print("Launching viser viewer...")
|
||||
ViserPlayViewer(env, policy).run()
|
||||
else:
|
||||
raise ValueError(f"Unknown viewer: {viewer}")
|
||||
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tyro.cli(main, config=mjlab.TYRO_FLAGS)
|
||||
Reference in New Issue
Block a user