tangger/issacdataengine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

init commit

Browse Source
This commit is contained in:
zyhe
2026-03-16 11:44:10 +00:00
commit 94384a93c9
552 changed files with 363038 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

107
workflows/simbox/core/loggers/__init__.py Normal file
View File

@@ -0,0 +1,107 @@
# pylint: skip-file
from abc import ABC, abstractmethod
from typing import Any, Dict, List
class BaseLogger(ABC):
def __init__(
self,
task_dir="Pick_up_the_object",
language_instruction="Pick up the object.",
detailed_language_instruction="Pick up the object with right gripper.",
collect_info="set1-1_collector1_20250715",
version="v1.0",
tpi_initial_info={},
):
self.log_num_steps = 0
self.task_dir = task_dir
self.language_instruction = language_instruction
self.detailed_language_instruction = detailed_language_instruction
self.version = version
self.collect_info = collect_info
self.tpi_initial_info = tpi_initial_info
self.json_data_logger: Dict[str, List[Any]] = {}
self.scalar_data_logger: Dict[str, List[Any]] = {}
self.action_data_logger: Dict[str, List[Any]] = {}
self.proprio_data_logger: Dict[str, List[Any]] = {}
self.object_data_logger: Dict[str, List[Any]] = {}
self.color_image_logger: Dict[str, List[Any]] = {}
self.depth_image_logger: Dict[str, List[Any]] = {}
def update_tpi_initial_info(self, tpi_initial_info):
self.tpi_initial_info = tpi_initial_info
def count_timestep(self):
self.log_num_steps += 1
def add_json_data(self, robot, key, data):
if robot not in self.json_data_logger:
self.json_data_logger[robot] = {}
self.json_data_logger[robot][key] = data
def add_proprio_data(self, robot, key, value):
if robot not in self.proprio_data_logger:
self.proprio_data_logger[robot] = {}
if key not in self.proprio_data_logger[robot]:
self.proprio_data_logger[robot][key] = []
self.proprio_data_logger[robot][key].append(value)
def add_action_data(self, robot, key, value):
if robot not in self.action_data_logger:
self.action_data_logger[robot] = {}
if key not in self.action_data_logger[robot]:
self.action_data_logger[robot][key] = []
self.action_data_logger[robot][key].append(value)
def add_object_data(self, robot, key, value):
if robot not in self.object_data_logger:
self.object_data_logger[robot] = {}
if key not in self.object_data_logger[robot]:
self.object_data_logger[robot][key] = []
self.object_data_logger[robot][key].append(value)
def add_scalar_data(self, robot, key, value):
if robot not in self.scalar_data_logger:
self.scalar_data_logger[robot] = {}
if key not in self.scalar_data_logger[robot]:
self.scalar_data_logger[robot][key] = []
self.scalar_data_logger[robot][key].append(value)
def add_color_image(self, robot, key, value):
if robot not in self.color_image_logger:
self.color_image_logger[robot] = {}
if key not in self.color_image_logger[robot]:
self.color_image_logger[robot][key] = []
self.color_image_logger[robot][key].append(value)
# def add_depth_image(self, key, value):
# if key not in self.depth_image_logger:
# self.depth_image_logger[key] = []
# self.depth_image_logger[key].append(value)
def clear(
self,
language_instruction="Pick up the object.",
detailed_language_instruction="Pick up the object with right gripper.",
):
self.language_instruction = language_instruction
self.detailed_language_instruction = detailed_language_instruction
self.last_qpos = None
self.last_ee_pose = None
self.log_num_steps = 0
self.tpi_initial_info = {}
self.json_data_logger = {}
self.proprio_data_logger = {}
self.action_data_logger = {}
self.object_data_logger = {}
self.scalar_data_logger = {}
self.color_image_logger = {}
self.depth_image_logger = {}
@abstractmethod
def close(self):
pass
@abstractmethod
def save(self):
pass

194
workflows/simbox/core/loggers/lmdb_logger.py Normal file
View File

