83 lines
4.4 KiB
Python
83 lines
4.4 KiB
Python
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()
|