Go2 DDS bridge is added.
This commit is contained in:
Rooholla-KhorramBakht
parent
a5749756c4
commit
030326f0b4
|
|
@ -0,0 +1,354 @@
|
|||
import os
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import onnxruntime
|
||||
from models import model_utils
|
||||
|
||||
class ActorStochasticMLP(nn.Module):
|
||||
def __init__(self):
|
||||
super(ActorStochasticMLP, self).__init__()
|
||||
|
||||
self.hidden_size = 256
|
||||
obs_dim = 49
|
||||
action_dim = 12
|
||||
self.gru = nn.GRU(obs_dim, self.hidden_size, batch_first=True)
|
||||
self.layer_dims = [obs_dim + self.hidden_size] + [256, 128, 64] + [action_dim]
|
||||
|
||||
init_ = lambda m: model_utils.init(m, nn.init.orthogonal_, lambda x: nn.init.
|
||||
constant_(x, 0), np.sqrt(2))
|
||||
|
||||
modules = []
|
||||
for i in range(len(self.layer_dims) - 1):
|
||||
modules.append(nn.Linear(self.layer_dims[i], self.layer_dims[i + 1]))
|
||||
if i < len(self.layer_dims) - 2:
|
||||
modules.append(model_utils.get_activation_func('elu'))
|
||||
modules.append(torch.nn.LayerNorm(self.layer_dims[i+1]))
|
||||
else:
|
||||
modules.append(model_utils.get_activation_func('identity'))
|
||||
|
||||
self.mu_net = nn.Sequential(*modules)
|
||||
|
||||
logstd = -1.0
|
||||
|
||||
self.logstd = torch.nn.Parameter(torch.ones(action_dim, dtype=torch.float32) * logstd)
|
||||
|
||||
self.action_dim = action_dim
|
||||
self.obs_dim = obs_dim
|
||||
|
||||
self.min_logstd = -1.427
|
||||
|
||||
print(self.mu_net)
|
||||
print(self.logstd)
|
||||
|
||||
def forward(self, obs, l):
|
||||
out, time_latent = self.gru(obs, l)
|
||||
x = torch.cat([obs, out], dim=-1)
|
||||
|
||||
mu = self.mu_net(x)
|
||||
|
||||
return mu, time_latent
|
||||
|
||||
class Policy:
|
||||
def __init__(self, checkpoint_path):
|
||||
self.agent = ActorStochasticMLP()
|
||||
agent_init, _, _, self.obs_rms, _, _ = torch.load(checkpoint_path, map_location="cpu")
|
||||
self.agent.load_state_dict(agent_init.state_dict())
|
||||
onnx_file_name = checkpoint_path.replace(".pt", ".onnx")
|
||||
self.actor_hidden_in = np.zeros((1, 256), dtype=np.float32)
|
||||
|
||||
dummy_input = torch.randn(1, 49)
|
||||
with torch.no_grad():
|
||||
torch_out, torch_hidden_out = self.agent(dummy_input, l = torch.tensor(self.actor_hidden_in))
|
||||
|
||||
torch.onnx.export(
|
||||
self.agent, # The model being converted
|
||||
(dummy_input, torch.tensor(self.actor_hidden_in)), # An example input for the model
|
||||
onnx_file_name, # Output file name
|
||||
export_params=True, # Store trained parameter weights inside the model file
|
||||
opset_version=11, # ONNX version (opset) to export to, adjust as needed
|
||||
do_constant_folding=True, # Whether to perform constant folding optimization
|
||||
input_names=['input', 'hidden_in'], # Name of the input in the ONNX graph (can be customized)
|
||||
output_names=['action', 'hidden_out'], # Name of the output (assuming get_action and get_value are key outputs)
|
||||
)
|
||||
|
||||
self.ort_session = onnxruntime.InferenceSession(onnx_file_name, providers=["CPUExecutionProvider"])
|
||||
ort_inputs = {'input': dummy_input.numpy(), 'hidden_in': self.actor_hidden_in}
|
||||
ort_outs = self.ort_session.run(None, ort_inputs)
|
||||
action, hidden_out = ort_outs[0], ort_outs[1]
|
||||
np.testing.assert_allclose(torch_out.numpy(), action, rtol=1e-03, atol=1e-05)
|
||||
np.testing.assert_allclose(torch_hidden_out.numpy(), hidden_out, rtol=1e-03, atol=1e-05)
|
||||
print("Exported model has been tested with ONNXRuntime, and the result looks good!")
|
||||
|
||||
def __call__(self, obs, info):
|
||||
obs = torch.cat(
|
||||
[torch.zeros((1, 3)), torch.tensor(obs[np.newaxis]), torch.zeros((1, 9))],
|
||||
dim = -1
|
||||
)
|
||||
obs = self.obs_rms.normalize(obs)[:, 3:-9].numpy().astype(np.float32)
|
||||
ort_inputs = {'input': obs, 'hidden_in': self.actor_hidden_in}
|
||||
ort_outs = self.ort_session.run(None, ort_inputs)
|
||||
self.actor_hidden_in = ort_outs[1]
|
||||
return ort_outs[0]
|
||||
|
||||
class CommandInterface:
|
||||
def __init__(self, limits=None):
|
||||
self.limits = limits
|
||||
self.x_vel_cmd, self.y_vel_cmd, self.yaw_vel_cmd = 0.0, 0.0, 0.0
|
||||
|
||||
def get_command(self):
|
||||
command = np.zeros((3,))
|
||||
command[0] = self.x_vel_cmd
|
||||
command[1] = self.y_vel_cmd
|
||||
command[2] = self.yaw_vel_cmd
|
||||
return command, False
|
||||
|
||||
|
||||
class NeuroDiffAgent:
|
||||
def __init__(self, command_profile, robot):
|
||||
self.robot = robot
|
||||
self.command_profile = command_profile
|
||||
# self.lcm_bridge = LCMBridgeClient(robot_name=self.robot_name)
|
||||
self.sim_dt = 0.001
|
||||
self.decimation = 20
|
||||
self.dt = self.sim_dt * self.decimation
|
||||
self.timestep = 0
|
||||
|
||||
self.device = "cpu"
|
||||
|
||||
joint_names = [
|
||||
"FL_hip_joint",
|
||||
"FL_thigh_joint",
|
||||
"FL_calf_joint",
|
||||
"FR_hip_joint",
|
||||
"FR_thigh_joint",
|
||||
"FR_calf_joint",
|
||||
"RL_hip_joint",
|
||||
"RL_thigh_joint",
|
||||
"RL_calf_joint",
|
||||
"RR_hip_joint",
|
||||
"RR_thigh_joint",
|
||||
"RR_calf_joint",
|
||||
]
|
||||
|
||||
policy_joint_names = joint_names
|
||||
|
||||
unitree_joint_names = [
|
||||
"FR_hip_joint",
|
||||
"FR_thigh_joint",
|
||||
"FR_calf_joint",
|
||||
"FL_hip_joint",
|
||||
"FL_thigh_joint",
|
||||
"FL_calf_joint",
|
||||
"RR_hip_joint",
|
||||
"RR_thigh_joint",
|
||||
"RR_calf_joint",
|
||||
"RL_hip_joint",
|
||||
"RL_thigh_joint",
|
||||
"RL_calf_joint",
|
||||
]
|
||||
|
||||
policy_to_unitree_map = []
|
||||
unitree_to_policy_map = []
|
||||
for i, policy_joint_name in enumerate(policy_joint_names):
|
||||
id = np.where([name == policy_joint_name for name in unitree_joint_names])[0][0]
|
||||
policy_to_unitree_map.append((i, id))
|
||||
self.policy_to_unitree_map = np.array(policy_to_unitree_map).astype(np.uint32)
|
||||
|
||||
for i, unitree_joint_name in enumerate(unitree_joint_names):
|
||||
id = np.where([name == unitree_joint_name for name in policy_joint_names])[0][0]
|
||||
unitree_to_policy_map.append((i, id))
|
||||
self.unitree_to_policy_map = np.array(unitree_to_policy_map).astype(np.uint32)
|
||||
|
||||
default_joint_angles = {
|
||||
"FL_hip_joint": 0.1,
|
||||
"RL_hip_joint": 0.1,
|
||||
"FR_hip_joint": -0.1,
|
||||
"RR_hip_joint": -0.1,
|
||||
"FL_thigh_joint": 0.8,
|
||||
"RL_thigh_joint": 1.0,
|
||||
"FR_thigh_joint": 0.8,
|
||||
"RR_thigh_joint": 1.0,
|
||||
"FL_calf_joint": -1.5,
|
||||
"RL_calf_joint": -1.5,
|
||||
"FR_calf_joint": -1.5,
|
||||
"RR_calf_joint": -1.5
|
||||
}
|
||||
self.default_dof_pos = np.array(
|
||||
[
|
||||
default_joint_angles[name]
|
||||
for name in joint_names
|
||||
]
|
||||
)
|
||||
|
||||
self.default_dof_pos = self.default_dof_pos
|
||||
|
||||
self.p_gains = np.zeros(12)
|
||||
self.d_gains = np.zeros(12)
|
||||
for i in range(12):
|
||||
self.p_gains[i] = 20.0
|
||||
self.d_gains[i] = 0.5
|
||||
|
||||
print(f"p_gains: {self.p_gains}")
|
||||
|
||||
self.commands = np.zeros(3)
|
||||
self.actions = np.zeros((1, 12))
|
||||
self.last_actions = np.zeros((1,12))
|
||||
self.gravity_vector = np.zeros(3)
|
||||
self.dof_pos = np.zeros(12)
|
||||
self.dof_vel = np.zeros(12)
|
||||
self.pos = np.zeros(3)
|
||||
self.body_linear_vel = np.zeros(3)
|
||||
self.body_angular_vel = np.zeros(3)
|
||||
self.joint_pos_target = np.zeros(12)
|
||||
self.joint_vel_target = np.zeros(12)
|
||||
self.prev_joint_acc = None
|
||||
self.torques = np.zeros(12)
|
||||
self.contact_state = np.ones(4)
|
||||
self.foot_contact_forces_mag = np.zeros(4)
|
||||
self.prev_foot_contact_forces_mag = np.zeros(4)
|
||||
self.test = 0
|
||||
|
||||
self.gait_indices = np.zeros(1, dtype=np.float32)
|
||||
self.clock_inputs = np.zeros(4, dtype=np.float32)
|
||||
|
||||
def wait_for_state(self):
|
||||
# return self.lcm_bridge.getStates(timeout=2)
|
||||
pass
|
||||
|
||||
def get_obs(self):
|
||||
cmds, reset_timer = self.command_profile.get_command()
|
||||
self.commands[:] = cmds
|
||||
|
||||
# self.state = self.wait_for_state()
|
||||
joint_state = self.robot.getJointStates()
|
||||
if joint_state is not None:
|
||||
self.gravity_vector = self.robot.getGravityInBody()
|
||||
self.prev_dof_pos = self.dof_pos.copy()
|
||||
self.dof_pos = np.array(joint_state['q'])[self.unitree_to_policy_map[:, 1]]
|
||||
self.prev_dof_vel = self.dof_vel.copy()
|
||||
self.dof_vel = np.array(joint_state['dq'])[self.unitree_to_policy_map[:, 1]]
|
||||
self.body_angular_vel = self.robot.getIMU()["gyro"]
|
||||
|
||||
try:
|
||||
self.foot_contact_forces_mag = self.robot.getFootContact()
|
||||
self.body_linear_vel = self.robot.getLinVel()
|
||||
self.pos, _ = self.robot.getPose()
|
||||
except:
|
||||
pass
|
||||
|
||||
ob = np.concatenate(
|
||||
(
|
||||
self.body_angular_vel * 0.25,
|
||||
self.commands * np.array([2.0, 2.0, 0.25]),
|
||||
self.gravity_vector[:, 0],
|
||||
self.dof_pos * 1.0,
|
||||
#((self.dof_pos - self.prev_dof_pos) / self.dt) * 0.05,
|
||||
self.dof_vel * 0.05,
|
||||
self.last_actions[0],
|
||||
self.clock_inputs
|
||||
),
|
||||
axis=0,
|
||||
)
|
||||
|
||||
#return torch.tensor(ob, device=self.device).float()
|
||||
return ob
|
||||
|
||||
def publish_action(self, action, hard_reset=False):
|
||||
# command_for_robot = UnitreeLowCommand()
|
||||
#self.joint_pos_target = (
|
||||
# action[0, :12].detach().cpu().numpy() * 0.25
|
||||
#).flatten()
|
||||
self.joint_pos_target = (
|
||||
action[0, :12] * 0.25
|
||||
).flatten()
|
||||
self.joint_pos_target += self.default_dof_pos
|
||||
self.joint_vel_target = np.zeros(12)
|
||||
# command_for_robot.q_des = self.joint_pos_target
|
||||
# command_for_robot.dq_des = self.joint_vel_target
|
||||
# command_for_robot.kp = self.p_gains
|
||||
# command_for_robot.kd = self.d_gains
|
||||
# command_for_robot.tau_ff = np.zeros(12)
|
||||
if hard_reset:
|
||||
command_for_robot.id = -1
|
||||
|
||||
self.torques = (self.joint_pos_target - self.dof_pos) * self.p_gains + (
|
||||
self.joint_vel_target - self.dof_vel
|
||||
) * self.d_gains
|
||||
# self.lcm_bridge.sendCommands(command_for_robot)
|
||||
self.robot.setCommands(self.joint_pos_target[self.policy_to_unitree_map[:, 1]],
|
||||
self.joint_vel_target[self.policy_to_unitree_map[:, 1]],
|
||||
self.p_gains[self.policy_to_unitree_map[:, 1]],
|
||||
self.d_gains[self.policy_to_unitree_map[:, 1]],
|
||||
np.zeros(12))
|
||||
|
||||
def reset(self):
|
||||
self.actions = torch.zeros((1, 12))
|
||||
self.time = time.time()
|
||||
self.timestep = 0
|
||||
return self.get_obs()
|
||||
|
||||
def step(self, actions, hard_reset=False):
|
||||
self.last_actions = self.actions[:]
|
||||
self.actions = actions
|
||||
self.publish_action(self.actions, hard_reset=hard_reset)
|
||||
# time.sleep(max(self.dt - (time.time() - self.time), 0))
|
||||
# if self.timestep % 100 == 0:
|
||||
# print(f"frq: {1 / (time.time() - self.time)} Hz")
|
||||
self.time = time.time()
|
||||
#obs = self.get_obs()
|
||||
|
||||
joint_acc = np.abs(self.prev_dof_vel - self.dof_vel) / self.dt
|
||||
if self.prev_joint_acc is None:
|
||||
self.prev_joint_acc = np.zeros_like(joint_acc)
|
||||
joint_jerk = np.abs(self.prev_joint_acc - joint_acc) / self.dt
|
||||
self.prev_joint_acc = joint_acc.copy()
|
||||
|
||||
# clock accounting
|
||||
frequencies = 3.
|
||||
phases = 0.5
|
||||
offsets = 0.
|
||||
bounds = 0
|
||||
self.gait_indices = np.remainder(
|
||||
self.gait_indices + self.dt * frequencies, 1.0
|
||||
)
|
||||
|
||||
self.foot_indices = [
|
||||
self.gait_indices + phases + offsets + bounds,
|
||||
self.gait_indices + offsets,
|
||||
self.gait_indices + bounds,
|
||||
self.gait_indices + phases,
|
||||
]
|
||||
|
||||
self.clock_inputs[0] = np.sin(2 * np.pi * self.foot_indices[0])
|
||||
self.clock_inputs[1] = np.sin(2 * np.pi * self.foot_indices[1])
|
||||
self.clock_inputs[2] = np.sin(2 * np.pi * self.foot_indices[2])
|
||||
self.clock_inputs[3] = np.sin(2 * np.pi * self.foot_indices[3])
|
||||
|
||||
foot_contact_rate = np.abs(self.foot_contact_forces_mag - self.prev_foot_contact_forces_mag)
|
||||
self.prev_foot_contact_forces_mag = self.foot_contact_forces_mag.copy()
|
||||
|
||||
infos = {
|
||||
"joint_pos": self.dof_pos[np.newaxis, :],
|
||||
"joint_vel": self.dof_vel[np.newaxis, :],
|
||||
"joint_pos_target": self.joint_pos_target[np.newaxis, :],
|
||||
"joint_vel_target": self.joint_vel_target[np.newaxis, :],
|
||||
"body_linear_vel": self.body_linear_vel[np.newaxis, :],
|
||||
"body_angular_vel": self.body_angular_vel[np.newaxis, :],
|
||||
"body_pos": self.pos[np.newaxis, :].copy(),
|
||||
"contact_state": self.contact_state[np.newaxis, :],
|
||||
"body_linear_vel_cmd": self.commands[np.newaxis, 0:2],
|
||||
"body_angular_vel_cmd": self.commands[np.newaxis, 2:],
|
||||
"torques": self.torques,
|
||||
"foot_contact_forces_mag": self.foot_contact_forces_mag.copy(),
|
||||
"joint_acc": joint_acc[np.newaxis, :],
|
||||
"joint_jerk": joint_jerk[np.newaxis, :],
|
||||
"foot_contact_rate": foot_contact_rate[np.newaxis, :],
|
||||
}
|
||||
|
||||
self.timestep += 1
|
||||
return None, None, None, infos
|
||||
|
|
@ -0,0 +1,329 @@
|
|||
# Copyright 2008 Willow Garage, Inc.
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# * Redistributions of source code must retain the above copyright
|
||||
# notice, this list of conditions and the following disclaimer.
|
||||
#
|
||||
# * Redistributions in binary form must reproduce the above copyright
|
||||
# notice, this list of conditions and the following disclaimer in the
|
||||
# documentation and/or other materials provided with the distribution.
|
||||
#
|
||||
# * Neither the name of the Willow Garage, Inc. nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
|
||||
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
|
||||
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
|
||||
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
|
||||
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
||||
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
|
||||
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
|
||||
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
||||
# POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
|
||||
"""
|
||||
Serialization of sensor_msgs.PointCloud2 messages.
|
||||
|
||||
Author: Tim Field
|
||||
ROS 2 port by Sebastian Grans
|
||||
File originally ported from:
|
||||
https://github.com/ros/common_msgs/blob/f48b00d43cdb82ed9367e0956db332484f676598/
|
||||
sensor_msgs/src/sensor_msgs/point_cloud2.py
|
||||
"""
|
||||
|
||||
import array
|
||||
from collections import namedtuple
|
||||
import sys
|
||||
from typing import Iterable, List, NamedTuple, Optional
|
||||
|
||||
import numpy as np
|
||||
try:
|
||||
from numpy.lib.recfunctions import (structured_to_unstructured, unstructured_to_structured)
|
||||
except ImportError:
|
||||
from sensor_msgs_py.numpy_compat import (structured_to_unstructured,
|
||||
unstructured_to_structured)
|
||||
|
||||
from sensor_msgs.msg import PointCloud2, PointField
|
||||
from std_msgs.msg import Header
|
||||
|
||||
|
||||
_DATATYPES = {}
|
||||
_DATATYPES[PointField.INT8] = np.dtype(np.int8)
|
||||
_DATATYPES[PointField.UINT8] = np.dtype(np.uint8)
|
||||
_DATATYPES[PointField.INT16] = np.dtype(np.int16)
|
||||
_DATATYPES[PointField.UINT16] = np.dtype(np.uint16)
|
||||
_DATATYPES[PointField.INT32] = np.dtype(np.int32)
|
||||
_DATATYPES[PointField.UINT32] = np.dtype(np.uint32)
|
||||
_DATATYPES[PointField.FLOAT32] = np.dtype(np.float32)
|
||||
_DATATYPES[PointField.FLOAT64] = np.dtype(np.float64)
|
||||
|
||||
DUMMY_FIELD_PREFIX = 'unnamed_field'
|
||||
|
||||
|
||||
def read_points(
|
||||
cloud: PointCloud2,
|
||||
field_names: Optional[List[str]] = None,
|
||||
skip_nans: bool = False,
|
||||
uvs: Optional[Iterable] = None,
|
||||
reshape_organized_cloud: bool = False) -> np.ndarray:
|
||||
"""
|
||||
Read points from a sensor_msgs.PointCloud2 message.
|
||||
|
||||
:param cloud: The point cloud to read from sensor_msgs.PointCloud2.
|
||||
:param field_names: The names of fields to read. If None, read all fields.
|
||||
(Type: Iterable, Default: None)
|
||||
:param skip_nans: If True, then don't return any point with a NaN value.
|
||||
(Type: Bool, Default: False)
|
||||
:param uvs: If specified, then only return the points at the given
|
||||
coordinates. (Type: Iterable, Default: None)
|
||||
:param reshape_organized_cloud: Returns the array as an 2D organized point cloud if set.
|
||||
:return: Structured NumPy array containing all points.
|
||||
"""
|
||||
assert isinstance(cloud, PointCloud2), \
|
||||
'Cloud is not a sensor_msgs.msg.PointCloud2'
|
||||
|
||||
# Cast bytes to numpy array
|
||||
points = np.ndarray(
|
||||
shape=(cloud.width * cloud.height, ),
|
||||
dtype=dtype_from_fields(cloud.fields, point_step=cloud.point_step),
|
||||
buffer=cloud.data)
|
||||
|
||||
# Keep only the requested fields
|
||||
if field_names is not None:
|
||||
assert all(field_name in points.dtype.names for field_name in field_names), \
|
||||
'Requests field is not in the fields of the PointCloud!'
|
||||
# Mask fields
|
||||
points = points[list(field_names)]
|
||||
|
||||
# Swap array if byte order does not match
|
||||
if bool(sys.byteorder != 'little') != bool(cloud.is_bigendian):
|
||||
points = points.byteswap(inplace=True)
|
||||
|
||||
# Check if we want to drop points with nan values
|
||||
if skip_nans and not cloud.is_dense:
|
||||
# Init mask which selects all points
|
||||
not_nan_mask = np.ones(len(points), dtype=bool)
|
||||
for field_name in points.dtype.names:
|
||||
# Only keep points without any non values in the mask
|
||||
not_nan_mask = np.logical_and(
|
||||
not_nan_mask, ~np.isnan(points[field_name]))
|
||||
# Select these points
|
||||
points = points[not_nan_mask]
|
||||
|
||||
# Select points indexed by the uvs field
|
||||
if uvs is not None:
|
||||
# Don't convert to numpy array if it is already one
|
||||
if not isinstance(uvs, np.ndarray):
|
||||
uvs = np.fromiter(uvs, int)
|
||||
# Index requested points
|
||||
points = points[uvs]
|
||||
|
||||
# Cast into 2d array if cloud is 'organized'
|
||||
if reshape_organized_cloud and cloud.height > 1:
|
||||
points = points.reshape(cloud.width, cloud.height)
|
||||
|
||||
return points
|
||||
|
||||
|
||||
def read_points_numpy(
|
||||
cloud: PointCloud2,
|
||||
field_names: Optional[List[str]] = None,
|
||||
skip_nans: bool = False,
|
||||
uvs: Optional[Iterable] = None,
|
||||
reshape_organized_cloud: bool = False) -> np.ndarray:
|
||||
"""
|
||||
Read equally typed fields from sensor_msgs.PointCloud2 message as a unstructured numpy array.
|
||||
|
||||
This method is better suited if one wants to perform math operations
|
||||
on e.g. all x,y,z fields.
|
||||
But it is limited to fields with the same dtype as unstructured numpy arrays
|
||||
only contain one dtype.
|
||||
|
||||
:param cloud: The point cloud to read from sensor_msgs.PointCloud2.
|
||||
:param field_names: The names of fields to read. If None, read all fields.
|
||||
(Type: Iterable, Default: None)
|
||||
:param skip_nans: If True, then don't return any point with a NaN value.
|
||||
(Type: Bool, Default: False)
|
||||
:param uvs: If specified, then only return the points at the given
|
||||
coordinates. (Type: Iterable, Default: None)
|
||||
:param reshape_organized_cloud: Returns the array as an 2D organized point cloud if set.
|
||||
:return: Numpy array containing all points.
|
||||
"""
|
||||
assert all(cloud.fields[0].datatype == field.datatype for field in cloud.fields[1:]), \
|
||||
'All fields need to have the same datatype. Use `read_points()` otherwise.'
|
||||
structured_numpy_array = read_points(
|
||||
cloud, field_names, skip_nans, uvs, reshape_organized_cloud)
|
||||
return structured_to_unstructured(structured_numpy_array)
|
||||
|
||||
|
||||
def read_points_list(
|
||||
cloud: PointCloud2,
|
||||
field_names: Optional[List[str]] = None,
|
||||
skip_nans: bool = False,
|
||||
uvs: Optional[Iterable] = None) -> List[NamedTuple]:
|
||||
"""
|
||||
Read points from a sensor_msgs.PointCloud2 message.
|
||||
|
||||
This function returns a list of namedtuples. It operates on top of the
|
||||
read_points method. For more efficient access use read_points directly.
|
||||
|
||||
:param cloud: The point cloud to read from. (Type: sensor_msgs.PointCloud2)
|
||||
:param field_names: The names of fields to read. If None, read all fields.
|
||||
(Type: Iterable, Default: None)
|
||||
:param skip_nans: If True, then don't return any point with a NaN value.
|
||||
(Type: Bool, Default: False)
|
||||
:param uvs: If specified, then only return the points at the given
|
||||
coordinates. (Type: Iterable, Default: None]
|
||||
:return: List of namedtuples containing the values for each point
|
||||
"""
|
||||
assert isinstance(cloud, PointCloud2), \
|
||||
'cloud is not a sensor_msgs.msg.PointCloud2'
|
||||
|
||||
if field_names is None:
|
||||
field_names = [f.name for f in cloud.fields]
|
||||
|
||||
Point = namedtuple('Point', field_names)
|
||||
|
||||
return [Point._make(p) for p in read_points(cloud, field_names,
|
||||
skip_nans, uvs)]
|
||||
|
||||
|
||||
def dtype_from_fields(fields: Iterable[PointField], point_step: Optional[int] = None) -> np.dtype:
|
||||
"""
|
||||
Convert a Iterable of sensor_msgs.msg.PointField messages to a np.dtype.
|
||||
|
||||
:param fields: The point cloud fields.
|
||||
(Type: iterable of sensor_msgs.msg.PointField)
|
||||
:param point_step: Point step size in bytes. Calculated from the given fields by default.
|
||||
(Type: optional of integer)
|
||||
:returns: NumPy datatype
|
||||
"""
|
||||
# Create a lists containing the names, offsets and datatypes of all fields
|
||||
field_names = []
|
||||
field_offsets = []
|
||||
field_datatypes = []
|
||||
for i, field in enumerate(fields):
|
||||
# Datatype as numpy datatype
|
||||
datatype = _DATATYPES[field.datatype]
|
||||
# Name field
|
||||
if field.name == '':
|
||||
name = f'{DUMMY_FIELD_PREFIX}_{i}'
|
||||
else:
|
||||
name = field.name
|
||||
# Handle fields with count > 1 by creating subfields with a suffix consiting
|
||||
# of "_" followed by the subfield counter [0 -> (count - 1)]
|
||||
assert field.count > 0, "Can't process fields with count = 0."
|
||||
for a in range(field.count):
|
||||
# Add suffix if we have multiple subfields
|
||||
if field.count > 1:
|
||||
subfield_name = f'{name}_{a}'
|
||||
else:
|
||||
subfield_name = name
|
||||
assert subfield_name not in field_names, 'Duplicate field names are not allowed!'
|
||||
field_names.append(subfield_name)
|
||||
# Create new offset that includes subfields
|
||||
field_offsets.append(field.offset + a * datatype.itemsize)
|
||||
field_datatypes.append(datatype.str)
|
||||
|
||||
# Create dtype
|
||||
dtype_dict = {
|
||||
'names': field_names,
|
||||
'formats': field_datatypes,
|
||||
'offsets': field_offsets
|
||||
}
|
||||
if point_step is not None:
|
||||
dtype_dict['itemsize'] = point_step
|
||||
return np.dtype(dtype_dict)
|
||||
|
||||
|
||||
def create_cloud(
|
||||
header: Header,
|
||||
fields: Iterable[PointField],
|
||||
points: Iterable) -> PointCloud2:
|
||||
"""
|
||||
Create a sensor_msgs.msg.PointCloud2 message.
|
||||
|
||||
:param header: The point cloud header. (Type: std_msgs.msg.Header)
|
||||
:param fields: The point cloud fields.
|
||||
(Type: iterable of sensor_msgs.msg.PointField)
|
||||
:param points: The point cloud points. List of iterables, i.e. one iterable
|
||||
for each point, with the elements of each iterable being the
|
||||
values of the fields for that point (in the same order as
|
||||
the fields parameter)
|
||||
:return: The point cloud as sensor_msgs.msg.PointCloud2
|
||||
"""
|
||||
# Check if input is numpy array
|
||||
if isinstance(points, np.ndarray):
|
||||
# Check if this is an unstructured array
|
||||
if points.dtype.names is None:
|
||||
assert all(fields[0].datatype == field.datatype for field in fields[1:]), \
|
||||
'All fields need to have the same datatype. Pass a structured NumPy array \
|
||||
with multiple dtypes otherwise.'
|
||||
# Convert unstructured to structured array
|
||||
points = unstructured_to_structured(
|
||||
points,
|
||||
dtype=dtype_from_fields(fields))
|
||||
else:
|
||||
assert points.dtype == dtype_from_fields(fields), \
|
||||
'PointFields and structured NumPy array dtype do not match for all fields! \
|
||||
Check their field order, names and types.'
|
||||
else:
|
||||
# Cast python objects to structured NumPy array (slow)
|
||||
points = np.array(
|
||||
# Points need to be tuples in the structured array
|
||||
list(map(tuple, points)),
|
||||
dtype=dtype_from_fields(fields))
|
||||
|
||||
# Handle organized clouds
|
||||
assert len(points.shape) <= 2, \
|
||||
'Too many dimensions for organized cloud! \
|
||||
Points can only be organized in max. two dimensional space'
|
||||
height = 1
|
||||
width = points.shape[0]
|
||||
# Check if input points are an organized cloud (2D array of points)
|
||||
if len(points.shape) == 2:
|
||||
height = points.shape[1]
|
||||
|
||||
# Convert numpy points to array.array
|
||||
memory_view = memoryview(points)
|
||||
casted = memory_view.cast('B')
|
||||
array_array = array.array('B')
|
||||
array_array.frombytes(casted)
|
||||
|
||||
# Put everything together
|
||||
cloud = PointCloud2(
|
||||
header=header,
|
||||
height=height,
|
||||
width=width,
|
||||
is_dense=False,
|
||||
is_bigendian=sys.byteorder != 'little',
|
||||
fields=fields,
|
||||
point_step=points.dtype.itemsize,
|
||||
row_step=(points.dtype.itemsize * width))
|
||||
# Set cloud via property instead of the constructor because of the bug described in
|
||||
# https://github.com/ros2/common_interfaces/issues/176
|
||||
cloud.data = array_array
|
||||
return cloud
|
||||
|
||||
|
||||
def create_cloud_xyz32(header: Header, points: Iterable) -> PointCloud2:
|
||||
"""
|
||||
Create a sensor_msgs.msg.PointCloud2 message with (x, y, z) fields.
|
||||
|
||||
:param header: The point cloud header. (Type: std_msgs.msg.Header)
|
||||
:param points: The point cloud points. (Type: Iterable)
|
||||
:return: The point cloud as sensor_msgs.msg.PointCloud2.
|
||||
"""
|
||||
fields = [PointField(name='x', offset=0,
|
||||
datatype=PointField.FLOAT32, count=1),
|
||||
PointField(name='y', offset=4,
|
||||
datatype=PointField.FLOAT32, count=1),
|
||||
PointField(name='z', offset=8,
|
||||
datatype=PointField.FLOAT32, count=1)]
|
||||
return create_cloud(header, fields, points)
|
||||
|
|
@ -9,6 +9,85 @@ from rclpy.qos import QoSProfile
|
|||
from rclpy.executors import MultiThreadedExecutor
|
||||
from geometry_msgs.msg import TransformStamped
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
from sensor_msgs.msg import PointCloud2
|
||||
from sensor_msgs_py import point_cloud2
|
||||
from sensor_msgs.msg import Image, CameraInfo
|
||||
from cv_bridge import CvBridge
|
||||
import cv2
|
||||
from Go2Py.utils.point_cloud2 import read_points_numpy
|
||||
import rclpy
|
||||
from rclpy.node import Node
|
||||
from nav_msgs.msg import Odometry
|
||||
|
||||
import rclpy
|
||||
from rclpy.node import Node
|
||||
from nav_msgs.msg import Odometry
|
||||
|
||||
import numpy as np
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
|
||||
class ROS2OdometryReader(Node):
|
||||
"""
|
||||
A class for subscribing to and retrieving Odometry messages.
|
||||
|
||||
Args:
|
||||
odom_topic (str): The topic name for the odometry data.
|
||||
node_name (str): A unique name for this node.
|
||||
|
||||
Attributes:
|
||||
odom (Odometry): The latest odometry message received from the topic.
|
||||
"""
|
||||
|
||||
def __init__(self, odom_topic, node_name):
|
||||
super().__init__(f'{node_name}_odom_reader')
|
||||
self.odom = None # Stores the latest odometry message
|
||||
|
||||
# Create a subscription to the specified odometry topic
|
||||
self.odom_subscriber = self.create_subscription(
|
||||
Odometry,
|
||||
odom_topic,
|
||||
self.odom_callback,
|
||||
10
|
||||
)
|
||||
|
||||
def odom_callback(self, msg):
|
||||
"""
|
||||
Callback function that stores the latest odometry message.
|
||||
|
||||
Args:
|
||||
msg (Odometry): The odometry message received.
|
||||
"""
|
||||
self.odom = msg
|
||||
|
||||
def get_odometry(self):
|
||||
"""
|
||||
Returns the latest odometry message received from the topic.
|
||||
|
||||
Returns:
|
||||
Odometry: The latest odometry message.
|
||||
"""
|
||||
return self.odom
|
||||
|
||||
def get_pose(self):
|
||||
if self.odom is not None:
|
||||
position = self.odom.pose.pose.position
|
||||
orientation = self.odom.pose.pose.orientation
|
||||
t = np.array([position.x, position.y, position.z])
|
||||
q = np.array([orientation.x, orientation.y, orientation.z, orientation.w])
|
||||
Rot = R.from_quat(q).as_matrix()
|
||||
T = np.eye(4)
|
||||
T[:3, :3] = Rot
|
||||
T[:3, -1] = t
|
||||
return T
|
||||
else:
|
||||
return None
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
Destroys the node, cleaning up resources.
|
||||
"""
|
||||
self.destroy_node()
|
||||
|
||||
def ros2_init(args=None):
|
||||
rclpy.init(args=args)
|
||||
|
||||
|
|
@ -49,10 +128,50 @@ class ROS2ExecutorManager:
|
|||
if self.executor_thread:
|
||||
self.executor_thread.join()
|
||||
|
||||
class ROS2TFInterface(Node):
|
||||
# class ROS2TFInterface(Node):
|
||||
|
||||
# def __init__(self, parent_name, child_name, node_name):
|
||||
# super().__init__(f'{node_name}_tf2_listener')
|
||||
# self.parent_name = parent_name
|
||||
# self.child_name = child_name
|
||||
# self.tfBuffer = tf2_ros.Buffer()
|
||||
# self.listener = tf2_ros.TransformListener(self.tfBuffer, self)
|
||||
# self.T = None
|
||||
# self.stamp = None
|
||||
# self.running = True
|
||||
# self.thread = threading.Thread(target=self.update_loop)
|
||||
# self.thread.start()
|
||||
# self.trans = None
|
||||
|
||||
# def update_loop(self):
|
||||
# while self.running:
|
||||
# try:
|
||||
# self.trans = self.tfBuffer.lookup_transform(self.parent_name, self.child_name, rclpy.time.Time(), rclpy.time.Duration(seconds=0.1))
|
||||
# except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
|
||||
# print(e)
|
||||
# time.sleep(0.01)
|
||||
|
||||
# def get_pose(self):
|
||||
# if self.trans is None:
|
||||
# return None
|
||||
# else:
|
||||
# translation = [self.trans.transform.translation.x, self.trans.transform.translation.y, self.trans.transform.translation.z]
|
||||
# rotation = [self.trans.transform.rotation.x, self.trans.transform.rotation.y, self.trans.transform.rotation.z, self.trans.transform.rotation.w]
|
||||
# self.T = np.eye(4)
|
||||
# self.T[0:3, 0:3] = R.from_quat(rotation).as_matrix()
|
||||
# self.T[:3, 3] = translation
|
||||
# self.stamp = self.trans.header.stamp.nanosec * 1e-9 + self.trans.header.stamp.sec
|
||||
# return self.T
|
||||
|
||||
# def close(self):
|
||||
# self.running = False
|
||||
# self.thread.join()
|
||||
# self.destroy_node()
|
||||
|
||||
|
||||
class ROS2TFInterface(Node):
|
||||
def __init__(self, parent_name, child_name, node_name):
|
||||
super().__init__(f'{node_name}_tf2_listener')
|
||||
super().__init__(f"{node_name}_tf2_listener")
|
||||
self.parent_name = parent_name
|
||||
self.child_name = child_name
|
||||
self.tfBuffer = tf2_ros.Buffer()
|
||||
|
|
@ -60,31 +179,137 @@ class ROS2TFInterface(Node):
|
|||
self.T = None
|
||||
self.stamp = None
|
||||
self.running = True
|
||||
self.thread = threading.Thread(target=self.update_loop)
|
||||
self.thread.start()
|
||||
self.trans = None
|
||||
|
||||
def update_loop(self):
|
||||
while self.running:
|
||||
try:
|
||||
self.trans = self.tfBuffer.lookup_transform(self.parent_name, self.child_name, rclpy.time.Time(), rclpy.time.Duration(seconds=0.1))
|
||||
except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
|
||||
print(e)
|
||||
time.sleep(0.01)
|
||||
|
||||
def get_pose(self):
|
||||
if self.trans is None:
|
||||
try:
|
||||
self.trans = self.tfBuffer.lookup_transform(
|
||||
self.parent_name,
|
||||
self.child_name,
|
||||
rclpy.time.Time(),
|
||||
rclpy.time.Duration(seconds=0.1),
|
||||
)
|
||||
except (
|
||||
tf2_ros.LookupException,
|
||||
tf2_ros.ConnectivityException,
|
||||
tf2_ros.ExtrapolationException,
|
||||
) as e:
|
||||
return None
|
||||
else:
|
||||
translation = [self.trans.transform.translation.x, self.trans.transform.translation.y, self.trans.transform.translation.z]
|
||||
rotation = [self.trans.transform.rotation.x, self.trans.transform.rotation.y, self.trans.transform.rotation.z, self.trans.transform.rotation.w]
|
||||
|
||||
translation = [
|
||||
self.trans.transform.translation.x,
|
||||
self.trans.transform.translation.y,
|
||||
self.trans.transform.translation.z,
|
||||
]
|
||||
rotation = [
|
||||
self.trans.transform.rotation.x,
|
||||
self.trans.transform.rotation.y,
|
||||
self.trans.transform.rotation.z,
|
||||
self.trans.transform.rotation.w,
|
||||
]
|
||||
self.T = np.eye(4)
|
||||
self.T[0:3, 0:3] = R.from_quat(rotation).as_matrix()
|
||||
self.T[:3, 3] = translation
|
||||
self.stamp = self.trans.header.stamp.nanosec * 1e-9 + self.trans.header.stamp.sec
|
||||
self.stamp = (
|
||||
self.trans.header.stamp.nanosec * 1e-9 + self.trans.header.stamp.sec
|
||||
)
|
||||
|
||||
return self.T
|
||||
|
||||
def close(self):
|
||||
self.running = False
|
||||
self.thread.join()
|
||||
self.destroy_node()
|
||||
|
||||
|
||||
class PoindCloud2Bridge(Node):
|
||||
def __init__(self, topic_name):
|
||||
super().__init__("point_cloud_listener")
|
||||
self.subscription = self.create_subscription(
|
||||
PointCloud2, topic_name, self.listener_callback, 1
|
||||
)
|
||||
self.data = None
|
||||
self.new_frame_flag = False
|
||||
self.points = None
|
||||
|
||||
self._points = None
|
||||
self.points = None
|
||||
self.ready = False
|
||||
self.running = True
|
||||
|
||||
def listener_callback(self, msg):
|
||||
# Read the x, y, z fields from the PointCloud2 message
|
||||
self._points = read_points_numpy(
|
||||
msg,
|
||||
field_names=("x", "y", "z"),
|
||||
skip_nans=True,
|
||||
reshape_organized_cloud=True,
|
||||
)
|
||||
# self._points = point_cloud2.read_points(
|
||||
# msg,
|
||||
# field_names=("x", "y", "z"),
|
||||
# skip_nans=True,
|
||||
# )
|
||||
# _points = np.reshape(_points, (-1, 16, 3))
|
||||
# self._points = np.reshape(_points, (-1, 16, 3))
|
||||
self.new_frame_flag = True
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
return self._points
|
||||
|
||||
class ROS2CameraReader(Node):
|
||||
"""
|
||||
A class for interacting with a ROS2 camera topics.
|
||||
|
||||
Args:
|
||||
image_topic (str): The topic name for the image stream.
|
||||
camera_info_topic (str, optional): The topic name for the camera information.
|
||||
K (numpy.ndarray, optional): The intrinsic camera matrix for the raw (distorted) images.
|
||||
D (numpy.ndarray, optional): The distortion coefficients.
|
||||
|
||||
Attributes:
|
||||
color_frame (numpy.ndarray): The latest image frame received from the topic.
|
||||
camera_info (CameraInfo): The latest camera information received from the topic.
|
||||
K (numpy.ndarray): The intrinsic camera matrix.
|
||||
D (numpy.ndarray): The distortion coefficients.
|
||||
"""
|
||||
def __init__(self, image_topic, node_name, camera_info_topic=None, K=None, D=None):
|
||||
super().__init__(f'{node_name}_camera_reader')
|
||||
self.bridge = CvBridge()
|
||||
self.color_frame = None
|
||||
self.camera_info = None
|
||||
self.K = K
|
||||
self.D = D
|
||||
self.node_name = node_name
|
||||
|
||||
self.image_subscriber = self.create_subscription(Image, image_topic, self.image_callback, 10)
|
||||
if camera_info_topic:
|
||||
self.camera_info_subscriber = self.create_subscription(CameraInfo, camera_info_topic, self.camera_info_callback, 10)
|
||||
|
||||
def image_callback(self, msg):
|
||||
self.color_frame = self.bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
|
||||
|
||||
def camera_info_callback(self, msg):
|
||||
self.camera_info = msg
|
||||
self.K = np.array(msg.k).reshape((3, 3))
|
||||
self.D = np.array(msg.d)
|
||||
|
||||
def get_image(self):
|
||||
"""
|
||||
Returns the latest image frame received from the topic.
|
||||
|
||||
Returns:
|
||||
numpy.ndarray: The latest image frame received from the topic.
|
||||
"""
|
||||
return self.color_frame
|
||||
|
||||
def get_intrinsics(self):
|
||||
"""
|
||||
Returns the intrinsic camera matrix.
|
||||
|
||||
Returns:
|
||||
numpy.ndarray: The intrinsic camera matrix.
|
||||
"""
|
||||
return {'K': self.K, 'D': self.D}
|
||||
def close(self):
|
||||
self.destroy_node()
|
||||
File diff suppressed because one or more lines are too long
|
|
@ -0,0 +1,43 @@
|
|||
module Go2DDSMsgs {
|
||||
struct PointCloud {
|
||||
sequence<float> x;
|
||||
sequence<float> y;
|
||||
sequence<float> z;
|
||||
string timestamp;
|
||||
};
|
||||
|
||||
struct DepthImage {
|
||||
uint16 width;
|
||||
uint16 height;
|
||||
sequence<uint16> data;
|
||||
string timestamp;
|
||||
};
|
||||
|
||||
struct RGBImage {
|
||||
uint16 width;
|
||||
uint16 height;
|
||||
sequence<uint8> data;
|
||||
string timestamp;
|
||||
};
|
||||
|
||||
struct Pose {
|
||||
double quat[4];
|
||||
double trans[3];
|
||||
string timestamp;
|
||||
};
|
||||
|
||||
struct HighCommand {
|
||||
double vx;
|
||||
double vy;
|
||||
double omega;
|
||||
string timestamp;
|
||||
};
|
||||
|
||||
struct Imu {
|
||||
double accel[3];
|
||||
double gyro[3];
|
||||
double quat[4];
|
||||
string timestamp;
|
||||
};
|
||||
|
||||
};
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
DDSMsgs
|
||||
|
||||
DepthImage
|
||||
HighCommand
|
||||
Imu
|
||||
PointCloud
|
||||
Pose
|
||||
RGBImage
|
||||
|
|
@ -0,0 +1,79 @@
|
|||
"""
|
||||
Generated by Eclipse Cyclone DDS idlc Python Backend
|
||||
Cyclone DDS IDL version: v0.10.5
|
||||
Module: Go2DDSMsgs
|
||||
IDL file: DDSMsgs.idl
|
||||
|
||||
"""
|
||||
|
||||
from enum import auto
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
import cyclonedds.idl as idl
|
||||
import cyclonedds.idl.annotations as annotate
|
||||
import cyclonedds.idl.types as types
|
||||
|
||||
# root module import for resolving types
|
||||
import Go2DDSMsgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@annotate.final
|
||||
@annotate.autoid("sequential")
|
||||
class PointCloud(idl.IdlStruct, typename="Go2DDSMsgs.PointCloud"):
|
||||
x: types.sequence[types.float32]
|
||||
y: types.sequence[types.float32]
|
||||
z: types.sequence[types.float32]
|
||||
timestamp: str
|
||||
|
||||
|
||||
@dataclass
|
||||
@annotate.final
|
||||
@annotate.autoid("sequential")
|
||||
class DepthImage(idl.IdlStruct, typename="Go2DDSMsgs.DepthImage"):
|
||||
width: types.uint16
|
||||
height: types.uint16
|
||||
data: types.sequence[types.uint16]
|
||||
timestamp: str
|
||||
|
||||
|
||||
@dataclass
|
||||
@annotate.final
|
||||
@annotate.autoid("sequential")
|
||||
class RGBImage(idl.IdlStruct, typename="Go2DDSMsgs.RGBImage"):
|
||||
width: types.uint16
|
||||
height: types.uint16
|
||||
data: types.sequence[types.uint8]
|
||||
timestamp: str
|
||||
|
||||
|
||||
@dataclass
|
||||
@annotate.final
|
||||
@annotate.autoid("sequential")
|
||||
class Pose(idl.IdlStruct, typename="Go2DDSMsgs.Pose"):
|
||||
quat: types.array[types.float64, 4]
|
||||
trans: types.array[types.float64, 3]
|
||||
timestamp: str
|
||||
|
||||
|
||||
@dataclass
|
||||
@annotate.final
|
||||
@annotate.autoid("sequential")
|
||||
class HighCommand(idl.IdlStruct, typename="Go2DDSMsgs.HighCommand"):
|
||||
vx: types.float64
|
||||
vy: types.float64
|
||||
omega: types.float64
|
||||
timestamp: str
|
||||
|
||||
|
||||
@dataclass
|
||||
@annotate.final
|
||||
@annotate.autoid("sequential")
|
||||
class Imu(idl.IdlStruct, typename="Go2DDSMsgs.Imu"):
|
||||
accel: types.array[types.float64, 3]
|
||||
gyro: types.array[types.float64, 3]
|
||||
quat: types.array[types.float64, 4]
|
||||
timestamp: str
|
||||
|
||||
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
"""
|
||||
Generated by Eclipse Cyclone DDS idlc Python Backend
|
||||
Cyclone DDS IDL version: v0.10.5
|
||||
Module: Go2DDSMsgs
|
||||
|
||||
"""
|
||||
|
||||
from ._DDSMsgs import DepthImage, HighCommand, Imu, PointCloud, Pose, RGBImage
|
||||
__all__ = ["DepthImage", "HighCommand", "Imu", "PointCloud", "Pose", "RGBImage", ]
|
||||
|
|
@ -0,0 +1,116 @@
|
|||
import sys
|
||||
sys.path.append('..')
|
||||
|
||||
from Go2Py.utils.ros2 import PoindCloud2Bridge, ros2_init, ros2_close, ROS2ExecutorManager, ROS2CameraReader
|
||||
from Go2Py.utils.ros2 import ROS2OdometryReader
|
||||
from Go2Py.robot.interface import GO2Real
|
||||
from dds_telemetry import Go2DDSServer
|
||||
import time
|
||||
import numpy as np
|
||||
import time
|
||||
import cv2
|
||||
|
||||
import pickle
|
||||
RECORD_DATASET = False
|
||||
DATASET_LENGTH = 120
|
||||
ros2_init()
|
||||
map_bridge = PoindCloud2Bridge('/Laser_map')
|
||||
fast_lio_odom = ROS2OdometryReader('/Odometry', 'odometry_subscriber')
|
||||
rgb_camera_bridge = ROS2CameraReader('/camera/color/image_raw', 'rgb_reader')
|
||||
depth_camera_bridge = ROS2CameraReader('/camera/depth/image_rect_raw', 'depth_reader')
|
||||
|
||||
ros2_exec_manager = ROS2ExecutorManager()
|
||||
ros2_exec_manager.add_node(map_bridge)
|
||||
ros2_exec_manager.add_node(fast_lio_odom)
|
||||
ros2_exec_manager.add_node(rgb_camera_bridge)
|
||||
ros2_exec_manager.add_node(depth_camera_bridge)
|
||||
ros2_exec_manager.start()
|
||||
|
||||
# The extrinsics of the MID360 LiDAR with respect to the robot body
|
||||
body_T_lidar = np.array([
|
||||
[ 0.9743701, 0.0, 0.2249511, 0.1870 ],
|
||||
[ 0.0, 1.0, 0.0, 0.0 ],
|
||||
[-0.2249511, 0.0, 0.9743701, 0.0803 ],
|
||||
[ 0.0, 0.0, 0.0, 1.0 ]
|
||||
])
|
||||
|
||||
robot = GO2Real(mode='highlevel')
|
||||
dds_server = Go2DDSServer(robot_name='go2')
|
||||
|
||||
def GetLidarPose():
|
||||
lio_T_lidar = fast_lio_odom.get_pose()
|
||||
if lio_T_lidar is not None:
|
||||
lio_T_body = lio_T_lidar @ np.linalg.inv(body_T_lidar)
|
||||
return lio_T_body
|
||||
else:
|
||||
return None
|
||||
|
||||
def GetLidarMap():
|
||||
pcd = map_bridge._points
|
||||
if pcd is not None:
|
||||
return pcd
|
||||
else:
|
||||
return None
|
||||
|
||||
print('Waitng for the ROS2 bridges to start up ...')
|
||||
time.sleep(5)
|
||||
print('Running the bridge loop')
|
||||
lio_pcd = None
|
||||
map_update_counter = 0
|
||||
|
||||
rgb_imgs = []
|
||||
depth_imgs = []
|
||||
lio_Ts_robot = []
|
||||
pcd_maps = []
|
||||
|
||||
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
while True:
|
||||
tic = time.time()
|
||||
lio_T_robot = GetLidarPose()
|
||||
if lio_T_robot is not None:
|
||||
dds_server.sendLidarOdom(lio_T_robot)
|
||||
if RECORD_DATASET:
|
||||
lio_Ts_robot.append(lio_T_robot)
|
||||
|
||||
|
||||
if map_update_counter%10==0:
|
||||
lio_pcd = GetLidarMap()
|
||||
if lio_pcd is not None:
|
||||
dds_server.sendMap(lio_pcd)
|
||||
if RECORD_DATASET:
|
||||
pcd_maps.append(lio_pcd.copy())
|
||||
|
||||
map_update_counter += 1
|
||||
|
||||
depth = depth_camera_bridge.get_image()
|
||||
rgb = rgb_camera_bridge.get_image()
|
||||
if depth is not None and rgb is not None:
|
||||
depth = cv2.resize(depth, (320, 240))
|
||||
rgb = cv2.resize(rgb, (320, 240))
|
||||
# print('sending image')
|
||||
dds_server.sendDepth(depth)
|
||||
dds_server.sendRGB(rgb)
|
||||
|
||||
if RECORD_DATASET:
|
||||
rgb_imgs.append(rgb.copy())
|
||||
depth_imgs.append(depth.copy())
|
||||
|
||||
|
||||
if time.time()-start_time >= DATASET_LENGTH:
|
||||
break
|
||||
|
||||
# Forward the rgb, depth, lio, lio_T_lidar to the computer through DDS
|
||||
toc = time.time()
|
||||
command = dds_server.getHighCmd()
|
||||
if command is not None:
|
||||
print('Recievied a command from the client. Forwarding to the robot...')
|
||||
robot.setCommandsHigh(command.vx, command.vy, command.omega)
|
||||
|
||||
print(f'{(toc-tic):0.3f}')
|
||||
while(time.time()-tic) < 0.1:
|
||||
time.sleep(0.001)
|
||||
|
||||
|
||||
|
|
@ -0,0 +1,111 @@
|
|||
import socket
|
||||
import struct
|
||||
import numpy as np
|
||||
import time
|
||||
import threading
|
||||
from cyclonedds.domain import DomainParticipant
|
||||
from cyclonedds.topic import Topic
|
||||
from cyclonedds.sub import DataReader
|
||||
from cyclonedds.pub import Publisher, DataWriter
|
||||
from dataclasses import dataclass
|
||||
from Go2DDSMsgs import PointCloud, RGBImage, DepthImage, Pose, HighCommand
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
|
||||
|
||||
def get_last_msg(reader, topic_type):
|
||||
""" """
|
||||
last_msg = reader.take()
|
||||
|
||||
if last_msg:
|
||||
while True:
|
||||
a = reader.take()
|
||||
if not a:
|
||||
break
|
||||
else:
|
||||
last_msg = a
|
||||
if last_msg:
|
||||
msg = last_msg[0]
|
||||
if type(msg) == topic_type:
|
||||
return msg
|
||||
else:
|
||||
return None
|
||||
|
||||
else:
|
||||
return None
|
||||
|
||||
class Go2DDSServer:
|
||||
def __init__(self, robot_name):
|
||||
self.robot_name = robot_name
|
||||
self.domain = DomainParticipant()
|
||||
self.rgb_topic = Topic(self.domain, f'{self.robot_name}_rgb', RGBImage)
|
||||
self.depth_topic = Topic(self.domain, f'{self.robot_name}_depth', DepthImage)
|
||||
self.map_topic = Topic(self.domain, f'{self.robot_name}_lio_map', PointCloud)
|
||||
self.lidar_odom_topic = Topic(self.domain, f'{self.robot_name}_lio_odom', Pose)
|
||||
self.high_cmd_topic = Topic(self.domain, f'{self.robot_name}_high_cmd', HighCommand)
|
||||
self.rgb_publisher = Publisher(self.domain)
|
||||
self.depth_publisher = Publisher(self.domain)
|
||||
self.map_publisher = Publisher(self.domain)
|
||||
self.lidar_odom_publisher = Publisher(self.domain)
|
||||
self.high_cmd_reader = DataReader(self.domain, self.high_cmd_topic)
|
||||
self.rgb_writer = DataWriter(self.rgb_publisher, self.rgb_topic)
|
||||
self.depth_writer = DataWriter(self.depth_publisher, self.depth_topic)
|
||||
self.map_writer = DataWriter(self.map_publisher, self.map_topic)
|
||||
self.lidar_odom_writer = DataWriter(self.lidar_odom_publisher, self.lidar_odom_topic)
|
||||
|
||||
|
||||
|
||||
|
||||
def sendRGB(self, rgb):
|
||||
self.rgb_writer.write(RGBImage(data=rgb.reshape(-1).tolist(), width=rgb.shape[1], height=rgb.shape[0], timestamp=''))
|
||||
|
||||
def sendDepth(self, depth):
|
||||
self.depth_writer.write(DepthImage(data=depth.reshape(-1).tolist(), width=depth.shape[1], height=depth.shape[0], timestamp=''))
|
||||
|
||||
def sendMap(self, pcd):
|
||||
map_pcd_mgs = PointCloud(x = pcd[:,0].tolist(),
|
||||
y = pcd[:,1].tolist(),
|
||||
z = pcd[:,2].tolist(),
|
||||
timestamp=''
|
||||
)
|
||||
self.map_writer.write(map_pcd_mgs)
|
||||
|
||||
def sendLidarOdom(self, lio_T_body):
|
||||
q = R.from_matrix(lio_T_body[:3,:3]).as_quat()
|
||||
t = lio_T_body[:3,3]
|
||||
pose_msg = Pose(quat=q.tolist(), trans=t.tolist(), timestamp='')
|
||||
self.lidar_odom_writer.write(pose_msg)
|
||||
|
||||
def getHighCmd(self):
|
||||
return get_last_msg(self.high_cmd_reader, HighCommand)
|
||||
|
||||
|
||||
class Go2DDSClient:
|
||||
def __init__(self, robot_name):
|
||||
self.robot_name = robot_name
|
||||
self.domain = DomainParticipant()
|
||||
self.rgb_topic = Topic(self.domain, f'{self.robot_name}_rgb', RGBImage)
|
||||
self.depth_topic = Topic(self.domain, f'{self.robot_name}_depth', DepthImage)
|
||||
self.map_topic = Topic(self.domain, f'{self.robot_name}_lio_map', PointCloud)
|
||||
self.lidar_odom_topic = Topic(self.domain, f'{self.robot_name}_lio_odom', Pose)
|
||||
self.high_cmd_topic = Topic(self.domain, f'{self.robot_name}_high_cmd', HighCommand)
|
||||
self.rgb_reader = DataReader(self.domain, self.rgb_topic)
|
||||
self.depth_reader = DataReader(self.domain, self.depth_topic)
|
||||
self.map_reader = DataReader(self.domain, self.map_topic)
|
||||
self.lidar_odom_reader = DataReader(self.domain, self.lidar_odom_topic)
|
||||
self.high_cmd_writer = DataWriter(self.domain, self.high_cmd_topic)
|
||||
|
||||
def getRGB(self):
|
||||
return get_last_msg(self.rgb_reader, RGBImage)
|
||||
|
||||
def getDepth(self):
|
||||
return get_last_msg(self.depth_reader, DepthImage)
|
||||
|
||||
def getMap(self):
|
||||
return get_last_msg(self.map_reader, PointCloud)
|
||||
|
||||
def getLidarOdom(self):
|
||||
return get_last_msg(self.lidar_odom_reader, Pose)
|
||||
|
||||
def sendHighCmd(self, vx, vy, omega):
|
||||
self.high_cmd_writer.write(HighCommand(vx=vx, vy=vy, omega=omega, timestamp=''))
|
||||
|
||||
|
|
@ -0,0 +1,53 @@
|
|||
from Go2Py.robot.fsm import FSM
|
||||
from Go2Py.robot.safety import SafetyHypervisor
|
||||
from Go2Py.robot.interface import GO2Real
|
||||
from Go2Py.control.neuro_diff import NeuroDiffAgent
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
import time
|
||||
|
||||
|
||||
class NeuroDiffController:
|
||||
def __init__(self, robot, remote, checkpoint):
|
||||
self.remote = remote
|
||||
self.robot = robot
|
||||
self.policy = Policy(checkpoint)
|
||||
self.command_profile = CommandInterface()
|
||||
self.agent = NeuroDiffSimAgent(self.command_profile, self.robot)
|
||||
self.hist_data = {}
|
||||
|
||||
def init(self):
|
||||
self.obs = self.agent.reset()
|
||||
self.policy_info = {}
|
||||
self.command_profile.yaw_vel_cmd = 0.0
|
||||
self.command_profile.x_vel_cmd = 0.0
|
||||
self.command_profile.y_vel_cmd = 0.0
|
||||
|
||||
def update(self, robot, remote):
|
||||
if not hasattr(self, "obs"):
|
||||
self.init()
|
||||
commands = getRemote(remote)
|
||||
self.command_profile.yaw_vel_cmd = 0.0
|
||||
self.command_profile.x_vel_cmd = 0.0
|
||||
self.command_profile.y_vel_cmd = 0.0
|
||||
|
||||
self.obs = self.agent.get_obs()
|
||||
action = self.policy(self.obs, self.policy_info)
|
||||
_, self.ret, self.done, self.info = self.agent.step(action)
|
||||
for key, value in self.info.items():
|
||||
if key in self.hist_data:
|
||||
self.hist_data[key].append(value)
|
||||
else:
|
||||
self.hist_data[key] = [value]
|
||||
|
||||
|
||||
|
||||
robot = GO2Real(mode='lowlevel')
|
||||
safety_hypervisor = SafetyHypervisor(robot)
|
||||
|
||||
time.sleep(3)
|
||||
|
||||
print(robot.getJointStates())
|
||||
|
||||
robot.close()
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
from cyclonedds.topic import Topic
|
||||
from cyclonedds.pub import DataWriter
|
||||
|
||||
from cyclonedds.domain import DomainParticipant
|
||||
from cyclonedds.topic import Topic
|
||||
from cyclonedds.sub import DataReader
|
||||
from cyclonedds.util import duration
|
||||
import numpy as np
|
||||
|
||||
from Go2DDSMsgs import PointCloud, RGBImage, DepthImage, Pose, HighCommand
|
||||
# from scipy.spatial.transform import Rotation as R
|
||||
|
||||
participant = DomainParticipant()
|
||||
topic = Topic(participant, 'test_topic', RGBImage)
|
||||
writer = DataWriter(participant, topic)
|
||||
|
||||
# domain = DomainParticipant()
|
||||
# rgb_topic = Topic(domain, 'go2_rgb', RGBImage)
|
||||
# rgb_publisher = Publisher(domain)
|
||||
# rgb_writer = DataWriter(rgb_publisher, rgb_topic)
|
||||
|
||||
|
||||
rgb = np.random.randint(0, 255, (10, 10, 3)).astype(np.uint8)
|
||||
|
||||
rgb_msg = RGBImage(data=[0, 1, 2], timestamp='')
|
||||
writer.write(rgb_msg)
|
||||
Loading…
Reference in New Issue