78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Any, Protocol
|
|
|
|
import numpy as np
|
|
|
|
|
|
@dataclass
|
|
class SimulationConfig:
|
|
realtime: bool = True
|
|
seed: int = 0
|
|
|
|
|
|
class ControlPolicy(Protocol):
|
|
def act(self, *, obs: np.ndarray | None = None, t: float = 0.0) -> np.ndarray: ...
|
|
|
|
|
|
def _default_observations(data: Any) -> np.ndarray:
|
|
qpos = np.asarray(data.qpos, dtype=np.float32).ravel()
|
|
qvel = np.asarray(data.qvel, dtype=np.float32).ravel()
|
|
return np.concatenate([qpos, qvel], axis=0)
|
|
|
|
|
|
def simulate_policy(
|
|
policy: ControlPolicy,
|
|
config: SimulationConfig,
|
|
state: Any | None = None,
|
|
) -> None:
|
|
"""Open MuJoCo's native viewer and step using actions from a policy.
|
|
|
|
This path drives MuJoCo physics directly (mj_step) and uses the policy output
|
|
as `data.ctrl`.
|
|
"""
|
|
|
|
import mujoco.viewer
|
|
|
|
if state is None:
|
|
raise ValueError("A valid environment state must be provided.")
|
|
|
|
model = state.mj_model
|
|
data = state.mj_data
|
|
|
|
start = time.time()
|
|
with mujoco.viewer.launch_passive(model, data) as viewer:
|
|
while viewer.is_running():
|
|
step_start = time.time()
|
|
|
|
t = time.time() - start
|
|
|
|
# Input vector for the policy
|
|
# TODO: custom input
|
|
obs = _default_observations(data)
|
|
|
|
# Policy action
|
|
ctrl = policy.act(obs=obs, t=t)
|
|
|
|
# Check if the policy output vector give an input for each actuator (nu)
|
|
# TODO: what if model trained on full morphology but we want to test on a damaged one?
|
|
# (nu mismatch)
|
|
if model.nu > 0:
|
|
ctrl = np.asarray(ctrl, dtype=np.float32).ravel()
|
|
if ctrl.shape != (model.nu,):
|
|
raise ValueError(
|
|
f"Policy returned ctrl shape {ctrl.shape}, expected ({model.nu},)"
|
|
)
|
|
data.ctrl[:] = ctrl
|
|
|
|
# Step the simulation and update the viewer
|
|
mujoco.mj_step(model, data)
|
|
viewer.sync()
|
|
|
|
# If we're running in realtime mode, sleep to maintain real-time pacing.
|
|
if config.realtime:
|
|
remaining = model.opt.timestep - (time.time() - step_start)
|
|
if remaining > 0:
|
|
time.sleep(remaining)
|