@@ -0,0 +1,194 @@
# pylint: skip-file
import json
import os
import pickle
from pathlib import Path
import cv2
import imageio
import lmdb
import numpy as np
from core.loggers import BaseLogger
from tqdm import tqdm
DEFAULT_RGB_SCALE_FACTOR = 256000.0
# pylint: disable=line-too-long,unused-argument
class LmdbLogger(BaseLogger):
def __init__(
self,
task_dir="Pick_up_the_object",
language_instruction="Pick up the object.",
detailed_language_instruction="Pick up the object with right gripper.",
collect_info: str = "set1-1_collector1_20250715",
version: str = "v1.0",
tpi_initial_info: dict = {},
max_size: int = 1,
image_quality: int = 100,
save_depth: bool = True,
min_inttype: int = 0,
max_inttype: int = 2**24 - 1,
):
super().__init__(
task_dir=task_dir,
language_instruction=language_instruction,
detailed_language_instruction=detailed_language_instruction,
collect_info=collect_info,
version=version,
tpi_initial_info=tpi_initial_info,
)
self.max_size = int(max_size * 1024**4)
self.image_quality = image_quality
self.min_inttype = min_inttype
self.max_inttype = max_inttype
def close(self):
pass
def save(self, this_save_dir, timestamp: str, save_img: bool = True):
for robot_idx, robot_name in enumerate(self.proprio_data_logger.keys()):
save_dir = Path(this_save_dir)
save_dir = save_dir / f"{robot_name}" / f"{self.task_dir}" / f"{self.collect_info}" / f"{timestamp}"
save_dir.mkdir(parents=True, exist_ok=True)
# Meta info
meta_info = {}
meta_info["keys"] = {}
meta_info["max_size"] = self.max_size
meta_info["language_instruction"] = self.language_instruction[robot_idx]
meta_info["detailed_language_instruction"] = self.detailed_language_instruction[robot_idx]
print("language_instruction :", meta_info["language_instruction"])
print("detailed_language_instruction :", meta_info["detailed_language_instruction"])
meta_info["tpi_initial_info"] = self.tpi_initial_info
meta_info["collect_info"] = self.collect_info
meta_info["version"] = self.version
# Lmdb
log_path_lmdb = save_dir / "lmdb"
lmdb_env = lmdb.open(str(log_path_lmdb), map_size=self.max_size)
txn = lmdb_env.begin(write=True)
# Save json data
with open(log_path_lmdb / "info.json", "w") as f:
json.dump(self.json_data_logger[robot_name], f)
txn.put("json_data".encode("utf-8"), pickle.dumps(self.json_data_logger[robot_name]))
meta_info["keys"]["json_data"] = ["json_data".encode("utf-8")]
# Save scalar data
meta_info["keys"]["scalar_data"] = []
for key, value in self.scalar_data_logger[robot_name].items():
txn.put(key.encode("utf-8"), pickle.dumps(value))
meta_info["keys"]["scalar_data"].append(key.encode("utf-8"))
# Save proprios
meta_info["keys"]["proprio_data"] = []
for key, value in self.proprio_data_logger[robot_name].items():
txn.put(key.encode("utf-8"), pickle.dumps(value))
meta_info["keys"]["proprio_data"].append(key.encode("utf-8"))
# Save objects
meta_info["keys"]["object_data"] = []
if robot_name in self.object_data_logger:
for key, value in self.object_data_logger[robot_name].items():
if "robotiq" in robot_name and key == "states.gripper.position":
value = [self.action_data_logger[robot_name]["master_actions.gripper.position"][0]] + (
self.action_data_logger[robot_name]["master_actions.gripper.position"]
)[:-1]
txn.put(key.encode("utf-8"), pickle.dumps(value))
meta_info["keys"]["object_data"].append(key.encode("utf-8"))
# Save master actions
meta_info["keys"]["action_data"] = []
for key, value in self.action_data_logger[robot_name].items():
# Update gripper action
if "gripper.position" in key:
value.pop(0)
value.append(value[-1]) # Use next gripper width as gripper action label
txn.put(key.encode("utf-8"), pickle.dumps(value))
meta_info["keys"]["action_data"].append(key.encode("utf-8"))
# Save actions
# Here we use next robot state as action
for key, value in self.proprio_data_logger[robot_name].items():
if "states." in key:
new_key = key.replace("states.", "actions.")
value.pop(0)
value.append(value[-1]) # Use next robot state as action
txn.put(new_key.encode("utf-8"), pickle.dumps(value))
meta_info["keys"]["action_data"].append(new_key.encode("utf-8"))
if (
"split_aloha" in robot_name
or "lift2" in robot_name
or "azure_loong" in robot_name
or "genie" in robot_name
):
left_gripper_openness = self.action_data_logger[robot_name]["master_actions.left_gripper.openness"]
right_gripper_openness = self.action_data_logger[robot_name]["master_actions.right_gripper.openness"]
txn.put("actions.left_gripper.openness".encode("utf-8"), pickle.dumps(left_gripper_openness))
meta_info["keys"]["action_data"].append("actions.left_gripper.openness".encode("utf-8"))
txn.put("actions.right_gripper.openness".encode("utf-8"), pickle.dumps(right_gripper_openness))
meta_info["keys"]["action_data"].append("actions.right_gripper.openness".encode("utf-8"))
elif "franka" in robot_name:
gripper_openness = self.action_data_logger[robot_name]["master_actions.gripper.openness"]
txn.put("actions.gripper.openness".encode("utf-8"), pickle.dumps(gripper_openness))
meta_info["keys"]["action_data"].append("actions.gripper.openness".encode("utf-8"))
# Save color images
if save_img:
for key, value in self.color_image_logger[robot_name].items():
root_img_path = save_dir / f"{key}"
root_img_path.mkdir(parents=True, exist_ok=True)
meta_info["keys"][key] = []
for i, image in enumerate(tqdm(value)):
step_id = str(i).zfill(4)
txn.put(
f"{key}/{step_id}".encode("utf-8"),
pickle.dumps(cv2.imencode(".jpg", image.astype(np.uint8))[1]),
)
meta_info["keys"][key].append(f"{key}/{step_id}".encode("utf-8"))
imageio.mimsave(os.path.join(root_img_path, "demo.mp4"), value, fps=15)
meta_info["num_steps"] = len(value)
txn.commit()
lmdb_env.close()
pickle.dump(meta_info, open(os.path.join(save_dir, "meta_info.pkl"), "wb"))
def dump(self):
logger_info = {}
logger_info["proprio_data_logger"] = self.proprio_data_logger
logger_info["max_size"] = self.max_size
logger_info["language_instruction"] = self.language_instruction
logger_info["detailed_language_instruction"] = self.detailed_language_instruction
logger_info["tpi_initial_info"] = self.tpi_initial_info
logger_info["collect_info"] = self.collect_info
logger_info["version"] = self.version
logger_info["json_data_logger"] = self.json_data_logger
logger_info["scalar_data_logger"] = self.scalar_data_logger
logger_info["object_data_logger"] = self.object_data_logger
logger_info["action_data_logger"] = self.action_data_logger
logger_info["log_num_steps"] = self.log_num_steps
return pickle.dumps(logger_info)
def dedump(self, ser):
logger_info = pickle.loads(ser)
self.proprio_data_logger = logger_info["proprio_data_logger"]
self.max_size = logger_info["max_size"]
self.language_instruction = logger_info["language_instruction"]
self.detailed_language_instruction = logger_info["detailed_language_instruction"]
self.tpi_initial_info = logger_info["tpi_initial_info"]
self.collect_info = logger_info["collect_info"]
self.version = logger_info["version"]
self.json_data_logger = logger_info["json_data_logger"]
self.scalar_data_logger = logger_info["scalar_data_logger"]
self.object_data_logger = logger_info["object_data_logger"]
self.action_data_logger = logger_info["action_data_logger"]
self.log_num_steps = logger_info["log_num_steps"]
return True

