Reference/API

Data

class[source]
graspkit.datasets.graspkit_data.GraspKitData(
scene_ids: list[str],
scene_descriptions: list[SceneDescription] | None = None,
first_indices: Int[Tensor, Bp1] | None = None,
pre_grasp_qpos: Float[Tensor, 'B K dof'] | Float[Tensor, 'N dof'] | None = None,
grasp_qpos: Float[Tensor, 'B K dof'] | Float[Tensor, 'N dof'] | None = None,
post_grasp_qpos: Float[Tensor, 'B K dof'] | Float[Tensor, 'N dof'] | None = None,
full_traj_qpos: Float[Tensor, 'B K T dof'] | Float[Tensor, 'N T dof'] | None = None,
pre_traj_penetration: Float[Tensor, 'B K'] | Float[Tensor, N] | None = None,
pointcloud: Float[Tensor, 'B T Npts 3'] | None = None,
pointcloud_rgb: Float[Tensor, 'B T Npts 3'] | None = None,
pointcloud_segmentation: Int[Tensor, 'B T Npts 1'] | None = None,
target_actor_segmentation_id: Int[Tensor, B] | None = None,
camera_poses: Float[Tensor, 'B T 7'] | None = None
)

Batched per-stage tensor payload; supports padded and packed layouts.

Parameters
scene_idslist[str]required
scene_descriptionslist[SceneDescription] | None= None
first_indicesInt[Tensor, Bp1] | None= None
pre_grasp_qposFloat[Tensor, 'B K dof'] | Float[Tensor, 'N dof'] | None= None
grasp_qposFloat[Tensor, 'B K dof'] | Float[Tensor, 'N dof'] | None= None
post_grasp_qposFloat[Tensor, 'B K dof'] | Float[Tensor, 'N dof'] | None= None
full_traj_qposFloat[Tensor, 'B K T dof'] | Float[Tensor, 'N T dof'] | None= None
pre_traj_penetrationFloat[Tensor, 'B K'] | Float[Tensor, N] | None= None
pointcloudFloat[Tensor, 'B T Npts 3'] | None= None
pointcloud_rgbFloat[Tensor, 'B T Npts 3'] | None= None
pointcloud_segmentationInt[Tensor, 'B T Npts 1'] | None= None
target_actor_segmentation_idInt[Tensor, B] | None= None
camera_posesFloat[Tensor, 'B T 7'] | None= None
Members
attribute[source]
GraspKitData.scene_ids: list[str]
attribute[source]
GraspKitData.scene_descriptions: list[SceneDescription] | None = None
attribute[source]
GraspKitData.first_indices: Int[Tensor, Bp1] | None = None
attribute[source]
GraspKitData.pre_grasp_qpos: Float[Tensor, 'B K dof'] | Float[Tensor, 'N dof'] | None = None
attribute[source]
GraspKitData.grasp_qpos: Float[Tensor, 'B K dof'] | Float[Tensor, 'N dof'] | None = None
attribute[source]
GraspKitData.post_grasp_qpos: Float[Tensor, 'B K dof'] | Float[Tensor, 'N dof'] | None = None
attribute[source]
GraspKitData.full_traj_qpos: Float[Tensor, 'B K T dof'] | Float[Tensor, 'N T dof'] | None = None
attribute[source]
GraspKitData.pre_traj_penetration: Float[Tensor, 'B K'] | Float[Tensor, N] | None = None
attribute[source]
GraspKitData.pointcloud: Float[Tensor, 'B T Npts 3'] | None = None
attribute[source]
GraspKitData.pointcloud_rgb: Float[Tensor, 'B T Npts 3'] | None = None
attribute[source]
GraspKitData.pointcloud_segmentation: Int[Tensor, 'B T Npts 1'] | None = None
attribute[source]
GraspKitData.target_actor_segmentation_id: Int[Tensor, B] | None = None
attribute[source]
GraspKitData.camera_poses: Float[Tensor, 'B T 7'] | None = None
property[source]
GraspKitData.num_scenes: int
property[source]
GraspKitData.is_packed: bool
property[source]
GraspKitData.is_padded: bool
property[source]
GraspKitData.num_grasps_per_scene: Tensor
method[source]
GraspKitData.to_padded() -> GraspKitData
Returns
GraspKitData
method[source]
GraspKitData.to_packed() -> GraspKitData
Returns
GraspKitData
method[source]
GraspKitData.filter(
mask: Tensor
) -> GraspKitData

Filter packed grasps or scenes by a boolean mask and drop empty scenes.

Parameters
maskTensorrequired
Returns
GraspKitData
method[source]
GraspKitData.save(
path: str,
append: bool = False
) -> None

Save scenes to individual H5 files or groups in a merged H5 file.

