File size: 2,256 Bytes
700dd75 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | from contextlib import contextmanager
import time
import gymnasium as gym
import numpy as np
from pinocchio.visualize import MeshcatVisualizer
from decoupled_wbc.control.base.env import Env
from decoupled_wbc.control.robot_model import RobotModel
class MeshcatVisualizerEnv(Env):
def __init__(self, robot_model: RobotModel):
self.robot_model = robot_model
self.viz = MeshcatVisualizer(
self.robot_model.pinocchio_wrapper.model,
self.robot_model.pinocchio_wrapper.collision_model,
self.robot_model.pinocchio_wrapper.visual_model,
)
try:
self.viz.initViewer(open=True)
except ImportError as err:
print("Error while initializing the viewer. It seems you should install Python meshcat")
print(err)
exit(0)
self.viz.loadViewerModel()
self.visualize(self.robot_model.pinocchio_wrapper.q0)
time.sleep(1.0)
self._observation_space = gym.spaces.Dict(
{
"q": gym.spaces.Box(
low=-2 * np.pi, high=2 * np.pi, shape=(self.robot_model.num_dofs,)
)
}
)
self._action_space = gym.spaces.Dict(
{
"q": gym.spaces.Box(
low=-2 * np.pi, high=2 * np.pi, shape=(self.robot_model.num_dofs,)
)
}
)
def visualize(self, robot_state: np.ndarray):
# visualize robot state
if robot_state is not None:
self.viz.display(robot_state)
def observe(self):
# Dummy observation
return {"q": self.robot_model.pinocchio_wrapper.q0}
def queue_action(self, action: dict[str, np.ndarray]):
self.visualize(action["q"])
def reset(self, **kwargs):
self.visualize(self.robot_model.pinocchio_wrapper.q0)
return {"q": self.robot_model.pinocchio_wrapper.q0}
def sensors(self) -> dict[str, any]:
return {}
def observation_space(self) -> gym.Space:
return self._observation_space
def action_space(self) -> gym.Space:
return self._action_space
def close(self):
return
@contextmanager
def activate(self):
yield
|