feat: add head control and joint data decoder
This commit is contained in:
377
tg3_data_collection/decode_joint_data.py
Executable file
377
tg3_data_collection/decode_joint_data.py
Executable file
@@ -0,0 +1,377 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Decode TG3 joint streams from one locally stored MCAP episode into CSV."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
|
||||
ARM_MOTOR_IDS = [*range(11, 18), *range(21, 28)]
|
||||
JOINT_NAMES = [
|
||||
*(f"left_joints_{index}" for index in range(7)),
|
||||
*(f"right_joints_{index}" for index in range(7)),
|
||||
]
|
||||
CUSTOM_TOPICS = {
|
||||
"/robot_state": "ros2_bridge_msgs/msg/RobotState",
|
||||
"/freq_change/arm_status": "ros2_bridge_msgs/msg/ArmStatus",
|
||||
"/data_logger/arm_status": "ros2_bridge_msgs/msg/ArmStatus",
|
||||
"/arm/cmd": "ros2_bridge_msgs/msg/ArmCtrl",
|
||||
}
|
||||
TOPICS = [
|
||||
"/robot_state",
|
||||
"/encoder_identical_joint",
|
||||
"/freq_change/arm_status",
|
||||
"/data_logger/arm_status",
|
||||
"/arm/cmd",
|
||||
"/tg3/data_collection/iarm_frame",
|
||||
]
|
||||
|
||||
|
||||
def _message_directory() -> Path:
|
||||
return (
|
||||
Path(__file__).resolve().parents[1]
|
||||
/ "tg3_local_teleop"
|
||||
/ "ros2_py"
|
||||
/ "src"
|
||||
/ "ros2_bridge_msgs"
|
||||
/ "msg"
|
||||
)
|
||||
|
||||
|
||||
def _read_definition(name: str) -> str:
|
||||
path = _message_directory() / f"{name}.msg"
|
||||
if not path.is_file():
|
||||
raise RuntimeError(f"required message definition is missing: {path}")
|
||||
return path.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def _definition_bundle(root_name: str, dependencies: list[tuple[str, str]]) -> str:
|
||||
if root_name == "ArmCtrl":
|
||||
text = (
|
||||
"std_msgs/Header header\n"
|
||||
"uint8 mode\n"
|
||||
"uint8 label\n"
|
||||
"uint8 reserved\n"
|
||||
"ros2_bridge_msgs/MotorCtrl[] ctrl\n"
|
||||
)
|
||||
else:
|
||||
text = _read_definition(root_name)
|
||||
for full_name, definition in dependencies:
|
||||
text += f"\n===\nMSG: {full_name}\n{definition}"
|
||||
return text
|
||||
|
||||
|
||||
def _custom_decoders() -> dict[str, Callable[[bytes], Any]]:
|
||||
try:
|
||||
from mcap_ros2._dynamic import generate_dynamic
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"missing MCAP decoder; install with: pip install mcap mcap-ros2-support"
|
||||
) from exc
|
||||
|
||||
header = "builtin_interfaces/Time stamp\nstring frame_id\n"
|
||||
motor_status = _read_definition("MotorStatus")
|
||||
common = [
|
||||
("std_msgs/msg/Header", header),
|
||||
("ros2_bridge_msgs/msg/MotorStatus", motor_status),
|
||||
]
|
||||
robot_dependencies = common + [
|
||||
(f"ros2_bridge_msgs/msg/{name}", _read_definition(name))
|
||||
for name in (
|
||||
"HeadStatus",
|
||||
"WaistStatus",
|
||||
"LegStatus",
|
||||
"ArmStatus",
|
||||
"ImuStatus",
|
||||
)
|
||||
]
|
||||
schemas = {
|
||||
"ros2_bridge_msgs/msg/RobotState": _definition_bundle(
|
||||
"RobotState", robot_dependencies
|
||||
),
|
||||
"ros2_bridge_msgs/msg/ArmStatus": _definition_bundle(
|
||||
"ArmStatus", common
|
||||
),
|
||||
"ros2_bridge_msgs/msg/ArmCtrl": _definition_bundle(
|
||||
"ArmCtrl",
|
||||
[
|
||||
("std_msgs/msg/Header", header),
|
||||
("ros2_bridge_msgs/msg/MotorCtrl", _read_definition("MotorCtrl")),
|
||||
],
|
||||
),
|
||||
}
|
||||
return {
|
||||
type_name: generate_dynamic(type_name, text)[type_name]
|
||||
for type_name, text in schemas.items()
|
||||
}
|
||||
|
||||
|
||||
def _stamp_ns(header: Any) -> int:
|
||||
return int(header.stamp.sec) * 1_000_000_000 + int(header.stamp.nanosec)
|
||||
|
||||
|
||||
def _side_and_index(motor_id: int) -> tuple[str, int]:
|
||||
if 11 <= motor_id <= 17:
|
||||
return "left", motor_id - 11
|
||||
if 21 <= motor_id <= 27:
|
||||
return "right", motor_id - 21
|
||||
return "unknown", -1
|
||||
|
||||
|
||||
def _open_csv(path: Path, columns: list[str]) -> tuple[Any, csv.DictWriter]:
|
||||
stream = path.open("x", encoding="utf-8", newline="")
|
||||
writer = csv.DictWriter(stream, fieldnames=columns)
|
||||
writer.writeheader()
|
||||
return stream, writer
|
||||
|
||||
|
||||
def decode_episode(episode: Path, output: Path) -> dict[str, int]:
|
||||
try:
|
||||
from mcap.reader import make_reader
|
||||
from mcap_ros2.decoder import DecoderFactory
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"missing MCAP decoder; install with: pip install mcap mcap-ros2-support"
|
||||
) from exc
|
||||
|
||||
bag_files = sorted((episode / "bag").glob("*.mcap"))
|
||||
if not bag_files:
|
||||
raise RuntimeError(f"no MCAP files found under {episode / 'bag'}")
|
||||
output.mkdir(parents=True, exist_ok=False)
|
||||
|
||||
feedback_columns = [
|
||||
"bag_time_ns",
|
||||
"ros_time_ns",
|
||||
"topic",
|
||||
"side",
|
||||
"joint_index",
|
||||
"motor_id",
|
||||
"position_rad",
|
||||
"speed_rad_s",
|
||||
"current_a",
|
||||
"temperature_c",
|
||||
"mos_temperature_c",
|
||||
"error",
|
||||
]
|
||||
target_columns = [
|
||||
"bag_time_ns",
|
||||
"ros_time_ns",
|
||||
"joint_index",
|
||||
"joint_name",
|
||||
"motor_id",
|
||||
"side",
|
||||
"position_rad",
|
||||
"velocity_rad_s",
|
||||
"effort",
|
||||
]
|
||||
command_columns = [
|
||||
"bag_time_ns",
|
||||
"ros_time_ns",
|
||||
"mode",
|
||||
"label",
|
||||
"motor_id",
|
||||
"side",
|
||||
"joint_index",
|
||||
"kp",
|
||||
"kd",
|
||||
"position_rad",
|
||||
"speed_rad_s",
|
||||
"torque_feedforward",
|
||||
"current_limit_a",
|
||||
]
|
||||
source_columns = [
|
||||
"bag_time_ns",
|
||||
"source_time_ms",
|
||||
"joint_index",
|
||||
"joint_name",
|
||||
"motor_id",
|
||||
"side",
|
||||
"position_rad",
|
||||
]
|
||||
streams: list[Any] = []
|
||||
feedback_stream, feedback_writer = _open_csv(
|
||||
output / "robot_arm_feedback.csv", feedback_columns
|
||||
)
|
||||
target_stream, target_writer = _open_csv(
|
||||
output / "teleop_joint_target.csv", target_columns
|
||||
)
|
||||
command_stream, command_writer = _open_csv(
|
||||
output / "vendor_arm_command.csv", command_columns
|
||||
)
|
||||
source_stream, source_writer = _open_csv(
|
||||
output / "iarm_source_joint.csv", source_columns
|
||||
)
|
||||
streams.extend((feedback_stream, target_stream, command_stream, source_stream))
|
||||
|
||||
counts = {
|
||||
"robot_arm_feedback_rows": 0,
|
||||
"teleop_joint_target_rows": 0,
|
||||
"vendor_arm_command_rows": 0,
|
||||
"iarm_source_joint_rows": 0,
|
||||
}
|
||||
custom_decoders = _custom_decoders()
|
||||
standard_factory = DecoderFactory()
|
||||
try:
|
||||
for bag_file in bag_files:
|
||||
with bag_file.open("rb") as bag_stream:
|
||||
reader = make_reader(bag_stream)
|
||||
for schema, channel, message in reader.iter_messages(topics=TOPICS):
|
||||
topic = channel.topic
|
||||
if topic in CUSTOM_TOPICS:
|
||||
decoded = custom_decoders[CUSTOM_TOPICS[topic]](message.data)
|
||||
else:
|
||||
decoder = standard_factory.decoder_for(
|
||||
channel.message_encoding, schema
|
||||
)
|
||||
if decoder is None:
|
||||
raise RuntimeError(
|
||||
f"no decoder for {topic} ({schema.name if schema else 'no schema'})"
|
||||
)
|
||||
decoded = decoder(message.data)
|
||||
|
||||
if topic in (
|
||||
"/robot_state",
|
||||
"/freq_change/arm_status",
|
||||
"/data_logger/arm_status",
|
||||
):
|
||||
status_message = decoded.arm if topic == "/robot_state" else decoded
|
||||
for motor in status_message.status:
|
||||
motor_id = int(motor.name)
|
||||
side, joint_index = _side_and_index(motor_id)
|
||||
feedback_writer.writerow(
|
||||
{
|
||||
"bag_time_ns": message.log_time,
|
||||
"ros_time_ns": _stamp_ns(decoded.header),
|
||||
"topic": topic,
|
||||
"side": side,
|
||||
"joint_index": joint_index,
|
||||
"motor_id": motor_id,
|
||||
"position_rad": motor.pos,
|
||||
"speed_rad_s": motor.speed,
|
||||
"current_a": motor.current,
|
||||
"temperature_c": motor.temperature,
|
||||
"mos_temperature_c": motor.mos_temperature,
|
||||
"error": motor.error,
|
||||
}
|
||||
)
|
||||
counts["robot_arm_feedback_rows"] += 1
|
||||
elif topic == "/encoder_identical_joint":
|
||||
for index, position in enumerate(decoded.position):
|
||||
motor_id = ARM_MOTOR_IDS[index] if index < 14 else -1
|
||||
side, _ = _side_and_index(motor_id)
|
||||
target_writer.writerow(
|
||||
{
|
||||
"bag_time_ns": message.log_time,
|
||||
"ros_time_ns": _stamp_ns(decoded.header),
|
||||
"joint_index": index,
|
||||
"joint_name": decoded.name[index]
|
||||
if index < len(decoded.name)
|
||||
else "",
|
||||
"motor_id": motor_id,
|
||||
"side": side,
|
||||
"position_rad": position,
|
||||
"velocity_rad_s": decoded.velocity[index]
|
||||
if index < len(decoded.velocity)
|
||||
else "",
|
||||
"effort": decoded.effort[index]
|
||||
if index < len(decoded.effort)
|
||||
else "",
|
||||
}
|
||||
)
|
||||
counts["teleop_joint_target_rows"] += 1
|
||||
elif topic == "/arm/cmd":
|
||||
for motor in decoded.ctrl:
|
||||
motor_id = int(motor.name)
|
||||
side, joint_index = _side_and_index(motor_id)
|
||||
command_writer.writerow(
|
||||
{
|
||||
"bag_time_ns": message.log_time,
|
||||
"ros_time_ns": _stamp_ns(decoded.header),
|
||||
"mode": decoded.mode,
|
||||
"label": decoded.label,
|
||||
"motor_id": motor_id,
|
||||
"side": side,
|
||||
"joint_index": joint_index,
|
||||
"kp": motor.kp,
|
||||
"kd": motor.kd,
|
||||
"position_rad": motor.pos,
|
||||
"speed_rad_s": motor.spd,
|
||||
"torque_feedforward": motor.tor,
|
||||
"current_limit_a": motor.cur,
|
||||
}
|
||||
)
|
||||
counts["vendor_arm_command_rows"] += 1
|
||||
elif topic == "/tg3/data_collection/iarm_frame":
|
||||
payload = json.loads(decoded.data)
|
||||
positions = payload.get("arm", {}).get("position", {})
|
||||
values = [
|
||||
*positions.get("left", []),
|
||||
*positions.get("right", []),
|
||||
]
|
||||
if len(values) != 14:
|
||||
continue
|
||||
source_time = payload.get("timestamp", "")
|
||||
for index, position in enumerate(values):
|
||||
motor_id = ARM_MOTOR_IDS[index]
|
||||
side, _ = _side_and_index(motor_id)
|
||||
source_writer.writerow(
|
||||
{
|
||||
"bag_time_ns": message.log_time,
|
||||
"source_time_ms": source_time,
|
||||
"joint_index": index,
|
||||
"joint_name": JOINT_NAMES[index],
|
||||
"motor_id": motor_id,
|
||||
"side": side,
|
||||
"position_rad": position,
|
||||
}
|
||||
)
|
||||
counts["iarm_source_joint_rows"] += 1
|
||||
finally:
|
||||
for stream in streams:
|
||||
stream.close()
|
||||
|
||||
(output / "summary.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"episode": episode.name,
|
||||
"joint_order": JOINT_NAMES,
|
||||
"motor_ids": ARM_MOTOR_IDS,
|
||||
"units": {
|
||||
"position": "rad",
|
||||
"speed": "rad/s",
|
||||
"current": "A",
|
||||
"temperature": "degC",
|
||||
},
|
||||
"rows": counts,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return counts
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("episode", type=Path, help="Data_Get episode directory")
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
help="new output directory (default: <episode>/decoded_joints)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
episode = args.episode.resolve()
|
||||
output = (args.output or episode / "decoded_joints").resolve()
|
||||
counts = decode_episode(episode, output)
|
||||
print(json.dumps({"output": str(output), "rows": counts}, indent=2))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user