82
workflows/simbox/core/loggers/utils.py Normal file
View File

@@ -0,0 +1,82 @@
from core.utils.transformation_utils import get_fk_solution, pose_to_6d
from .lmdb_logger import LmdbLogger
# pylint: disable=line-too-long,unused-argument
def log_dual_obs(logger: LmdbLogger, obs, action_dict, controllers, step_idx=0):
# Add robots' proprio
for robot_name, robot_infos in obs["robots"].items():
for key in robot_infos.keys():
logger.add_proprio_data(robot_name, key, robot_infos[key])
# Add objects' data (if exists)
if "objects" in obs:
for object_name in obs["objects"].keys():
for attr_name, attr_value in obs["objects"][object_name].items():
logger.add_object_data(robot_name, f"{object_name}/{attr_name}", attr_value)
# Add robots' action data (very very important)
if "split_aloha" in robot_name or "lift2" in robot_name or "azure_loong" in robot_name or "genie" in robot_name:
left_joint_position = obs["robots"][robot_name]["states.left_joint.position"]
right_joint_position = obs["robots"][robot_name]["states.right_joint.position"]
left_gripper_position = obs["robots"][robot_name]["states.left_gripper.position"]
right_gripper_position = obs["robots"][robot_name]["states.right_gripper.position"]
left_gripper_openness = (
1.0 if controllers[robot_name]["left"]._gripper_state > 0.0 else 0.0
) # 1.0 open, 0.0 close
right_gripper_openness = (
1.0 if controllers[robot_name]["right"]._gripper_state > 0.0 else 0.0
) # 1.0 open, 0.0 close
# Use raw action to udpate if one arm is not static
robot_action = action_dict.get(robot_name, None)
if robot_action is not None:
raw_action = robot_action["raw_action"]
for action in raw_action:
lr_name = action["lr_name"]
if lr_name == "left":
arm_action = action["arm_action"]
left_joint_position = arm_action
elif lr_name == "right":
arm_action = action["arm_action"]
right_joint_position = arm_action
else:
pass
logger.add_action_data(robot_name, "master_actions.left_joint.position", left_joint_position)
logger.add_action_data(robot_name, "master_actions.right_joint.position", right_joint_position)
logger.add_action_data(robot_name, "master_actions.left_gripper.position", left_gripper_position)
logger.add_action_data(robot_name, "master_actions.right_gripper.position", right_gripper_position)
logger.add_action_data(robot_name, "master_actions.left_gripper.openness", left_gripper_openness)
logger.add_action_data(robot_name, "master_actions.right_gripper.openness", right_gripper_openness)
elif "franka" in robot_name:
joint_position = obs["robots"][robot_name]["states.joint.position"]
gripper_pose = obs["robots"][robot_name]["states.gripper.pose"]
gripper_openness = (
1.0 if controllers[robot_name]["left"]._gripper_state > 0.0 else 0.0
) # 1.0 open, 0.0 close
gripper_position = obs["robots"][robot_name]["states.gripper.position"]
# Use raw action to udpate if one arm is not static
robot_action = action_dict.get(robot_name, None)
if robot_action is not None:
raw_action = robot_action["raw_action"]
for action in raw_action:
lr_name = action["lr_name"]
if lr_name == "left":
arm_action = action["arm_action"]
joint_position = arm_action
gripper_pose = pose_to_6d(get_fk_solution(joint_position[:7]))
else:
pass
logger.add_action_data(robot_name, "master_actions.joint.position", joint_position)
logger.add_action_data(robot_name, "master_actions.gripper.position", gripper_position)
logger.add_action_data(robot_name, "master_actions.gripper.openness", gripper_openness)
logger.add_action_data(robot_name, "master_actions.gripper.pose", gripper_pose)
else:
raise NotImplementedError
# Count time steps
logger.count_timestep()