import gymnasium as gym
import matplotlib.pyplot as plt
import newton
import numpy as np
import warp as wp
from gymnasium import spaces
from lwmr.utils import create_viewer_viser
from tqdm.auto import trange
wp.config.quiet = TrueGym Interface
class RotatingCubeEnv(gym.Env):
metadata = {"render_modes": ["viser", "none"], "render_fps": 60}
def __init__(
self,
render_mode=None,
frame_step=1.0 / 60.0,
sim_substeps=4,
max_episode_steps=400,
max_target_velocity=10.0,
desired_velocity=8.0,
target_kd=100.0,
use_graph=False,
seed=None,
):
super().__init__()
self.render_mode = render_mode
self.frame_step = float(frame_step)
self.sim_substeps = int(sim_substeps)
self.time_step = self.frame_step / self.sim_substeps
self.max_episode_steps = int(max_episode_steps)
self.max_target_velocity = float(max_target_velocity)
self.desired_velocity = float(desired_velocity)
self.target_kd = float(target_kd)
self.use_graph = bool(use_graph)
self.sim_time = 0.0
self.step_count = 0
self.viewer = None
self.graph = None
self.np_random = None
self._seed = seed
# One normalized velocity command.
self.action_space = spaces.Box(
low=np.array([-1.0], dtype=np.float32),
high=np.array([1.0], dtype=np.float32),
dtype=np.float32,
)
# joint position, joint velocity
self.observation_space = spaces.Box(
low=np.array([-np.inf, -np.inf], dtype=np.float32),
high=np.array([np.inf, np.inf], dtype=np.float32),
dtype=np.float32,
)
self._build_model()
# ---------------------------------------------------------------------
# region Model
# ---------------------------------------------------------------------
def _build_model(self):
builder = newton.ModelBuilder()
builder.add_ground_plane()
xform = wp.transform(p=wp.vec3(0.0, 0.0, 1.0))
body = builder.add_link(mass=1.0)
joint = builder.add_joint_revolute(
parent=-1,
child=body,
parent_xform=xform,
axis=wp.vec3(0.0, 0.0, 1.0),
actuator_mode=newton.JointTargetMode.VELOCITY,
target_kd=self.target_kd,
)
builder.add_articulation([joint])
builder.add_shape_box(body)
# self.builder = builder
# self.body = body
# self.joint = joint
self.model = builder.finalize()
self.state_0 = self.model.state()
self.state_1 = self.model.state()
self.control = self.model.control()
self.pipeline = newton.CollisionPipeline(self.model)
self.contacts = self.pipeline.contacts()
self.solver = newton.solvers.SolverMuJoCo(self.model)
self.joint_index = builder.joint_qd_start[joint]
self.joint_target_vels = np.zeros(self.control.joint_target_qd.shape, dtype=np.float32)
if self.use_graph and self.model.device.is_cuda and wp.get_device().is_cuda:
with wp.ScopedCapture() as capture:
self._simulate()
self.graph = capture.graph
else:
self.graph = None
if self.render_mode == "viser":
# self.viewer = create_viewer("gym_cube", self.model)
self.viewer = create_viewer_viser("spinning_cube", self.model, quiet=False, overwrite=False)
# ---------------------------------------------------------------------
# region Gym
# ---------------------------------------------------------------------
def reset(self, *, seed=None, options=None):
super().reset(seed=seed)
if self.step_count == 0:
return self._get_obs(), {}
if seed is not None:
self._seed = seed
self.sim_time = 0.0
self.step_count = 0
self._build_model()
obs = self._get_obs()
info = {}
return obs, info
def step(self, action):
action = np.asarray(action, dtype=np.float32)
action = np.clip(action, self.action_space.low, self.action_space.high) # type: ignore
target_vel = float(action[0]) * self.max_target_velocity
self._set_target_velocity(target_vel)
if self.graph is not None:
wp.capture_launch(self.graph)
else:
self._simulate()
self.sim_time += self.frame_step
self.step_count += 1
obs = self._get_obs()
reward, reward_info = self._compute_reward(action=action, target_vel=target_vel, obs=obs)
terminated = False
truncated = self.step_count >= self.max_episode_steps
info = {
"sim_time": self.sim_time,
"step_count": self.step_count,
"target_velocity": target_vel,
**reward_info,
}
if self.render_mode == "viser":
self.render()
return obs, reward, terminated, truncated, info
def render(self):
if self.viewer:
self.viewer.begin_frame(self.sim_time)
self.viewer.log_state(self.state_0)
self.viewer.end_frame()
def close(self):
if self.viewer:
self.viewer.close()
# ---------------------------------------------------------------------
# Simulation helpers
# ---------------------------------------------------------------------
def _simulate(self):
for _ in range(self.sim_substeps):
self.state_0.clear_forces()
self.pipeline.collide(self.state_0, self.contacts)
self.solver.step(self.state_0, self.state_1, self.control, self.contacts, self.time_step)
self.state_0, self.state_1 = self.state_1, self.state_0
def _set_target_velocity(self, target_vel):
self.joint_target_vels.fill(0.0)
self.joint_target_vels[self.joint_index] = target_vel
self.control.joint_target_qd.assign(self.joint_target_vels)
def _get_obs(self):
joint_q = self.state_0.joint_q.numpy()
joint_qd = self.state_0.joint_qd.numpy()
q = float(joint_q[self.joint_index]) if joint_q.size > 0 else 0.0
qd = float(joint_qd[self.joint_index]) if joint_qd.size > 0 else 0.0
obs = np.array([q, qd], dtype=np.float32)
if not np.all(np.isfinite(obs)):
raise RuntimeError(f"Non-finite observation: {obs}")
return obs
def _compute_reward(self, action, target_vel, obs):
q, qd = float(obs[0]), float(obs[1])
# Velocity tracking reward.
vel_error = qd - self.desired_velocity
# Bounded reward around desired velocity.
tracking_reward = -0.01 * vel_error * vel_error
# Small action penalty.
action_penalty = -0.001 * float(np.sum(np.square(action)))
reward = tracking_reward + action_penalty
info = {
"joint_position": q,
"joint_velocity": qd,
"desired_velocity": self.desired_velocity,
"velocity_error": vel_error,
"tracking_reward": tracking_reward,
"action_penalty": action_penalty,
"reward": reward,
}
return float(reward), infoenv = RotatingCubeEnv(render_mode="viser")
obs, info = env.reset()
obs_history = [obs.copy()]
time = 0.0
for _ in trange(400):
vel = np.sin(2 * np.pi * 0.5 * time) * 0.5 + 0.5
action = np.array([vel], dtype=np.float32)
obs, reward, terminated, truncated, info = env.step(action)
obs_history.append(obs.copy())
time += env.frame_step
if terminated or truncated:
obs, info = env.reset()
# env.close()
env.unwrapped.viewer.show_notebook() # type: ignoreWarp 1.18.0 initialized:
CUDA not enabled in this build
Devices:
"cpu" : "arm"
Kernel cache:
~/Library/Caches/warp/1.18.0
Module validate_and_correct_inertia_kernel_1ff2ac62 6ad45b8 load on device 'cpu' took 1.25 ms (cached)
Module newton._src.geometry.bvh 5fa2076 load on device 'cpu' took 0.75 ms (cached)
Module newton._src.solvers.mujoco.kernels c5847ab load on device 'cpu' took 0.87 ms (cached)
Module mujoco_warp._src.io 593bb82 load on device 'cpu' took 1.31 ms (cached)
Module mujoco_warp._src.smooth 5a7824c load on device 'cpu' took 0.77 ms (cached)
Module _small_cholesky_factorize_block__locals__kernel_d04adac2 ad56182 load on device 'cpu' took 0.62 ms (cached)
Module _small_cholesky_solve_block__locals__kernel_b5e4dd90 91a86cb load on device 'cpu' took 0.45 ms (cached)
Recording to docs/_static/spinning_cube_04.viser...╭────── viser (listening *:8080) ───────╮ │ ╷ │ │ HTTP │ http://localhost:8080 │ │ Websocket │ ws://localhost:8080 │ │ ╵ │ ╰───────────────────────────────────────╯
Module newton._src.viewer.kernels 3942d1f load on device 'cpu' took 0.55 ms (cached)
0%| | 0/400 [00:00<?, ?it/s]Module newton._src.sim.collide a0e9da6 load on device 'cpu' took 0.74 ms (cached)
Module newton._src.geometry.broad_phase_nxn 4413549 load on device 'cpu' took 1.10 ms (cached)
Module narrow_phase_primitive_write_contact 066b49c load on device 'cpu' took 0.45 ms (cached)
Module narrow_phase_gjk_mpr_True_write_contact_support_map_lean_post_process_minkowski_only 2915b37 load on device 'cpu' took 0.52 ms (cached)
Module newton._src.geometry.narrow_phase 72cccae load on device 'cpu' took 0.50 ms (cached)
Module _nxn_broadphase__locals__kernel_5ea22428 e8b65c2 load on device 'cpu' took 0.57 ms (cached)
Module _primitive_narrowphase__locals__primitive_narrowphase_84c518b0 cd90311 load on device 'cpu' took 0.75 ms (cached)
Module mujoco_warp._src.constraint 90a64f5 load on device 'cpu' took 0.49 ms (cached)
Module _friction_dof__locals__kernel_eadd710f ffecde8 load on device 'cpu' took 0.84 ms (cached)
Module _efc_contact_init__locals__kernel_574ccaaf e931a35 load on device 'cpu' took 0.74 ms (cached)
⋮ (26 progress updates omitted)Module _efc_contact_jac_dense__locals__kernel_cf20828f ddac338 load on device 'cpu' took 0.82 ms (cached)
Module _efc_contact_update__locals__kernel_7d5c4160 5ca7f62 load on device 'cpu' took 0.50 ms (cached)
Module mujoco_warp._src.forward b904ce8 load on device 'cpu' took 1.94 ms (cached)
Module mujoco_warp._src.passive e5677c9 load on device 'cpu' took 0.73 ms (cached)
Module _qfrc_smooth__locals__kernel_8b823980 59d943b load on device 'cpu' took 0.46 ms (cached)
Module mujoco_warp._src.support 79faa31 load on device 'cpu' took 0.41 ms (cached)
Module _small_cholesky_factorize_solve_block__locals__kernel_75195291 69430d9 load on device 'cpu' took 1.04 ms (cached)
Module _solve_init_dof__locals__kernel_1d7ecfa4 b626280 load on device 'cpu' took 0.38 ms (cached)
Module mujoco_warp._src.solver e441f3d load on device 'cpu' took 0.42 ms (cached)
Module _solve_init_jaref_kernel__locals__kernel_2587957b 6b7a8d0 load on device 'cpu' took 0.42 ms (cached)
Module mul_m_kernel__locals___mul_m_d74bd8d9 ec4adf4 load on device 'cpu' took 0.46 ms (cached)
Module _update_constraint_efc__locals__kernel_4b172e8b 1b383ee load on device 'cpu' took 0.41 ms (cached)
Module _update_constraint_init_qfrc_constraint_dense__locals__kernel_c29c2cfc 711ed07 load on device 'cpu' took 0.40 ms (cached)
Module _update_gradient_zero_grad_dot__locals__kernel_93442399 d6cac50 load on device 'cpu' took 0.39 ms (cached)
Module _update_gradient_grad__locals__kernel_378e7517 71dbe85 load on device 'cpu' took 0.38 ms (cached)
Module _update_gradient_JTDAJ_dense_tiled__locals__kernel_d4d12986 cc4fd44 load on device 'cpu' took 0.40 ms (cached)
Module _update_gradient_cholesky__locals__kernel_103bfa90 035ac3e load on device 'cpu' took 0.36 ms (cached)
Module _linesearch_iterative_kernel__locals__kernel_07b6d583 880e174 load on device 'cpu' took 0.42 ms (cached)
Module _update_constraint_efc__locals__kernel_523090d1 2255645 load on device 'cpu' took 0.41 ms (cached)
Module _update_constraint_init_qfrc_constraint_dense__locals__kernel_f3231cfe 951fff2 load on device 'cpu' took 0.40 ms (cached)
Module _update_gradient_zero_grad_dot__locals__kernel_721e3cea 93a9a23 load on device 'cpu' took 0.42 ms (cached)
Module _update_gradient_grad__locals__kernel_5363029b 67d94ed load on device 'cpu' took 0.39 ms (cached)
Module _update_gradient_cholesky__locals__kernel_d2bec93a d50b456 load on device 'cpu' took 0.39 ms (cached)
Module mujoco_warp._src.sensor 8b0053e load on device 'cpu' took 0.67 ms (cached)
Module mujoco_warp._src.derivative 65e089c load on device 'cpu' took 0.58 ms (cached)
Module _next_time_builder__locals___next_time_f604e0b3 324f4ea load on device 'cpu' took 0.55 ms (cached)
Module newton._src.sim.articulation 1e9d119 load on device 'cpu' took 0.47 ms (cached)
Recording to docs/_static/spinning_cube_04.viser...╭────── viser (listening *:8081) ───────╮ │ ╷ │ │ HTTP │ http://localhost:8081 │ │ Websocket │ ws://localhost:8081 │ │ ╵ │ ╰───────────────────────────────────────╯
100%|██████████| 400/400 [00:02<00:00, 138.19it/s]plt.plot(np.array(obs_history)[:, 0], label="joint_position")
plt.plot(np.array(obs_history)[:, 1], label="joint_velocity")
plt.legend()
plt.title("Observation History")
plt.xlabel("Step");