Parameters
pathstrrequired
appendbool= False
Returns
None
method[source]
GraspKitData.load(
path: str,
scene_ids: list[str] | None = None,
packed: bool = True
) -> GraspKitData

Load scenes in packed form or padded form with uniform grasp counts.

Parameters
pathstrrequired
scene_idslist[str] | None= None
packedbool= True
Returns
GraspKitData
class[source]
graspkit.datasets.replay_data.ReplayData(
qpos: Float[np.ndarray, 'T dof'],
obj_pose: Float[np.ndarray, 'T 7'],
robot_path: str,
mesh_dir: str,
base_type: str,
obj_mesh_path: str,
obj_scale: float,
stage: str = 'sim-valid',
label: str = 'success',
root_pos: Float[np.ndarray, 3] = (lambda: np.zeros(3, np.float32))(),
root_quat: Float[np.ndarray, 4] = (lambda: np.array([1, 0, 0, 0], np.float32))(),
table_size: Float[np.ndarray, 3] = (lambda: np.zeros(3, np.float32))(),
table_pos: Float[np.ndarray, 3] = (lambda: np.zeros(3, np.float32))(),
table_quat: Float[np.ndarray, 4] = (lambda: np.array([1, 0, 0, 0], np.float32))()
)

One recorded trajectory for the replay viewer: robot qpos, object pose and the assets to draw them.

Parameters
qposFloat[np.ndarray, 'T dof']required
obj_poseFloat[np.ndarray, 'T 7']required
robot_pathstrrequired
mesh_dirstrrequired
base_typestrrequired
obj_mesh_pathstrrequired
obj_scalefloatrequired
stagestr= 'sim-valid'
labelstr= 'success'
root_posFloat[np.ndarray, 3]= (lambda: np.zeros(3, np.float32))()
root_quatFloat[np.ndarray, 4]= (lambda: np.array([1, 0, 0, 0], np.float32))()
table_sizeFloat[np.ndarray, 3]= (lambda: np.zeros(3, np.float32))()
table_posFloat[np.ndarray, 3]= (lambda: np.zeros(3, np.float32))()
table_quatFloat[np.ndarray, 4]= (lambda: np.array([1, 0, 0, 0], np.float32))()
Members
attribute[source]
ReplayData.qpos: Float[np.ndarray, 'T dof']

FK variables: floating [xyz, euler_rxyz, joints] or fixed joints.

attribute[source]
ReplayData.obj_pose: Float[np.ndarray, 'T 7']

Object world pose [xyz, quat_wxyz].

attribute[source]
ReplayData.robot_path: str

URDF or MJCF the viewer loads into ViserUrdf / ViserMjcf.

attribute[source]
ReplayData.mesh_dir: str

URDF mesh directory; empty if none.

attribute[source]
ReplayData.base_type: str

"floating" or "fixed".

attribute[source]
ReplayData.obj_mesh_path: str

Object collision mesh (loaded via trimesh).

attribute[source]
ReplayData.obj_scale: float
attribute[source]
ReplayData.stage: str = 'sim-valid'

Viewer partition, e.g. optimize_grasps.py --verbose's "sample (unfiltered)" ... "sim".

attribute[source]
ReplayData.label: str = 'success'

Free filter within a stage, e.g. "success" or "failure".

attribute[source]
ReplayData.root_pos: Float[np.ndarray, 3] = field(default_factory=lambda: np.zeros(3, np.float32))
attribute[source]
ReplayData.root_quat: Float[np.ndarray, 4] = field(default_factory=lambda: np.array([1, 0, 0, 0], np.float32))
attribute[source]
ReplayData.table_size: Float[np.ndarray, 3] = field(default_factory=lambda: np.zeros(3, np.float32))
attribute[source]
ReplayData.table_pos: Float[np.ndarray, 3] = field(default_factory=lambda: np.zeros(3, np.float32))
attribute[source]
ReplayData.table_quat: Float[np.ndarray, 4] = field(default_factory=lambda: np.array([1, 0, 0, 0], np.float32))
method[source]
ReplayData.from_traj(
qpos: Float[np.ndarray, 'T dof'],
obj_pose: Float[np.ndarray, 'T 7'],
robot_config: RobotConfig,
scene_desc: SceneDescription,
stage: str = 'sim-valid',
label: str = 'success'
) -> ReplayData
Parameters
qposFloat[np.ndarray, 'T dof']required
obj_poseFloat[np.ndarray, 'T 7']required
robot_configRobotConfigrequired
scene_descSceneDescriptionrequired
stagestr= 'sim-valid'
labelstr= 'success'
Returns
ReplayData
method[source]
ReplayData.save(
replays: list[ReplayData],
path: str
) -> None
Parameters
replayslist[ReplayData]required
pathstrrequired
Returns
None
method[source]
ReplayData.load(
path: str
) -> list[ReplayData]
Parameters
pathstrrequired
Returns
list[ReplayData]