add natural language model and app

app.py

@@ -10,13 +10,15 @@ from omegaconf import OmegaConf
 import sys
 sys.path.append('./src')
 
-from StructDiffusion.data.semantic_arrangement_demo import SemanticArrangementDataset
+from StructDiffusion.data.semantic_arrangement_language_demo import SemanticArrangementDataset
 from StructDiffusion.language.tokenizer import Tokenizer
-from StructDiffusion.models.pl_models import ConditionalPoseDiffusionModel
-from StructDiffusion.diffusion.sampler import Sampler
+from StructDiffusion.models.pl_models import ConditionalPoseDiffusionModel, PairwiseCollisionModel
+from StructDiffusion.diffusion.sampler import Sampler, SamplerV2
 from StructDiffusion.diffusion.pose_conversion import get_struct_objs_poses
 from StructDiffusion.utils.files import get_checkpoint_path_from_dir
-from StructDiffusion.utils.rearrangement import show_pcs_with_trimesh
+from StructDiffusion.utils.rearrangement import show_pcs_with_trimesh, get_trimesh_scene_with_table
+import StructDiffusion.utils.transformations as tra
+from StructDiffusion.language.sentence_encoder import SentenceBertEncoder
 import StructDiffusion.utils.transformations as tra
 
 
@@ -65,23 +67,31 @@ class Infer_Wrapper:
 
     def __init__(self, args, cfg):
 
+        self.num_pts = cfg.DATASET.num_pts
+
         # load
         pl.seed_everything(args.eval_random_seed)
         self.device = (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
 
-        checkpoint_dir = os.path.join(cfg.WANDB.save_dir, cfg.WANDB.project, args.checkpoint_id, "checkpoints")
-        checkpoint_path = get_checkpoint_path_from_dir(checkpoint_dir)
+        diffusion_checkpoint_path = get_checkpoint_path_from_dir(os.path.join(cfg.WANDB.save_dir, cfg.WANDB.project, args.diffusion_checkpoint_id, "checkpoints"))
+        collision_checkpoint_path = get_checkpoint_path_from_dir(os.path.join(cfg.WANDB.save_dir, cfg.WANDB.project, args.collision_checkpoint_id, "checkpoints"))
 
         self.tokenizer = Tokenizer(cfg.DATASET.vocab_dir)
         # override ignore_rgb for visualization
         cfg.DATASET.ignore_rgb = False
         self.dataset = SemanticArrangementDataset(tokenizer=self.tokenizer, **cfg.DATASET)
 
-        self.sampler = Sampler(ConditionalPoseDiffusionModel, checkpoint_path, self.device)
+        self.sampler = SamplerV2(ConditionalPoseDiffusionModel, diffusion_checkpoint_path,
+                                 PairwiseCollisionModel, collision_checkpoint_path, self.device)
+
+        self.sentence_encoder = SentenceBertEncoder()
+
+        self.session_id_to_obj_xyzs = {}
 
     def visualize_scene(self, di, session_id):
-        raw_datum = self.dataset.get_raw_data(di)
-        language_command = self.tokenizer.convert_structure_params_to_natural_language(raw_datum["sentence"])
+
+        raw_datum = self.dataset.get_raw_data(di, inference_mode=True, shuffle_object_index=True)
+        language_command = raw_datum["template_sentence"]
 
         obj_xyz = raw_datum["pcs"]
         scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in obj_xyz], [xyz[:, 3:] for xyz in obj_xyz], return_scene=True)
@@ -93,20 +103,77 @@ class Infer_Wrapper:
 
         return language_command, scene_filename
 
-    def infer(self, di, session_id, progress=gr.Progress()):
+    def build_scene(self, mesh_filename_1, x_1, y_1, z_1, ai_1, aj_1, ak_1, scale_1,
+                                                mesh_filename_2, x_2, y_2, z_2, ai_2, aj_2, ak_2, scale_2,
+                                                mesh_filename_3, x_3, y_3, z_3, ai_3, aj_3, ak_3, scale_3,
+                                                mesh_filename_4, x_4, y_4, z_4, ai_4, aj_4, ak_4, scale_4,
+                                                mesh_filename_5, x_5, y_5, z_5, ai_5, aj_5, ak_5, scale_5, session_id):
+
+        object_list = [(mesh_filename_1, x_1, y_1, z_1, ai_1, aj_1, ak_1, scale_1),
+                       (mesh_filename_2, x_2, y_2, z_2, ai_2, aj_2, ak_2, scale_2),
+                       (mesh_filename_3, x_3, y_3, z_3, ai_3, aj_3, ak_3, scale_3),
+                       (mesh_filename_4, x_4, y_4, z_4, ai_4, aj_4, ak_4, scale_4),
+                       (mesh_filename_5, x_5, y_5, z_5, ai_5, aj_5, ak_5, scale_5)]
+
+        scene = get_trimesh_scene_with_table()
+
+        obj_xyzs = []
+        for mesh_filename, x, y, z, ai, aj, ak, scale in object_list:
+            if mesh_filename is None:
+                continue
+            obj_mesh = trimesh.load(mesh_filename)
+            obj_mesh.apply_scale(scale)
+            z_min = obj_mesh.bounds[0, 2]
+            tform = tra.euler_matrix(ai, aj, ak)
+            tform[:3, 3] = [x, y, z - z_min]
+            obj_mesh.apply_transform(tform)
+            obj_xyz = obj_mesh.sample(self.num_pts)
+            obj = trimesh.PointCloud(obj_xyz)
+            scene.add_geometry(obj)
+
+            obj_xyzs.append(obj_xyz)
+
+        self.session_id_to_obj_xyzs[session_id] = obj_xyzs
+
+        # scene.show()
+
+        # obj_file = "/home/weiyu/data_drive/StructDiffusion/housekeep_custom_handpicked_small/visual/book_Eat_to_Live_The_Amazing_NutrientRich_Program_for_Fast_and_Sustained_Weight_Loss_Revised_Edition_Book_L/model.obj"
+        # obj = trimesh.load(obj_file)
+        #
+        # scene = get_trimesh_scene_with_table()
+        # scene.add_geometry(obj)
+        #
+        # scene.show()
+
+        # raw_datum = self.dataset.get_raw_data(di, inference_mode=True, shuffle_object_index=True)
+        # language_command = raw_datum["template_sentence"]
+        #
+        # obj_xyz = raw_datum["pcs"]
+        # scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in obj_xyz], [xyz[:, 3:] for xyz in obj_xyz],
+        #                               return_scene=True)
+
+        scene.apply_transform(tra.euler_matrix(np.pi, 0, np.pi / 2))
+        scene_filename = "./tmp_data/input_scene_{}.glb".format(session_id)
+        scene.export(scene_filename)
+
+        return scene_filename
+
+        # return language_command, scene_filename
+
+    def infer(self, language_command, session_id, progress=gr.Progress()):
+
+        obj_xyzs = self.session_id_to_obj_xyzs[session_id]
 
-        # di = np.random.choice(len(self.dataset))
+        sentence_embedding = self.sentence_encoder.encode([language_command]).flatten()
 
-        raw_datum = self.dataset.get_raw_data(di)
-        print(self.tokenizer.convert_structure_params_to_natural_language(raw_datum["sentence"]))
-        datum = self.dataset.convert_to_tensors(raw_datum, self.tokenizer)
+        raw_datum = self.dataset.build_data_from_xyzs(obj_xyzs, sentence_embedding)
+        datum = self.dataset.convert_to_tensors(raw_datum, self.tokenizer, use_sentence_embedding=True)
         batch = self.dataset.single_datum_to_batch(datum, args.num_samples, self.device, inference_mode=True)
 
-        num_poses = datum["goal_poses"].shape[0]
-        xs = self.sampler.sample(batch, num_poses, progress)
+        num_poses = raw_datum["num_goal_poses"]
+        struct_pose, pc_poses_in_struct = self.sampler.sample(batch, num_poses, args.num_elites, args.discriminator_batch_size)
 
-        struct_pose, pc_poses_in_struct = get_struct_objs_poses(xs[0])
-        new_obj_xyzs = move_pc_and_create_scene_simple(batch["pcs"], struct_pose, pc_poses_in_struct)
+        new_obj_xyzs = move_pc_and_create_scene_simple(batch["pcs"][:args.num_elites], struct_pose, pc_poses_in_struct)
 
         # vis
         vis_obj_xyzs = new_obj_xyzs[:3]
@@ -115,18 +182,11 @@ class Infer_Wrapper:
                 vis_obj_xyzs = vis_obj_xyzs.detach().cpu()
             vis_obj_xyzs = vis_obj_xyzs.numpy()
 
-        # for bi, vis_obj_xyz in enumerate(vis_obj_xyzs):
-        #     if verbose:
-        #         print("example {}".format(bi))
-        #         print(vis_obj_xyz.shape)
-        #
-        #     if trimesh:
-        #         show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], [xyz[:, 3:] for xyz in vis_obj_xyz])
         vis_obj_xyz = vis_obj_xyzs[0]
-        scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], [xyz[:, 3:] for xyz in vis_obj_xyz], return_scene=True)
-
+        # scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], [xyz[:, 3:] for xyz in vis_obj_xyz], return_scene=True)
+        scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], obj_rgbs=None, return_scene=True)
+        scene.show()
         scene.apply_transform(tra.euler_matrix(np.pi, 0, np.pi/2))
-
         scene_filename = "./tmp_data/output_scene_{}.glb".format(session_id)
         scene.export(scene_filename)
 
@@ -167,10 +227,13 @@ class Infer_Wrapper:
 
 args = OmegaConf.create()
 args.base_config_file = "./configs/base.yaml"
-args.config_file = "./configs/conditional_pose_diffusion.yaml"
-args.checkpoint_id = "ConditionalPoseDiffusion"
+args.config_file = "./configs/conditional_pose_diffusion_language.yaml"
+args.diffusion_checkpoint_id = "ConditionalPoseDiffusionLanguage"
+args.collision_checkpoint_id = "CollisionDiscriminator"
 args.eval_random_seed = 42
-args.num_samples = 1
+args.num_samples = 50
+args.num_elites = 3
+args.discriminator_batch_size = 10
 
 base_cfg = OmegaConf.load(args.base_config_file)
 cfg = OmegaConf.load(args.config_file)
@@ -178,34 +241,123 @@ cfg = OmegaConf.merge(base_cfg, cfg)
 
 infer_wrapper = Infer_Wrapper(args, cfg)
 
-# version 0
-# demo = gr.Interface(
-#     fn=infer_wrapper.run,
-#     inputs=gr.Slider(0, len(infer_wrapper.dataset)),
-#     # clear color range [0-1.0]
-#     outputs=gr.Model3D(clear_color=[0, 0, 0, 0],  label="3D Model")
-# )
+# # version 1
+# demo = gr.Blocks(theme=gr.themes.Soft())
+# with demo:
+#     gr.Markdown("<p style='text-align:center;font-size:18px'><b>StructDiffusion Demo</b></p>")
+#     # font-size:18px
+#     gr.Markdown("<p style='text-align:center'>StructDiffusion combines a diffusion model and an object-centric transformer to construct structures given partial-view point clouds and high-level language goals.<br><a href='https://structdiffusion.github.io/'>Website</a> | <a href='https://github.com/StructDiffusion/StructDiffusion'>Code</a></p>")
 #
+#     session_id = gr.State(value=np.random.randint(0, 1000))
+#     data_selection = gr.Number(label="Example No.", minimum=0, maximum=len(infer_wrapper.dataset) - 1, precision=0)
+#     input_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Input 3D Scene")
+#     language_command = gr.Textbox(label="Input Language Command")
+#     output_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Generated 3D Structure")
+#
+#     b1 = gr.Button("Show Input Language and Scene")
+#     b2 = gr.Button("Generate 3D Structure")
+#
+#     b1.click(infer_wrapper.visualize_scene, inputs=[data_selection, session_id], outputs=[language_command, input_scene])
+#     b2.click(infer_wrapper.infer, inputs=[data_selection, session_id], outputs=output_scene)
+#
+# demo.queue(concurrency_count=10)
 # demo.launch()
 
 # version 1
-demo = gr.Blocks(theme=gr.themes.Soft())
+# demo = gr.Blocks(theme=gr.themes.Soft())
+demo = gr.Blocks()
 with demo:
     gr.Markdown("<p style='text-align:center;font-size:18px'><b>StructDiffusion Demo</b></p>")
     # font-size:18px
     gr.Markdown("<p style='text-align:center'>StructDiffusion combines a diffusion model and an object-centric transformer to construct structures given partial-view point clouds and high-level language goals.<br><a href='https://structdiffusion.github.io/'>Website</a> | <a href='https://github.com/StructDiffusion/StructDiffusion'>Code</a></p>")
 
     session_id = gr.State(value=np.random.randint(0, 1000))
-    data_selection = gr.Number(label="Example No.", minimum=0, maximum=len(infer_wrapper.dataset) - 1, precision=0)
-    input_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Input 3D Scene")
+    with gr.Tab("Object 1"):
+        with gr.Column(scale=1, min_width=600):
+            mesh_filename_1 = gr.Model3D(clear_color=[0, 0, 0, 0], label="Load 3D Object")
+            with gr.Row():
+                x_1 = gr.Slider(0, 1, label="x")
+                y_1 = gr.Slider(-0.5, 0.5, label="y")
+                z_1 = gr.Slider(0, 0.5, label="z")
+            with gr.Row():
+                ai_1 = gr.Slider(0, np.pi * 2, label="roll")
+                aj_1 = gr.Slider(0, np.pi * 2, label="pitch")
+                ak_1 = gr.Slider(0, np.pi * 2, label="yaw")
+            scale_1 = gr.Slider(0, 1)
+    with gr.Tab("Object 2"):
+        with gr.Column(scale=1, min_width=600):
+            mesh_filename_2 = gr.Model3D(clear_color=[0, 0, 0, 0], label="Load 3D Object")
+            with gr.Row():
+                x_2 = gr.Slider(0, 1, label="x")
+                y_2 = gr.Slider(-0.5, 0.5, label="y")
+                z_2 = gr.Slider(0, 0.5, label="z")
+            with gr.Row():
+                ai_2 = gr.Slider(0, np.pi * 2, label="roll")
+                aj_2 = gr.Slider(0, np.pi * 2, label="pitch")
+                ak_2 = gr.Slider(0, np.pi * 2, label="yaw")
+            scale_2 = gr.Slider(0, 1)
+    with gr.Tab("Object 3"):
+        with gr.Column(scale=1, min_width=600):
+            mesh_filename_3 = gr.Model3D(clear_color=[0, 0, 0, 0], label="Load 3D Object")
+            with gr.Row():
+                x_3 = gr.Slider(0, 1, label="x")
+                y_3 = gr.Slider(-0.5, 0.5, label="y")
+                z_3 = gr.Slider(0, 0.5, label="z")
+            with gr.Row():
+                ai_3 = gr.Slider(0, np.pi * 2, label="roll")
+                aj_3 = gr.Slider(0, np.pi * 2, label="pitch")
+                ak_3 = gr.Slider(0, np.pi * 2, label="yaw")
+            scale_3 = gr.Slider(0, 1)
+    with gr.Tab("Object 4"):
+        with gr.Column(scale=1, min_width=600):
+            mesh_filename_4 = gr.Model3D(clear_color=[0, 0, 0, 0], label="Load 3D Object")
+            with gr.Row():
+                x_4 = gr.Slider(0, 1, label="x")
+                y_4 = gr.Slider(-0.5, 0.5, label="y")
+                z_4 = gr.Slider(0, 0.5, label="z")
+            with gr.Row():
+                ai_4 = gr.Slider(0, np.pi * 2, label="roll")
+                aj_4 = gr.Slider(0, np.pi * 2, label="pitch")
+                ak_4 = gr.Slider(0, np.pi * 2, label="yaw")
+            scale_4 = gr.Slider(0, 1)
+    with gr.Tab("Object 5"):
+        with gr.Column(scale=1, min_width=600):
+            mesh_filename_5 = gr.Model3D(clear_color=[0, 0, 0, 0], label="Load 3D Object")
+            with gr.Row():
+                x_5 = gr.Slider(0, 1, label="x")
+                y_5 = gr.Slider(-0.5, 0.5, label="y")
+                z_5 = gr.Slider(0, 0.5, label="z")
+            with gr.Row():
+                ai_5 = gr.Slider(0, np.pi * 2, label="roll")
+                aj_5 = gr.Slider(0, np.pi * 2, label="pitch")
+                ak_5 = gr.Slider(0, np.pi * 2, label="yaw")
+            scale_5 = gr.Slider(0, 1)
+
+    b1 = gr.Button("Build Initial Scene")
+
+    initial_scene = gr.Model3D(clear_color=[0, 0, 0, 0], label="Initial 3D Scene")
     language_command = gr.Textbox(label="Input Language Command")
-    output_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Generated 3D Structure")
 
-    b1 = gr.Button("Show Input Language and Scene")
     b2 = gr.Button("Generate 3D Structure")
 
-    b1.click(infer_wrapper.visualize_scene, inputs=[data_selection, session_id], outputs=[language_command, input_scene])
-    b2.click(infer_wrapper.infer, inputs=[data_selection, session_id], outputs=output_scene)
+    output_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Generated 3D Structure")
+
+    # data_selection = gr.Number(label="Example No.", minimum=0, maximum=len(infer_wrapper.dataset) - 1, precision=0)
+    # input_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Input 3D Scene")
+    # language_command = gr.Textbox(label="Input Language Command")
+    # output_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Generated 3D Structure")
+    #
+    # b1 = gr.Button("Show Input Language and Scene")
+    # b2 = gr.Button("Generate 3D Structure")
+
+    b1.click(infer_wrapper.build_scene, inputs=[mesh_filename_1, x_1, y_1, z_1, ai_1, aj_1, ak_1, scale_1,
+                                                mesh_filename_2, x_2, y_2, z_2, ai_2, aj_2, ak_2, scale_2,
+                                                mesh_filename_3, x_3, y_3, z_3, ai_3, aj_3, ak_3, scale_3,
+                                                mesh_filename_4, x_4, y_4, z_4, ai_4, aj_4, ak_4, scale_4,
+                                                mesh_filename_5, x_5, y_5, z_5, ai_5, aj_5, ak_5, scale_5,
+                                                session_id], outputs=[initial_scene])
+
+    b2.click(infer_wrapper.infer, inputs=[language_command, session_id], outputs=output_scene)
 
 demo.queue(concurrency_count=10)
 demo.launch()

app_v0.py

@@ -0,0 +1,282 @@
+import os
+import argparse
+import torch
+import trimesh
+import numpy as np
+import pytorch_lightning as pl
+import gradio as gr
+from omegaconf import OmegaConf
+
+import sys
+sys.path.append('./src')
+
+from StructDiffusion.data.semantic_arrangement_demo import SemanticArrangementDataset
+from StructDiffusion.language.tokenizer import Tokenizer
+from StructDiffusion.models.pl_models import ConditionalPoseDiffusionModel
+from StructDiffusion.diffusion.sampler import Sampler
+from StructDiffusion.diffusion.pose_conversion import get_struct_objs_poses
+from StructDiffusion.utils.files import get_checkpoint_path_from_dir
+from StructDiffusion.utils.rearrangement import show_pcs_with_trimesh
+import StructDiffusion.utils.transformations as tra
+
+
+def move_pc_and_create_scene_simple(obj_xyzs, struct_pose, pc_poses_in_struct):
+
+    device = obj_xyzs.device
+
+    # obj_xyzs: B, N, P, 3 or 6
+    # struct_pose: B, 1, 4, 4
+    # pc_poses_in_struct: B, N, 4, 4
+
+    B, N, _, _ = pc_poses_in_struct.shape
+    _, _, P, _ = obj_xyzs.shape
+
+    current_pc_poses = torch.eye(4).repeat(B, N, 1, 1).to(device)  # B, N, 4, 4
+    # print(torch.mean(obj_xyzs, dim=2).shape)
+    current_pc_poses[:, :, :3, 3] = torch.mean(obj_xyzs[:, :, :, :3], dim=2)  # B, N, 4, 4
+    current_pc_poses = current_pc_poses.reshape(B * N, 4, 4)  # B x N, 4, 4
+
+    struct_pose = struct_pose.repeat(1, N, 1, 1) # B, N, 4, 4
+    struct_pose = struct_pose.reshape(B * N, 4, 4)  # B x 1, 4, 4
+    pc_poses_in_struct = pc_poses_in_struct.reshape(B * N, 4, 4)  # B x N, 4, 4
+
+    goal_pc_pose = struct_pose @ pc_poses_in_struct  # B x N, 4, 4
+    # print("goal pc poses")
+    # print(goal_pc_pose)
+    goal_pc_transform = goal_pc_pose @ torch.inverse(current_pc_poses)  # B x N, 4, 4
+
+    # # important: pytorch3d uses row-major ordering, need to transpose each transformation matrix
+    # transpose = tra3d.Transform3d(matrix=goal_pc_transform.transpose(1, 2))
+    # new_obj_xyzs = obj_xyzs.reshape(B * N, P, -1)  # B x N, P, 3
+    # new_obj_xyzs[:, :, :3] = transpose.transform_points(new_obj_xyzs[:, :, :3])
+
+    # a verision that does not rely on pytorch3d
+    new_obj_xyzs = obj_xyzs.reshape(B * N, P, -1)[:, :, :3]  # B x N, P, 3
+    new_obj_xyzs = torch.concat([new_obj_xyzs, torch.ones(B * N, P, 1).to(device)], dim=-1) # B x N, P, 4
+    new_obj_xyzs = torch.einsum('bij,bkj->bki', goal_pc_transform, new_obj_xyzs)[:, :, :3]  # # B x N, P, 3
+
+    # put it back to B, N, P, 3
+    obj_xyzs[:, :, :, :3] = new_obj_xyzs.reshape(B, N, P, -1)
+
+    return obj_xyzs
+
+
+class Infer_Wrapper:
+
+    def __init__(self, args, cfg):
+
+        # load
+        pl.seed_everything(args.eval_random_seed)
+        self.device = (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
+
+        checkpoint_dir = os.path.join(cfg.WANDB.save_dir, cfg.WANDB.project, args.checkpoint_id, "checkpoints")
+        checkpoint_path = get_checkpoint_path_from_dir(checkpoint_dir)
+
+        self.tokenizer = Tokenizer(cfg.DATASET.vocab_dir)
+        # override ignore_rgb for visualization
+        cfg.DATASET.ignore_rgb = False
+        self.dataset = SemanticArrangementDataset(tokenizer=self.tokenizer, **cfg.DATASET)
+
+        self.sampler = Sampler(ConditionalPoseDiffusionModel, checkpoint_path, self.device)
+
+    def visualize_scene(self, di, session_id):
+        raw_datum = self.dataset.get_raw_data(di)
+        language_command = self.tokenizer.convert_structure_params_to_natural_language(raw_datum["sentence"])
+
+        obj_xyz = raw_datum["pcs"]
+        scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in obj_xyz], [xyz[:, 3:] for xyz in obj_xyz], return_scene=True)
+
+        scene.apply_transform(tra.euler_matrix(np.pi, 0, np.pi/2))
+
+        scene_filename = "./tmp_data/input_scene_{}.glb".format(session_id)
+        scene.export(scene_filename)
+
+        return language_command, scene_filename
+
+    def infer(self, di, session_id, progress=gr.Progress()):
+
+        # di = np.random.choice(len(self.dataset))
+
+        raw_datum = self.dataset.get_raw_data(di)
+        print(self.tokenizer.convert_structure_params_to_natural_language(raw_datum["sentence"]))
+        datum = self.dataset.convert_to_tensors(raw_datum, self.tokenizer)
+        batch = self.dataset.single_datum_to_batch(datum, args.num_samples, self.device, inference_mode=True)
+
+        num_poses = datum["goal_poses"].shape[0]
+        xs = self.sampler.sample(batch, num_poses, progress)
+
+        struct_pose, pc_poses_in_struct = get_struct_objs_poses(xs[0])
+        new_obj_xyzs = move_pc_and_create_scene_simple(batch["pcs"], struct_pose, pc_poses_in_struct)
+
+        # vis
+        vis_obj_xyzs = new_obj_xyzs[:3]
+        if torch.is_tensor(vis_obj_xyzs):
+            if vis_obj_xyzs.is_cuda:
+                vis_obj_xyzs = vis_obj_xyzs.detach().cpu()
+            vis_obj_xyzs = vis_obj_xyzs.numpy()
+
+        # for bi, vis_obj_xyz in enumerate(vis_obj_xyzs):
+        #     if verbose:
+        #         print("example {}".format(bi))
+        #         print(vis_obj_xyz.shape)
+        #
+        #     if trimesh:
+        #         show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], [xyz[:, 3:] for xyz in vis_obj_xyz])
+        vis_obj_xyz = vis_obj_xyzs[0]
+        scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], [xyz[:, 3:] for xyz in vis_obj_xyz], return_scene=True)
+
+        scene.apply_transform(tra.euler_matrix(np.pi, 0, np.pi/2))
+
+        scene_filename = "./tmp_data/output_scene_{}.glb".format(session_id)
+        scene.export(scene_filename)
+
+        # pc_filename = "/home/weiyu/Research/StructDiffusion/StructDiffusion/interactive_demo/tmp_data/pc.glb"
+        # scene_filename = "/home/weiyu/Research/StructDiffusion/StructDiffusion/interactive_demo/tmp_data/scene.glb"
+        #
+        # vis_obj_xyz = vis_obj_xyz.reshape(-1, 6)
+        # vis_pc = trimesh.PointCloud(vis_obj_xyz[:, :3], colors=np.concatenate([vis_obj_xyz[:, 3:] * 255, np.ones([vis_obj_xyz.shape[0], 1]) * 255], axis=-1))
+        # vis_pc.export(pc_filename)
+        #
+        # scene = trimesh.Scene()
+        # # add the coordinate frame first
+        # # geom = trimesh.creation.axis(0.01)
+        # # scene.add_geometry(geom)
+        # table = trimesh.creation.box(extents=[1.0, 1.0, 0.02])
+        # table.apply_translation([0.5, 0, -0.01])
+        # table.visual.vertex_colors = [150, 111, 87, 125]
+        # scene.add_geometry(table)
+        # # bounds = trimesh.creation.box(extents=[4.0, 4.0, 4.0])
+        # # bounds = trimesh.creation.icosphere(subdivisions=3, radius=3.1)
+        # # bounds.apply_translation([0, 0, 0])
+        # # bounds.visual.vertex_colors = [30, 30, 30, 30]
+        # # scene.add_geometry(bounds)
+        # # RT_4x4 = np.array([[-0.39560353822208355, -0.9183993826406329, 0.006357240869497738, 0.2651463080169481],
+        # #                    [-0.797630370081598, 0.3401340617616391, -0.4980909683511864, 0.2225696480721997],
+        # #                    [0.45528412367406523, -0.2021172778236285, -0.8671014777611122, 0.9449050652025951],
+        # #                    [0.0, 0.0, 0.0, 1.0]])
+        # # RT_4x4 = np.linalg.inv(RT_4x4)
+        # # RT_4x4 = RT_4x4 @ np.diag([1, -1, -1, 1])
+        # # scene.camera_transform = RT_4x4
+        #
+        # mesh_list = trimesh.util.concatenate(scene.dump())
+        # print(mesh_list)
+        # trimesh.io.export.export_mesh(mesh_list, scene_filename, file_type='obj')
+
+        return scene_filename
+
+    def infer_new(self, di, session_id, progress=gr.Progress()):
+
+        # di = np.random.choice(len(self.dataset))
+
+        raw_datum = self.dataset.get_raw_data(di)
+        print(self.tokenizer.convert_structure_params_to_natural_language(raw_datum["sentence"]))
+        datum = self.dataset.convert_to_tensors(raw_datum, self.tokenizer)
+        batch = self.dataset.single_datum_to_batch(datum, args.num_samples, self.device, inference_mode=True)
+
+        num_poses = datum["goal_poses"].shape[0]
+        xs = self.sampler.sample(batch, num_poses, progress)
+
+        struct_pose, pc_poses_in_struct = get_struct_objs_poses(xs[0])
+        new_obj_xyzs = move_pc_and_create_scene_simple(batch["pcs"], struct_pose, pc_poses_in_struct)
+
+        # vis
+        vis_obj_xyzs = new_obj_xyzs[:3]
+        if torch.is_tensor(vis_obj_xyzs):
+            if vis_obj_xyzs.is_cuda:
+                vis_obj_xyzs = vis_obj_xyzs.detach().cpu()
+            vis_obj_xyzs = vis_obj_xyzs.numpy()
+
+        # for bi, vis_obj_xyz in enumerate(vis_obj_xyzs):
+        #     if verbose:
+        #         print("example {}".format(bi))
+        #         print(vis_obj_xyz.shape)
+        #
+        #     if trimesh:
+        #         show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], [xyz[:, 3:] for xyz in vis_obj_xyz])
+        vis_obj_xyz = vis_obj_xyzs[0]
+        scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], [xyz[:, 3:] for xyz in vis_obj_xyz], return_scene=True)
+
+        scene.apply_transform(tra.euler_matrix(np.pi, 0, np.pi/2))
+
+        scene_filename = "./tmp_data/output_scene_{}.glb".format(session_id)
+        scene.export(scene_filename)
+
+        # pc_filename = "/home/weiyu/Research/StructDiffusion/StructDiffusion/interactive_demo/tmp_data/pc.glb"
+        # scene_filename = "/home/weiyu/Research/StructDiffusion/StructDiffusion/interactive_demo/tmp_data/scene.glb"
+        #
+        # vis_obj_xyz = vis_obj_xyz.reshape(-1, 6)
+        # vis_pc = trimesh.PointCloud(vis_obj_xyz[:, :3], colors=np.concatenate([vis_obj_xyz[:, 3:] * 255, np.ones([vis_obj_xyz.shape[0], 1]) * 255], axis=-1))
+        # vis_pc.export(pc_filename)
+        #
+        # scene = trimesh.Scene()
+        # # add the coordinate frame first
+        # # geom = trimesh.creation.axis(0.01)
+        # # scene.add_geometry(geom)
+        # table = trimesh.creation.box(extents=[1.0, 1.0, 0.02])
+        # table.apply_translation([0.5, 0, -0.01])
+        # table.visual.vertex_colors = [150, 111, 87, 125]
+        # scene.add_geometry(table)
+        # # bounds = trimesh.creation.box(extents=[4.0, 4.0, 4.0])
+        # # bounds = trimesh.creation.icosphere(subdivisions=3, radius=3.1)
+        # # bounds.apply_translation([0, 0, 0])
+        # # bounds.visual.vertex_colors = [30, 30, 30, 30]
+        # # scene.add_geometry(bounds)
+        # # RT_4x4 = np.array([[-0.39560353822208355, -0.9183993826406329, 0.006357240869497738, 0.2651463080169481],
+        # #                    [-0.797630370081598, 0.3401340617616391, -0.4980909683511864, 0.2225696480721997],
+        # #                    [0.45528412367406523, -0.2021172778236285, -0.8671014777611122, 0.9449050652025951],
+        # #                    [0.0, 0.0, 0.0, 1.0]])
+        # # RT_4x4 = np.linalg.inv(RT_4x4)
+        # # RT_4x4 = RT_4x4 @ np.diag([1, -1, -1, 1])
+        # # scene.camera_transform = RT_4x4
+        #
+        # mesh_list = trimesh.util.concatenate(scene.dump())
+        # print(mesh_list)
+        # trimesh.io.export.export_mesh(mesh_list, scene_filename, file_type='obj')
+
+        return scene_filename
+
+
+args = OmegaConf.create()
+args.base_config_file = "./configs/base.yaml"
+args.config_file = "./configs/conditional_pose_diffusion.yaml"
+args.checkpoint_id = "ConditionalPoseDiffusion"
+args.eval_random_seed = 42
+args.num_samples = 1
+
+base_cfg = OmegaConf.load(args.base_config_file)
+cfg = OmegaConf.load(args.config_file)
+cfg = OmegaConf.merge(base_cfg, cfg)
+
+infer_wrapper = Infer_Wrapper(args, cfg)
+
+# version 0
+# demo = gr.Interface(
+#     fn=infer_wrapper.run,
+#     inputs=gr.Slider(0, len(infer_wrapper.dataset)),
+#     # clear color range [0-1.0]
+#     outputs=gr.Model3D(clear_color=[0, 0, 0, 0],  label="3D Model")
+# )
+#
+# demo.launch()
+
+# version 1
+demo = gr.Blocks(theme=gr.themes.Soft())
+with demo:
+    gr.Markdown("<p style='text-align:center;font-size:18px'><b>StructDiffusion Demo</b></p>")
+    # font-size:18px
+    gr.Markdown("<p style='text-align:center'>StructDiffusion combines a diffusion model and an object-centric transformer to construct structures given partial-view point clouds and high-level language goals.<br><a href='https://structdiffusion.github.io/'>Website</a> | <a href='https://github.com/StructDiffusion/StructDiffusion'>Code</a></p>")
+
+    session_id = gr.State(value=np.random.randint(0, 1000))
+    data_selection = gr.Number(label="Example No.", minimum=0, maximum=len(infer_wrapper.dataset) - 1, precision=0)
+    input_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Input 3D Scene")
+    language_command = gr.Textbox(label="Input Language Command")
+    output_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Generated 3D Structure")
+
+    b1 = gr.Button("Show Input Language and Scene")
+    b2 = gr.Button("Generate 3D Structure")
+
+    b1.click(infer_wrapper.visualize_scene, inputs=[data_selection, session_id], outputs=[language_command, input_scene])
+    b2.click(infer_wrapper.infer, inputs=[data_selection, session_id], outputs=output_scene)
+
+demo.queue(concurrency_count=10)
+demo.launch()

app_v1.py

@@ -0,0 +1,217 @@
+import os
+import argparse
+import torch
+import trimesh
+import numpy as np
+import pytorch_lightning as pl
+import gradio as gr
+from omegaconf import OmegaConf
+
+import sys
+sys.path.append('./src')
+
+from StructDiffusion.data.semantic_arrangement_language_demo import SemanticArrangementDataset
+from StructDiffusion.language.tokenizer import Tokenizer
+from StructDiffusion.models.pl_models import ConditionalPoseDiffusionModel, PairwiseCollisionModel
+from StructDiffusion.diffusion.sampler import Sampler, SamplerV2
+from StructDiffusion.diffusion.pose_conversion import get_struct_objs_poses
+from StructDiffusion.utils.files import get_checkpoint_path_from_dir
+from StructDiffusion.utils.rearrangement import show_pcs_with_trimesh, get_trimesh_scene_with_table
+import StructDiffusion.utils.transformations as tra
+from StructDiffusion.language.sentence_encoder import SentenceBertEncoder
+
+
+def move_pc_and_create_scene_simple(obj_xyzs, struct_pose, pc_poses_in_struct):
+
+    device = obj_xyzs.device
+
+    # obj_xyzs: B, N, P, 3 or 6
+    # struct_pose: B, 1, 4, 4
+    # pc_poses_in_struct: B, N, 4, 4
+
+    B, N, _, _ = pc_poses_in_struct.shape
+    _, _, P, _ = obj_xyzs.shape
+
+    current_pc_poses = torch.eye(4).repeat(B, N, 1, 1).to(device)  # B, N, 4, 4
+    # print(torch.mean(obj_xyzs, dim=2).shape)
+    current_pc_poses[:, :, :3, 3] = torch.mean(obj_xyzs[:, :, :, :3], dim=2)  # B, N, 4, 4
+    current_pc_poses = current_pc_poses.reshape(B * N, 4, 4)  # B x N, 4, 4
+
+    struct_pose = struct_pose.repeat(1, N, 1, 1) # B, N, 4, 4
+    struct_pose = struct_pose.reshape(B * N, 4, 4)  # B x 1, 4, 4
+    pc_poses_in_struct = pc_poses_in_struct.reshape(B * N, 4, 4)  # B x N, 4, 4
+
+    goal_pc_pose = struct_pose @ pc_poses_in_struct  # B x N, 4, 4
+    # print("goal pc poses")
+    # print(goal_pc_pose)
+    goal_pc_transform = goal_pc_pose @ torch.inverse(current_pc_poses)  # B x N, 4, 4
+
+    # # important: pytorch3d uses row-major ordering, need to transpose each transformation matrix
+    # transpose = tra3d.Transform3d(matrix=goal_pc_transform.transpose(1, 2))
+    # new_obj_xyzs = obj_xyzs.reshape(B * N, P, -1)  # B x N, P, 3
+    # new_obj_xyzs[:, :, :3] = transpose.transform_points(new_obj_xyzs[:, :, :3])
+
+    # a verision that does not rely on pytorch3d
+    new_obj_xyzs = obj_xyzs.reshape(B * N, P, -1)[:, :, :3]  # B x N, P, 3
+    new_obj_xyzs = torch.concat([new_obj_xyzs, torch.ones(B * N, P, 1).to(device)], dim=-1) # B x N, P, 4
+    new_obj_xyzs = torch.einsum('bij,bkj->bki', goal_pc_transform, new_obj_xyzs)[:, :, :3]  # # B x N, P, 3
+
+    # put it back to B, N, P, 3
+    obj_xyzs[:, :, :, :3] = new_obj_xyzs.reshape(B, N, P, -1)
+
+    return obj_xyzs
+
+
+class Infer_Wrapper:
+
+    def __init__(self, args, cfg):
+
+        # load
+        pl.seed_everything(args.eval_random_seed)
+        self.device = (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
+
+        diffusion_checkpoint_path = get_checkpoint_path_from_dir(os.path.join(cfg.WANDB.save_dir, cfg.WANDB.project, args.diffusion_checkpoint_id, "checkpoints"))
+        collision_checkpoint_path = get_checkpoint_path_from_dir(os.path.join(cfg.WANDB.save_dir, cfg.WANDB.project, args.collision_checkpoint_id, "checkpoints"))
+
+        self.tokenizer = Tokenizer(cfg.DATASET.vocab_dir)
+        # override ignore_rgb for visualization
+        cfg.DATASET.ignore_rgb = False
+        self.dataset = SemanticArrangementDataset(tokenizer=self.tokenizer, **cfg.DATASET)
+
+        self.sampler = SamplerV2(ConditionalPoseDiffusionModel, diffusion_checkpoint_path,
+                                 PairwiseCollisionModel, collision_checkpoint_path, self.device)
+
+    def visualize_scene(self, di, session_id):
+
+        raw_datum = self.dataset.get_raw_data(di, inference_mode=True, shuffle_object_index=True)
+        language_command = raw_datum["template_sentence"]
+
+        obj_xyz = raw_datum["pcs"]
+        scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in obj_xyz], [xyz[:, 3:] for xyz in obj_xyz], return_scene=True)
+
+        scene.apply_transform(tra.euler_matrix(np.pi, 0, np.pi/2))
+
+        scene_filename = "./tmp_data/input_scene_{}.glb".format(session_id)
+        scene.export(scene_filename)
+
+        return language_command, scene_filename
+
+
+    def infer(self, di, session_id, progress=gr.Progress()):
+
+        # di = np.random.choice(len(self.dataset))
+
+        raw_datum = self.dataset.get_raw_data(di, inference_mode=True, shuffle_object_index=True)
+        print(raw_datum["template_sentence"])
+        datum = self.dataset.convert_to_tensors(raw_datum, self.tokenizer, use_sentence_embedding=self.dataset.use_sentence_embedding)
+        batch = self.dataset.single_datum_to_batch(datum, args.num_samples, self.device, inference_mode=True)
+
+        num_poses = datum["goal_poses"].shape[0]
+        struct_pose, pc_poses_in_struct = self.sampler.sample(batch, num_poses, args.num_elites, args.discriminator_batch_size)
+
+        new_obj_xyzs = move_pc_and_create_scene_simple(batch["pcs"][:args.num_elites], struct_pose, pc_poses_in_struct)
+
+        # vis
+        vis_obj_xyzs = new_obj_xyzs[:3]
+        if torch.is_tensor(vis_obj_xyzs):
+            if vis_obj_xyzs.is_cuda:
+                vis_obj_xyzs = vis_obj_xyzs.detach().cpu()
+            vis_obj_xyzs = vis_obj_xyzs.numpy()
+
+        # for bi, vis_obj_xyz in enumerate(vis_obj_xyzs):
+        #     if verbose:
+        #         print("example {}".format(bi))
+        #         print(vis_obj_xyz.shape)
+        #
+        #     if trimesh:
+        #         show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], [xyz[:, 3:] for xyz in vis_obj_xyz])
+        vis_obj_xyz = vis_obj_xyzs[0]
+        scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], [xyz[:, 3:] for xyz in vis_obj_xyz], return_scene=True)
+
+        scene.apply_transform(tra.euler_matrix(np.pi, 0, np.pi/2))
+
+        scene_filename = "./tmp_data/output_scene_{}.glb".format(session_id)
+        scene.export(scene_filename)
+
+        # pc_filename = "/home/weiyu/Research/StructDiffusion/StructDiffusion/interactive_demo/tmp_data/pc.glb"
+        # scene_filename = "/home/weiyu/Research/StructDiffusion/StructDiffusion/interactive_demo/tmp_data/scene.glb"
+        #
+        # vis_obj_xyz = vis_obj_xyz.reshape(-1, 6)
+        # vis_pc = trimesh.PointCloud(vis_obj_xyz[:, :3], colors=np.concatenate([vis_obj_xyz[:, 3:] * 255, np.ones([vis_obj_xyz.shape[0], 1]) * 255], axis=-1))
+        # vis_pc.export(pc_filename)
+        #
+        # scene = trimesh.Scene()
+        # # add the coordinate frame first
+        # # geom = trimesh.creation.axis(0.01)
+        # # scene.add_geometry(geom)
+        # table = trimesh.creation.box(extents=[1.0, 1.0, 0.02])
+        # table.apply_translation([0.5, 0, -0.01])
+        # table.visual.vertex_colors = [150, 111, 87, 125]
+        # scene.add_geometry(table)
+        # # bounds = trimesh.creation.box(extents=[4.0, 4.0, 4.0])
+        # # bounds = trimesh.creation.icosphere(subdivisions=3, radius=3.1)
+        # # bounds.apply_translation([0, 0, 0])
+        # # bounds.visual.vertex_colors = [30, 30, 30, 30]
+        # # scene.add_geometry(bounds)
+        # # RT_4x4 = np.array([[-0.39560353822208355, -0.9183993826406329, 0.006357240869497738, 0.2651463080169481],
+        # #                    [-0.797630370081598, 0.3401340617616391, -0.4980909683511864, 0.2225696480721997],
+        # #                    [0.45528412367406523, -0.2021172778236285, -0.8671014777611122, 0.9449050652025951],
+        # #                    [0.0, 0.0, 0.0, 1.0]])
+        # # RT_4x4 = np.linalg.inv(RT_4x4)
+        # # RT_4x4 = RT_4x4 @ np.diag([1, -1, -1, 1])
+        # # scene.camera_transform = RT_4x4
+        #
+        # mesh_list = trimesh.util.concatenate(scene.dump())
+        # print(mesh_list)
+        # trimesh.io.export.export_mesh(mesh_list, scene_filename, file_type='obj')
+
+        return scene_filename
+
+
+args = OmegaConf.create()
+args.base_config_file = "./configs/base.yaml"
+args.config_file = "./configs/conditional_pose_diffusion_language.yaml"
+args.diffusion_checkpoint_id = "ConditionalPoseDiffusionLanguage"
+args.collision_checkpoint_id = "CollisionDiscriminator"
+args.eval_random_seed = 42
+args.num_samples = 50
+args.num_elites = 3
+args.discriminator_batch_size = 10
+
+base_cfg = OmegaConf.load(args.base_config_file)
+cfg = OmegaConf.load(args.config_file)
+cfg = OmegaConf.merge(base_cfg, cfg)
+
+infer_wrapper = Infer_Wrapper(args, cfg)
+
+# version 0
+# demo = gr.Interface(
+#     fn=infer_wrapper.run,
+#     inputs=gr.Slider(0, len(infer_wrapper.dataset)),
+#     # clear color range [0-1.0]
+#     outputs=gr.Model3D(clear_color=[0, 0, 0, 0],  label="3D Model")
+# )
+#
+# demo.launch()
+
+# version 1
+demo = gr.Blocks(theme=gr.themes.Soft())
+with demo:
+    gr.Markdown("<p style='text-align:center;font-size:18px'><b>StructDiffusion Demo</b></p>")
+    # font-size:18px
+    gr.Markdown("<p style='text-align:center'>StructDiffusion combines a diffusion model and an object-centric transformer to construct structures given partial-view point clouds and high-level language goals.<br><a href='https://structdiffusion.github.io/'>Website</a> | <a href='https://github.com/StructDiffusion/StructDiffusion'>Code</a></p>")
+
+    session_id = gr.State(value=np.random.randint(0, 1000))
+    data_selection = gr.Number(label="Example No.", minimum=0, maximum=len(infer_wrapper.dataset) - 1, precision=0)
+    input_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Input 3D Scene")
+    language_command = gr.Textbox(label="Input Language Command")
+    output_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Generated 3D Structure")
+
+    b1 = gr.Button("Show Input Language and Scene")
+    b2 = gr.Button("Generate 3D Structure")
+
+    b1.click(infer_wrapper.visualize_scene, inputs=[data_selection, session_id], outputs=[language_command, input_scene])
+    b2.click(infer_wrapper.infer, inputs=[data_selection, session_id], outputs=output_scene)
+
+demo.queue(concurrency_count=10)
+demo.launch()

configs/conditional_pose_diffusion_language.yaml

@@ -0,0 +1,92 @@
+random_seed: 1
+
+WANDB:
+  project: StructDiffusion
+  save_dir: ${base_dirs.wandb_dir}
+  name: conditional_pose_diffusion_language_shuffle
+
+DATASET:
+  data_root: ${base_dirs.data}
+  vocab_dir: ${base_dirs.data}/type_vocabs_coarse.json
+
+  # important
+  use_virtual_structure_frame: True
+  ignore_distractor_objects: True
+  ignore_rgb: True
+
+  # the following are determined by the dataset
+  max_num_target_objects: 7
+  max_num_distractor_objects: 5
+  # set to 1 because we use sentence embedding, which only takes one spot in the input seq to transformer diffusion
+  max_num_shape_parameters: 1
+  # set to zeros because they are not used for now
+  max_num_rearrange_features: 0
+  max_num_anchor_features: 0
+
+  # language
+  sentence_embedding_file: ${base_dirs.data}/template_sentence_data.pkl
+  use_incomplete_sentence: True
+
+  # shuffle
+  shuffle_object_index: True
+
+  num_pts: 1024
+  filter_num_moved_objects_range:
+  data_augmentation: False
+
+DATALOADER:
+  batch_size: 64
+  num_workers: 8
+  pin_memory: True
+
+MODEL:
+  # transformer encoder
+  encoder_input_dim: 256
+  num_attention_heads: 8
+  encoder_hidden_dim: 512
+  encoder_dropout: 0.0
+  encoder_activation: relu
+  encoder_num_layers: 8
+  # output head
+  structure_dropout: 0
+  object_dropout: 0
+  # pc encoder
+  ignore_rgb: ${DATASET.ignore_rgb}
+  pc_emb_dim: 256
+  posed_pc_emb_dim: 80
+  # pose encoder
+  pose_emb_dim: 80
+  # language
+  word_emb_dim: 160
+  # diffusion step
+  time_emb_dim: 80
+  # sequence embeddings
+  # max_num_target_objects (+ max_num_distractor_objects if not ignore_distractor_objects)
+  max_seq_size: 7
+  max_token_type_size: 4
+  seq_pos_emb_dim: 8
+  seq_type_emb_dim: 8
+  # virtual frame
+  use_virtual_structure_frame: ${DATASET.use_virtual_structure_frame}
+  # language
+  use_sentence_embedding: True
+  sentence_embedding_dim: 384
+
+NOISE_SCHEDULE:
+  timesteps: 200
+
+LOSS:
+  type: huber
+
+OPTIMIZER:
+  lr: 0.0001
+  weight_decay: 0  #0.0001
+  # lr_restart: 3000
+  # warmup: 10
+
+TRAINER:
+  max_epochs: 200
+  gradient_clip_val: 1.0
+  gpus: 1
+  deterministic: False
+  # enable_progress_bar: False

data/template_sentence_data.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bbb41fd847ca6ee48d484c346fd8b0bbf478dfcc4559bf9646e3b7e7bf9fe83b
+size 5680159

requirements.txt

@@ -7,4 +7,5 @@ pyglet==1.5.0
 openpyxl
 pytorch_lightning==1.6.1
 wandb===0.13.10
-omegaconf==2.2.2
+omegaconf==2.2.2
+sentence-transformers

src/StructDiffusion/data/pairwise_collision.py

@@ -32,11 +32,27 @@ def load_pairwise_collision_data(h5_filename):
     return data_dict
 
 
+def replace_root_directory(original_filename: str, new_root: str) -> str:
+    # Split the original filename into a list by directory
+    original_parts = original_filename.split('/')
+
+    # Find the index of the "data_new_objects" part
+    data_index = original_parts.index('data_new_objects')
+
+    # Split the new root into a list by directory
+    new_root_parts = new_root.split('/')
+
+    # Combine the new root with the rest of the original filename
+    updated_filename = '/'.join(new_root_parts + original_parts[data_index + 1:])
+
+    return updated_filename
+
+
 class PairwiseCollisionDataset(torch.utils.data.Dataset):
 
     def __init__(self, urdf_pc_idx_file, collision_data_dir, random_rotation=True,
                  num_pts=1024, normalize_pc=True, num_scene_pts=2048, data_augmentation=False,
-                 debug=False):
+                 debug=False, new_data_root=None):
 
         # load dictionary mapping from urdf to list of pc data, each sample is
         #   {"step_t": step_t, "obj": obj, "filename": filename}
@@ -49,6 +65,8 @@ class PairwiseCollisionDataset(torch.utils.data.Dataset):
                 filename =  pd["filename"]
                 if "data00026058" in filename or "data00011415" in filename or "data00026061" in filename or "data00700565" in filename or "data00505290" in filename:
                     continue
+                if new_data_root:
+                    pd["filename"] = replace_root_directory(pd["filename"], new_data_root)
                 valid_pc_data.append(pd)
             if valid_pc_data:
                 self.urdf_to_pc_data[urdf] = valid_pc_data
@@ -297,65 +315,3 @@ class PairwiseCollisionDataset(torch.utils.data.Dataset):
             "label": torch.FloatTensor([label]),
         }
         return datum
-
-    # @staticmethod
-    # def collate_fn(data):
-    #     """
-    #     :param data:
-    #     :return:
-    #     """
-    #
-    #     batched_data_dict = {}
-    #     for key in ["is_circle"]:
-    #         batched_data_dict[key] = torch.cat([dict[key] for dict in data], dim=0)
-    #     for key in ["scene_xyz"]:
-    #         batched_data_dict[key] = torch.stack([dict[key] for dict in data], dim=0)
-    #
-    #     return batched_data_dict
-    #
-    # # def create_pair_xyzs_from_obj_xyzs(self, new_obj_xyzs, debug=False):
-    # #
-    # #     new_obj_xyzs = [xyz.cpu().numpy() for xyz in new_obj_xyzs]
-    # #
-    # #     # compute pairwise collision
-    # #     scene_xyzs = []
-    # #     obj_xyz_pair_idxs = list(itertools.combinations(range(len(new_obj_xyzs)), 2))
-    # #
-    # #     for obj_xyz_pair_idx in obj_xyz_pair_idxs:
-    # #         obj_xyz_pair = [new_obj_xyzs[obj_xyz_pair_idx[0]], new_obj_xyzs[obj_xyz_pair_idx[1]]]
-    # #         num_indicator = 2
-    # #         obj_xyz_pair_ind = []
-    # #         for oi, obj_xyz in enumerate(obj_xyz_pair):
-    # #             obj_xyz = np.concatenate([obj_xyz, np.tile(np.eye(num_indicator)[oi], (obj_xyz.shape[0], 1))], axis=1)
-    # #             obj_xyz_pair_ind.append(obj_xyz)
-    # #         pair_scene_xyz = np.concatenate(obj_xyz_pair_ind, axis=0)
-    # #
-    # #         # subsampling and normalizing pc
-    # #         rand_idx = np.random.randint(0, pair_scene_xyz.shape[0], self.num_scene_pts)
-    # #         pair_scene_xyz = pair_scene_xyz[rand_idx]
-    # #         if self.normalize_pc:
-    # #             pair_scene_xyz[:, 0:3] = pc_normalize(pair_scene_xyz[:, 0:3])
-    # #
-    # #         scene_xyzs.append(array_to_tensor(pair_scene_xyz))
-    # #
-    # #     if debug:
-    # #         for scene_xyz in scene_xyzs:
-    # #             show_pcs([scene_xyz[:, 0:3]], [np.tile(np.array([0, 1, 0], dtype=np.float), (scene_xyz.shape[0], 1))],
-    # #                      add_coordinate_frame=True)
-    # #
-    # #     return scene_xyzs
-
-
-if __name__ == "__main__":
-    dataset = PairwiseCollisionDataset(urdf_pc_idx_file="/home/weiyu/data_drive/StructDiffusion/pairwise_collision_data/urdf_pc_idx.pkl",
-                      collision_data_dir="/home/weiyu/data_drive/StructDiffusion/pairwise_collision_data",
-                      debug=False)
-
-    for i in tqdm.tqdm(np.random.permutation(len(dataset))):
-        # print(i)
-        d = dataset[i]
-        # print(d["label"])
-
-    # dl = torch.utils.data.DataLoader(dataset, batch_size=32, num_workers=8)
-    # for b in tqdm.tqdm(dl):
-    #     pass

src/StructDiffusion/data/semantic_arrangement.py

@@ -533,47 +533,4 @@ def compute_min_max(dataloader):
         current_min, _ = torch.min(goal_poses, dim=0)
         max_value[max_value < current_max] = current_max[max_value < current_max]
         max_value[max_value > current_min] = current_min[max_value > current_min]
-    print(f"{min_value} - {max_value}")
-
-
-if __name__ == "__main__":
-
-    tokenizer = Tokenizer("/home/weiyu/data_drive/data_new_objects/type_vocabs_coarse.json")
-
-    data_roots = []
-    index_roots = []
-    for shape, index in [("circle", "index_10k"), ("line", "index_10k"), ("stacking", "index_10k"), ("dinner", "index_10k")]:
-        data_roots.append("/home/weiyu/data_drive/data_new_objects/examples_{}_new_objects/result".format(shape))
-        index_roots.append(index)
-
-    dataset = SemanticArrangementDataset(data_roots=data_roots,
-                                         index_roots=index_roots,
-                                         split="valid", tokenizer=tokenizer,
-                                         max_num_target_objects=7,
-                                         max_num_distractor_objects=5,
-                                         max_num_shape_parameters=5,
-                                         max_num_rearrange_features=0,
-                                         max_num_anchor_features=0,
-                                         num_pts=1024,
-                                         use_virtual_structure_frame=True,
-                                         ignore_distractor_objects=True,
-                                         ignore_rgb=True,
-                                         filter_num_moved_objects_range=None,  # [5, 5]
-                                         data_augmentation=False,
-                                         shuffle_object_index=False,
-                                         debug=False)
-
-    # print(len(dataset))
-    # for d in dataset:
-    #     print("\n\n" + "="*100)
-
-    dataloader = DataLoader(dataset, batch_size=64, shuffle=False, num_workers=8)
-    for i, d in enumerate(tqdm(dataloader)):
-        pass
-        # for k in d:
-        #     if isinstance(d[k], torch.Tensor):
-        #         print("--size", k, d[k].shape)
-        # for k in d:
-        #     print(k, d[k])
-        #
-        # input("next?")
+    print(f"{min_value} - {max_value}")

src/StructDiffusion/data/semantic_arrangement_language.py

@@ -0,0 +1,633 @@
+import copy
+import cv2
+import h5py
+import numpy as np
+import os
+import trimesh
+import torch
+from tqdm import tqdm
+import json
+import random
+import pickle
+
+from torch.utils.data import DataLoader
+
+# Local imports
+from StructDiffusion.utils.rearrangement import show_pcs, get_pts, combine_and_sample_xyzs
+from StructDiffusion.language.tokenizer import Tokenizer
+
+import StructDiffusion.utils.brain2.camera as cam
+import StructDiffusion.utils.brain2.image as img
+import StructDiffusion.utils.transformations as tra
+
+
+class SemanticArrangementDataset(torch.utils.data.Dataset):
+
+    def __init__(self, data_roots, index_roots, split, tokenizer,
+                 max_num_target_objects=11, max_num_distractor_objects=5,
+                 max_num_shape_parameters=7, max_num_rearrange_features=1, max_num_anchor_features=3,
+                 num_pts=1024,
+                 use_virtual_structure_frame=True, ignore_distractor_objects=True, ignore_rgb=True,
+                 filter_num_moved_objects_range=None, shuffle_object_index=False,
+                 sentence_embedding_file=None, use_incomplete_sentence=False,
+                 data_augmentation=True, debug=False, **kwargs):
+        """
+
+        Note: setting filter_num_moved_objects_range=[k, k] and max_num_objects=k will create no padding for target objs
+
+        :param data_root:
+        :param split: train, valid, or test
+        :param shuffle_object_index: whether to shuffle the positions of target objects and other objects in the sequence
+        :param debug:
+        :param max_num_shape_parameters:
+        :param max_num_objects:
+        :param max_num_rearrange_features:
+        :param max_num_anchor_features:
+        :param num_pts:
+        :param use_stored_arrangement_indices:
+        :param kwargs:
+        """
+
+        self.use_virtual_structure_frame = use_virtual_structure_frame
+        self.ignore_distractor_objects = ignore_distractor_objects
+        self.ignore_rgb = ignore_rgb and not debug
+
+        self.num_pts = num_pts
+        self.debug = debug
+
+        self.max_num_objects = max_num_target_objects
+        self.max_num_other_objects = max_num_distractor_objects
+        self.max_num_shape_parameters = max_num_shape_parameters
+        self.max_num_rearrange_features = max_num_rearrange_features
+        self.max_num_anchor_features = max_num_anchor_features
+        self.shuffle_object_index = shuffle_object_index
+
+        # used to tokenize the language part
+        self.tokenizer = tokenizer
+
+        # retrieve data
+        self.data_roots = data_roots
+        self.arrangement_data = []
+        arrangement_steps = []
+        for ddx in range(len(data_roots)):
+            data_root = data_roots[ddx]
+            index_root = index_roots[ddx]
+            arrangement_indices_file = os.path.join(data_root, index_root, "{}_arrangement_indices_file_all.txt".format(split))
+            if os.path.exists(arrangement_indices_file):
+                with open(arrangement_indices_file, "r") as fh:
+                    arrangement_steps.extend([(os.path.join(data_root, f[0]), f[1]) for f in eval(fh.readline().strip())])
+            else:
+                print("{} does not exist".format(arrangement_indices_file))
+        # only keep the goal, ignore the intermediate steps
+        for filename, step_t in arrangement_steps:
+            if step_t == 0:
+                if "data00026058" in filename or "data00011415" in filename or "data00026061" in filename or "data00700565" in filename:
+                    continue
+                self.arrangement_data.append((filename, step_t))
+        # if specified, filter data
+        if filter_num_moved_objects_range is not None:
+            self.arrangement_data = self.filter_based_on_number_of_moved_objects(filter_num_moved_objects_range)
+        print("{} valid sequences".format(len(self.arrangement_data)))
+
+        # language
+        if sentence_embedding_file:
+            assert max_num_shape_parameters == 1
+            # since we do not use them right now, ignore them
+            # assert max_num_rearrange_features == 0
+            # assert max_num_anchor_features == 0
+            with open(sentence_embedding_file, "rb") as fh:
+                template_sentence_data = pickle.load(fh)
+            self.use_sentence_embedding = True
+            self.type_value_tuple_to_template_sentences = template_sentence_data["type_value_tuple_to_template_sentences"]
+            self.template_sentence_to_embedding = template_sentence_data["template_sentence_to_embedding"]
+            self.use_incomplete_sentence = use_incomplete_sentence
+            print("use sentence embedding")
+            print(len(self.type_value_tuple_to_template_sentences))
+            print(len(self.template_sentence_to_embedding))
+        else:
+            self.use_sentence_embedding = False
+
+        # Data Aug
+        self.data_augmentation = data_augmentation
+        # additive noise
+        self.gp_rescale_factor_range = [12, 20]
+        self.gaussian_scale_range = [0., 0.003]
+        # multiplicative noise
+        self.gamma_shape = 1000.
+        self.gamma_scale = 0.001
+
+    def filter_based_on_number_of_moved_objects(self, filter_num_moved_objects_range):
+        assert len(list(filter_num_moved_objects_range)) == 2
+        min_num, max_num = filter_num_moved_objects_range
+        print("Remove scenes that have less than {} or more than {} objects being moved".format(min_num, max_num))
+        ok_data = []
+        for filename, step_t in self.arrangement_data:
+            h5 = h5py.File(filename, 'r')
+            moved_objs = h5['moved_objs'][()].split(',')
+            if min_num <= len(moved_objs) <= max_num:
+                ok_data.append((filename, step_t))
+        print("{} valid sequences left".format(len(ok_data)))
+        return ok_data
+
+    def get_data_idx(self, idx):
+        # Create the datum to return
+        file_idx = np.argmax(idx < self.file_to_count)
+        data = h5py.File(self.data_files[file_idx], 'r')
+        if file_idx > 0:
+            # for lang2sym, idx is always 0
+            idx = idx - self.file_to_count[file_idx - 1]
+        return data, idx, file_idx
+
+    def add_noise_to_depth(self, depth_img):
+        """ add depth noise """
+        multiplicative_noise = np.random.gamma(self.gamma_shape, self.gamma_scale)
+        depth_img = multiplicative_noise * depth_img
+        return depth_img
+
+    def add_noise_to_xyz(self, xyz_img, depth_img):
+        """ TODO: remove this code or at least celean it up"""
+        xyz_img = xyz_img.copy()
+        H, W, C = xyz_img.shape
+        gp_rescale_factor = np.random.randint(self.gp_rescale_factor_range[0],
+                                              self.gp_rescale_factor_range[1])
+        gp_scale = np.random.uniform(self.gaussian_scale_range[0],
+                                     self.gaussian_scale_range[1])
+        small_H, small_W = (np.array([H, W]) / gp_rescale_factor).astype(int)
+        additive_noise = np.random.normal(loc=0.0, scale=gp_scale, size=(small_H, small_W, C))
+        additive_noise = cv2.resize(additive_noise, (W, H), interpolation=cv2.INTER_CUBIC)
+        xyz_img[depth_img > 0, :] += additive_noise[depth_img > 0, :]
+        return xyz_img
+
+    def random_index(self):
+        return self[np.random.randint(len(self))]
+
+    def _get_rgb(self, h5, idx, ee=True):
+        RGB = "ee_rgb" if ee else "rgb"
+        rgb1 = img.PNGToNumpy(h5[RGB][idx])[:, :, :3] / 255.  # remove alpha
+        return rgb1
+
+    def _get_depth(self, h5, idx, ee=True):
+        DEPTH = "ee_depth" if ee else "depth"
+
+    def _get_images(self, h5, idx, ee=True):
+        if ee:
+            RGB, DEPTH, SEG = "ee_rgb", "ee_depth", "ee_seg"
+            DMIN, DMAX = "ee_depth_min", "ee_depth_max"
+        else:
+            RGB, DEPTH, SEG = "rgb", "depth", "seg"
+            DMIN, DMAX = "depth_min", "depth_max"
+        dmin = h5[DMIN][idx]
+        dmax = h5[DMAX][idx]
+        rgb1 = img.PNGToNumpy(h5[RGB][idx])[:, :, :3] / 255.  # remove alpha
+        depth1 = h5[DEPTH][idx] / 20000. * (dmax - dmin) + dmin
+        seg1 = img.PNGToNumpy(h5[SEG][idx])
+
+        valid1 = np.logical_and(depth1 > 0.1, depth1 < 2.)
+
+        # proj_matrix = h5['proj_matrix'][()]
+        camera = cam.get_camera_from_h5(h5)
+        if self.data_augmentation:
+            depth1 = self.add_noise_to_depth(depth1)
+
+        xyz1 = cam.compute_xyz(depth1, camera)
+        if self.data_augmentation:
+            xyz1 = self.add_noise_to_xyz(xyz1, depth1)
+
+        # Transform the point cloud
+        # Here it is...
+        # CAM_POSE = "ee_cam_pose" if ee else "cam_pose"
+        CAM_POSE = "ee_camera_view" if ee else "camera_view"
+        cam_pose = h5[CAM_POSE][idx]
+        if ee:
+            # ee_camera_view has 0s for x, y, z
+            cam_pos = h5["ee_cam_pose"][:][:3, 3]
+            cam_pose[:3, 3] = cam_pos
+
+        # Get transformed point cloud
+        h, w, d = xyz1.shape
+        xyz1 = xyz1.reshape(h * w, -1)
+        xyz1 = trimesh.transform_points(xyz1, cam_pose)
+        xyz1 = xyz1.reshape(h, w, -1)
+
+        scene1 = rgb1, depth1, seg1, valid1, xyz1
+
+        return scene1
+
+    def __len__(self):
+        return len(self.arrangement_data)
+
+    def _get_ids(self, h5):
+        """
+        get object ids
+
+        @param h5:
+        @return:
+        """
+        ids = {}
+        for k in h5.keys():
+            if k.startswith("id_"):
+                ids[k[3:]] = h5[k][()]
+        return ids
+
+    def get_positive_ratio(self):
+        num_pos = 0
+        for d in self.arrangement_data:
+            filename, step_t = d
+            if step_t == 0:
+                num_pos += 1
+        return (len(self.arrangement_data) - num_pos) * 1.0 / num_pos
+
+    def get_object_position_vocab_sizes(self):
+        return self.tokenizer.get_object_position_vocab_sizes()
+
+    def get_vocab_size(self):
+        return self.tokenizer.get_vocab_size()
+
+    def get_data_index(self, idx):
+        filename = self.arrangement_data[idx]
+        return filename
+
+    def get_raw_data(self, idx, inference_mode=False, shuffle_object_index=False):
+        """
+
+        :param idx:
+        :param inference_mode:
+        :param shuffle_object_index: used to test different orders of objects
+        :return:
+        """
+
+        filename, _ = self.arrangement_data[idx]
+
+        h5 = h5py.File(filename, 'r')
+        ids = self._get_ids(h5)
+        all_objs = sorted([o for o in ids.keys() if "object_" in o])
+        goal_specification = json.loads(str(np.array(h5["goal_specification"])))
+        num_rearrange_objs = len(goal_specification["rearrange"]["objects"])
+        num_other_objs = len(goal_specification["anchor"]["objects"] + goal_specification["distract"]["objects"])
+        assert len(all_objs) == num_rearrange_objs + num_other_objs, "{}, {}".format(len(all_objs), num_rearrange_objs + num_other_objs)
+        assert num_rearrange_objs <= self.max_num_objects
+        assert num_other_objs <= self.max_num_other_objects
+
+        # important: only using the last step
+        step_t = num_rearrange_objs
+
+        target_objs = all_objs[:num_rearrange_objs]
+        other_objs = all_objs[num_rearrange_objs:]
+
+        structure_parameters = goal_specification["shape"]
+
+        # Important: ensure the order is correct
+        if structure_parameters["type"] == "circle" or structure_parameters["type"] == "line":
+            target_objs = target_objs[::-1]
+        elif structure_parameters["type"] == "tower" or structure_parameters["type"] == "dinner":
+            target_objs = target_objs
+        else:
+            raise KeyError("{} structure is not recognized".format(structure_parameters["type"]))
+        all_objs = target_objs + other_objs
+
+        ###################################
+        # getting scene images and point clouds
+        scene = self._get_images(h5, step_t, ee=True)
+        rgb, depth, seg, valid, xyz = scene
+        if inference_mode:
+            initial_scene = scene
+
+        # getting object point clouds
+        obj_pcs = []
+        obj_pad_mask = []
+        current_pc_poses = []
+        other_obj_pcs = []
+        other_obj_pad_mask = []
+        for obj in all_objs:
+            obj_mask = np.logical_and(seg == ids[obj], valid)
+            if np.sum(obj_mask) <= 0:
+                raise Exception
+            ok, obj_xyz, obj_rgb, _ = get_pts(xyz, rgb, obj_mask, num_pts=self.num_pts)
+            if not ok:
+                raise Exception
+
+            if obj in target_objs:
+                if self.ignore_rgb:
+                    obj_pcs.append(obj_xyz)
+                else:
+                    obj_pcs.append(torch.concat([obj_xyz, obj_rgb], dim=-1))
+                obj_pad_mask.append(0)
+                pc_pose = np.eye(4)
+                pc_pose[:3, 3] = torch.mean(obj_xyz, dim=0).numpy()
+                current_pc_poses.append(pc_pose)
+            elif obj in other_objs:
+                if self.ignore_rgb:
+                    other_obj_pcs.append(obj_xyz)
+                else:
+                    other_obj_pcs.append(torch.concat([obj_xyz, obj_rgb], dim=-1))
+                other_obj_pad_mask.append(0)
+            else:
+                raise Exception
+
+        ###################################
+        # computes goal positions for objects
+        # Important: because of the noises we added to point clouds, the rearranged point clouds will not be perfect
+        if self.use_virtual_structure_frame:
+            goal_structure_pose = tra.euler_matrix(structure_parameters["rotation"][0], structure_parameters["rotation"][1],
+                                              structure_parameters["rotation"][2])
+            goal_structure_pose[:3, 3] = [structure_parameters["position"][0], structure_parameters["position"][1],
+                                     structure_parameters["position"][2]]
+            goal_structure_pose_inv = np.linalg.inv(goal_structure_pose)
+
+        goal_obj_poses = []
+        current_obj_poses = []
+        goal_pc_poses = []
+        for obj, current_pc_pose in zip(target_objs, current_pc_poses):
+            goal_pose = h5[obj][0]
+            current_pose = h5[obj][step_t]
+            if inference_mode:
+                goal_obj_poses.append(goal_pose)
+                current_obj_poses.append(current_pose)
+
+            goal_pc_pose = goal_pose @ np.linalg.inv(current_pose) @ current_pc_pose
+            if self.use_virtual_structure_frame:
+                goal_pc_pose = goal_structure_pose_inv @ goal_pc_pose
+            goal_pc_poses.append(goal_pc_pose)
+
+        # transform current object point cloud to the goal point cloud in the world frame
+        if self.debug:
+            new_obj_pcs = [copy.deepcopy(pc.numpy()) for pc in obj_pcs]
+            for i, obj_pc in enumerate(new_obj_pcs):
+
+                current_pc_pose = current_pc_poses[i]
+                goal_pc_pose = goal_pc_poses[i]
+                if self.use_virtual_structure_frame:
+                    goal_pc_pose = goal_structure_pose @ goal_pc_pose
+                print("current pc pose", current_pc_pose)
+                print("goal pc pose", goal_pc_pose)
+
+                goal_pc_transform = goal_pc_pose @ np.linalg.inv(current_pc_pose)
+                print("transform", goal_pc_transform)
+                new_obj_pc = copy.deepcopy(obj_pc)
+                new_obj_pc[:, :3] = trimesh.transform_points(obj_pc[:, :3], goal_pc_transform)
+                print(new_obj_pc.shape)
+
+                # visualize rearrangement sequence (new_obj_xyzs), the current object before moving (obj_xyz), and other objects
+                new_obj_pcs[i] = new_obj_pc
+                new_obj_pcs[i][:, 3:] = np.tile(np.array([1, 0, 0], dtype=np.float), (new_obj_pc.shape[0], 1))
+                new_obj_rgb_current = np.tile(np.array([0, 1, 0], dtype=np.float), (new_obj_pc.shape[0], 1))
+                show_pcs([pc[:, :3] for pc in new_obj_pcs] + [pc[:, :3] for pc in other_obj_pcs] + [obj_pc[:, :3]],
+                         [pc[:, 3:] for pc in new_obj_pcs] + [pc[:, 3:] for pc in other_obj_pcs] + [new_obj_rgb_current],
+                         add_coordinate_frame=True)
+            show_pcs([pc[:, :3] for pc in new_obj_pcs], [pc[:, 3:] for pc in new_obj_pcs], add_coordinate_frame=True)
+
+        # pad data
+        for i in range(self.max_num_objects - len(target_objs)):
+            obj_pcs.append(torch.zeros_like(obj_pcs[0], dtype=torch.float32))
+            obj_pad_mask.append(1)
+        for i in range(self.max_num_other_objects - len(other_objs)):
+            other_obj_pcs.append(torch.zeros_like(obj_pcs[0], dtype=torch.float32))
+            other_obj_pad_mask.append(1)
+
+        ###################################
+        # preparing sentence
+        sentence = []
+        sentence_pad_mask = []
+
+        # structure parameters
+        # 5 parameters
+        structure_parameters = goal_specification["shape"]
+        if structure_parameters["type"] == "circle" or structure_parameters["type"] == "line":
+            sentence.append((structure_parameters["type"], "shape"))
+            sentence.append((structure_parameters["rotation"][2], "rotation"))
+            sentence.append((structure_parameters["position"][0], "position_x"))
+            sentence.append((structure_parameters["position"][1], "position_y"))
+            if structure_parameters["type"] == "circle":
+                sentence.append((structure_parameters["radius"], "radius"))
+            elif structure_parameters["type"] == "line":
+                sentence.append((structure_parameters["length"] / 2.0, "radius"))
+            if not self.use_sentence_embedding:
+                for _ in range(5):
+                    sentence_pad_mask.append(0)
+        else:
+            sentence.append((structure_parameters["type"], "shape"))
+            sentence.append((structure_parameters["rotation"][2], "rotation"))
+            sentence.append((structure_parameters["position"][0], "position_x"))
+            sentence.append((structure_parameters["position"][1], "position_y"))
+            if not self.use_sentence_embedding:
+                for _ in range(4):
+                    sentence_pad_mask.append(0)
+                sentence.append(("PAD", None))
+                sentence_pad_mask.append(1)
+
+        if self.use_sentence_embedding:
+
+            if self.use_incomplete_sentence:
+                token_idxs = np.random.permutation(len(sentence))
+                token_idxs = token_idxs[:np.random.randint(1, len(sentence) + 1)]
+                token_idxs = sorted(token_idxs)
+                incomplete_sentence = [sentence[ti] for ti in token_idxs]
+            else:
+                incomplete_sentence = sentence
+
+            type_value_tuple = self.tokenizer.convert_structure_params_to_type_value_tuple(incomplete_sentence)
+            template_sentence = np.random.choice(self.type_value_tuple_to_template_sentences[type_value_tuple])
+            sentence_embedding = self.template_sentence_to_embedding[template_sentence]
+            sentence_pad_mask = [0]
+
+        ###################################
+        # paddings
+        for i in range(self.max_num_objects - len(target_objs)):
+            goal_pc_poses.append(np.eye(4))
+
+        ###################################
+        if self.debug:
+            print("---")
+            print("all objects:", all_objs)
+            print("target objects:", target_objs)
+            print("other objects:", other_objs)
+            print("goal specification:", goal_specification)
+            print("sentence:", sentence)
+            if self.use_sentence_embedding:
+                print("use sentence embedding")
+                if self.use_incomplete_sentence:
+                    print("incomplete_sentence:", incomplete_sentence)
+                print("template sentence:", template_sentence)
+            show_pcs([pc[:, :3] for pc in obj_pcs + other_obj_pcs], [pc[:, 3:] for pc in obj_pcs + other_obj_pcs], add_coordinate_frame=True)
+
+        assert len(obj_pcs) == len(goal_pc_poses)
+        ###################################
+
+        # shuffle the position of objects
+        # important: only shuffle for dinner
+        if shuffle_object_index and structure_parameters["type"] == "dinner":
+            num_target_objs = len(target_objs)
+            shuffle_target_object_indices = list(range(num_target_objs))
+            random.shuffle(shuffle_target_object_indices)
+            shuffle_object_indices = shuffle_target_object_indices + list(range(num_target_objs, self.max_num_objects))
+            obj_pcs = [obj_pcs[i] for i in shuffle_object_indices]
+            goal_pc_poses = [goal_pc_poses[i] for i in shuffle_object_indices]
+            if inference_mode:
+                goal_obj_poses = [goal_obj_poses[i] for i in shuffle_object_indices[:num_target_objs]]
+                current_obj_poses = [current_obj_poses[i] for i in shuffle_object_indices[:num_target_objs]]
+                target_objs = [target_objs[i] for i in shuffle_target_object_indices[:num_target_objs]]
+                current_pc_poses = [current_pc_poses[i] for i in shuffle_object_indices[:num_target_objs]]
+
+        ###################################
+        if self.use_virtual_structure_frame:
+            if self.ignore_distractor_objects:
+                # language, structure virtual frame, target objects
+                pcs = obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [2] + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + [0] + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + [0] + obj_pad_mask
+            else:
+                # language, distractor objects, structure virtual frame, target objects
+                pcs = other_obj_pcs + obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [1] * self.max_num_other_objects + [2] + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + list(range(self.max_num_other_objects)) + [0] + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + other_obj_pad_mask + [0] + obj_pad_mask
+            goal_poses = [goal_structure_pose] + goal_pc_poses
+        else:
+            if self.ignore_distractor_objects:
+                # language, target objects
+                pcs = obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + obj_pad_mask
+            else:
+                # language, distractor objects, target objects
+                pcs = other_obj_pcs + obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [1] * self.max_num_other_objects + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + list(range(self.max_num_other_objects)) + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + other_obj_pad_mask + obj_pad_mask
+            goal_poses = goal_pc_poses
+
+        datum = {
+            "pcs": pcs,
+            "goal_poses": goal_poses,
+            "type_index": type_index,
+            "position_index": position_index,
+            "pad_mask": pad_mask,
+            "t": step_t,
+            "filename": filename
+        }
+        if self.use_sentence_embedding:
+            datum["sentence"] = sentence_embedding
+        else:
+            datum["sentence"] = sentence
+
+        if inference_mode:
+            datum["rgb"] = rgb
+            datum["goal_obj_poses"] = goal_obj_poses
+            datum["current_obj_poses"] = current_obj_poses
+            datum["target_objs"] = target_objs
+            datum["initial_scene"] = initial_scene
+            datum["ids"] = ids
+            datum["goal_specification"] = goal_specification
+            datum["current_pc_poses"] = current_pc_poses
+            if self.use_sentence_embedding:
+                datum["template_sentence"] = template_sentence
+
+        return datum
+
+    @staticmethod
+    def convert_to_tensors(datum, tokenizer, use_sentence_embedding=False):
+        tensors = {
+            "pcs": torch.stack(datum["pcs"], dim=0),
+            "goal_poses": torch.FloatTensor(np.array(datum["goal_poses"])),
+            "type_index": torch.LongTensor(np.array(datum["type_index"])),
+            "position_index": torch.LongTensor(np.array(datum["position_index"])),
+            "pad_mask": torch.LongTensor(np.array(datum["pad_mask"])),
+            "t": datum["t"],
+            "filename": datum["filename"]
+        }
+        if use_sentence_embedding:
+            tensors["sentence"] = torch.FloatTensor(datum["sentence"])  # after batching, B x sentence embed dim
+        else:
+            tensors["sentence"] = torch.LongTensor(np.array([tokenizer.tokenize(*i) for i in datum["sentence"]]))
+        return tensors
+
+    def __getitem__(self, idx):
+
+        datum = self.convert_to_tensors(self.get_raw_data(idx, shuffle_object_index=self.shuffle_object_index),
+                                        self.tokenizer,
+                                        self.use_sentence_embedding)
+
+        return datum
+
+    def single_datum_to_batch(self, x, num_samples, device, inference_mode=True):
+        tensor_x = {}
+
+        tensor_x["pcs"] = x["pcs"].to(device)[None, :, :, :].repeat(num_samples, 1, 1, 1)
+        tensor_x["sentence"] = x["sentence"].to(device)[None, :].repeat(num_samples, 1)
+        if not inference_mode:
+            tensor_x["goal_poses"] = x["goal_poses"].to(device)[None, :, :, :].repeat(num_samples, 1, 1, 1)
+
+        tensor_x["type_index"] = x["type_index"].to(device)[None, :].repeat(num_samples, 1)
+        tensor_x["position_index"] = x["position_index"].to(device)[None, :].repeat(num_samples, 1)
+        tensor_x["pad_mask"] = x["pad_mask"].to(device)[None, :].repeat(num_samples, 1)
+
+        return tensor_x
+
+
+def compute_min_max(dataloader):
+
+    # tensor([-0.3557, -0.3847,  0.0000, -1.0000, -1.0000, -0.4759, -1.0000, -1.0000,
+    #         -0.9079, -0.8668, -0.9105, -0.4186])
+    # tensor([0.3915, 0.3494, 0.3267, 1.0000, 1.0000, 0.8961, 1.0000, 1.0000, 0.8194,
+    #         0.4787, 0.6421, 1.0000])
+    # tensor([0.0918, -0.3758, 0.0000, -1.0000, -1.0000, 0.0000, -1.0000, -1.0000,
+    #         -0.0000, 0.0000, 0.0000, 1.0000])
+    # tensor([0.9199, 0.3710, 0.0000, 1.0000, 1.0000, 0.0000, 1.0000, 1.0000, -0.0000,
+    #         0.0000, 0.0000, 1.0000])
+
+    min_value = torch.ones(16) * 10000
+    max_value = torch.ones(16) * -10000
+    for d in tqdm(dataloader):
+        goal_poses = d["goal_poses"]
+        goal_poses = goal_poses.reshape(-1, 16)
+        current_max, _ = torch.max(goal_poses, dim=0)
+        current_min, _ = torch.min(goal_poses, dim=0)
+        max_value[max_value < current_max] = current_max[max_value < current_max]
+        max_value[max_value > current_min] = current_min[max_value > current_min]
+    print(f"{min_value} - {max_value}")
+
+
+if __name__ == "__main__":
+
+    tokenizer = Tokenizer("/home/weiyu/data_drive/data_new_objects/type_vocabs_coarse.json")
+
+    data_roots = []
+    index_roots = []
+    for shape, index in [("circle", "index_10k"), ("line", "index_10k"), ("stacking", "index_10k"), ("dinner", "index_10k")]:
+        data_roots.append("/home/weiyu/data_drive/data_new_objects/examples_{}_new_objects/result".format(shape))
+        index_roots.append(index)
+
+    dataset = SemanticArrangementDataset(data_roots=data_roots,
+                                         index_roots=index_roots,
+                                         split="valid", tokenizer=tokenizer,
+                                         max_num_target_objects=7,
+                                         max_num_distractor_objects=5,
+                                         max_num_shape_parameters=1,
+                                         max_num_rearrange_features=0,
+                                         max_num_anchor_features=0,
+                                         num_pts=1024,
+                                         use_virtual_structure_frame=True,
+                                         ignore_distractor_objects=True,
+                                         ignore_rgb=True,
+                                         filter_num_moved_objects_range=None,  # [5, 5]
+                                         data_augmentation=False,
+                                         shuffle_object_index=True,
+                                         sentence_embedding_file="/home/weiyu/Research/StructDiffusion/old/StructDiffusion/src/StructDiffusion/language/template_sentence_data.pkl",
+                                         use_incomplete_sentence=True,
+                                         debug=False)
+
+    # print(len(dataset))
+    # for d in dataset:
+    #     print("\n\n" + "="*100)
+
+    dataloader = DataLoader(dataset, batch_size=64, shuffle=False, num_workers=8)
+    for i, d in enumerate(tqdm(dataloader)):
+        for k in d:
+            if isinstance(d[k], torch.Tensor):
+                print("--size", k, d[k].shape)
+        for k in d:
+            print(k, d[k])
+
+        input("next?")

src/StructDiffusion/data/semantic_arrangement_language_demo.py

@@ -0,0 +1,693 @@
+import copy
+import cv2
+import h5py
+import numpy as np
+import os
+import trimesh
+import torch
+from tqdm import tqdm
+import json
+import random
+import pickle
+
+from torch.utils.data import DataLoader
+
+# Local imports
+from StructDiffusion.utils.rearrangement import show_pcs, get_pts, combine_and_sample_xyzs
+from StructDiffusion.language.tokenizer import Tokenizer
+
+import StructDiffusion.utils.brain2.camera as cam
+import StructDiffusion.utils.brain2.image as img
+import StructDiffusion.utils.transformations as tra
+
+
+class SemanticArrangementDataset(torch.utils.data.Dataset):
+
+    def __init__(self, data_root, tokenizer,
+                 max_num_target_objects=11, max_num_distractor_objects=5,
+                 max_num_shape_parameters=7, max_num_rearrange_features=1, max_num_anchor_features=3,
+                 num_pts=1024,
+                 use_virtual_structure_frame=True, ignore_distractor_objects=True, ignore_rgb=True,
+                 filter_num_moved_objects_range=None, shuffle_object_index=False,
+                 sentence_embedding_file=None, use_incomplete_sentence=False,
+                 data_augmentation=True, debug=False, **kwargs):
+        """
+
+        Note: setting filter_num_moved_objects_range=[k, k] and max_num_objects=k will create no padding for target objs
+
+        :param data_root:
+        :param split: train, valid, or test
+        :param shuffle_object_index: whether to shuffle the positions of target objects and other objects in the sequence
+        :param debug:
+        :param max_num_shape_parameters:
+        :param max_num_objects:
+        :param max_num_rearrange_features:
+        :param max_num_anchor_features:
+        :param num_pts:
+        :param use_stored_arrangement_indices:
+        :param kwargs:
+        """
+
+        self.use_virtual_structure_frame = use_virtual_structure_frame
+        self.ignore_distractor_objects = ignore_distractor_objects
+        self.ignore_rgb = ignore_rgb and not debug
+
+        self.num_pts = num_pts
+        self.debug = debug
+
+        self.max_num_objects = max_num_target_objects
+        self.max_num_other_objects = max_num_distractor_objects
+        self.max_num_shape_parameters = max_num_shape_parameters
+        self.max_num_rearrange_features = max_num_rearrange_features
+        self.max_num_anchor_features = max_num_anchor_features
+        self.shuffle_object_index = shuffle_object_index
+
+        # used to tokenize the language part
+        self.tokenizer = tokenizer
+
+        # retrieve data
+        self.data_root = data_root
+        self.arrangement_data = []
+        for filename in os.listdir(data_root):
+            if ".h5" in filename:
+                self.arrangement_data.append((os.path.join(data_root, filename), 0))
+        print("{} valid sequences".format(len(self.arrangement_data)))
+
+        # language
+        if sentence_embedding_file:
+            assert max_num_shape_parameters == 1
+            # since we do not use them right now, ignore them
+            # assert max_num_rearrange_features == 0
+            # assert max_num_anchor_features == 0
+            with open(sentence_embedding_file, "rb") as fh:
+                template_sentence_data = pickle.load(fh)
+            self.use_sentence_embedding = True
+            self.type_value_tuple_to_template_sentences = template_sentence_data["type_value_tuple_to_template_sentences"]
+            self.template_sentence_to_embedding = template_sentence_data["template_sentence_to_embedding"]
+            self.use_incomplete_sentence = use_incomplete_sentence
+            print("use sentence embedding")
+            print(len(self.type_value_tuple_to_template_sentences))
+            print(len(self.template_sentence_to_embedding))
+        else:
+            self.use_sentence_embedding = False
+
+        # Data Aug
+        self.data_augmentation = data_augmentation
+        # additive noise
+        self.gp_rescale_factor_range = [12, 20]
+        self.gaussian_scale_range = [0., 0.003]
+        # multiplicative noise
+        self.gamma_shape = 1000.
+        self.gamma_scale = 0.001
+
+    def filter_based_on_number_of_moved_objects(self, filter_num_moved_objects_range):
+        assert len(list(filter_num_moved_objects_range)) == 2
+        min_num, max_num = filter_num_moved_objects_range
+        print("Remove scenes that have less than {} or more than {} objects being moved".format(min_num, max_num))
+        ok_data = []
+        for filename, step_t in self.arrangement_data:
+            h5 = h5py.File(filename, 'r')
+            moved_objs = h5['moved_objs'][()].split(',')
+            if min_num <= len(moved_objs) <= max_num:
+                ok_data.append((filename, step_t))
+        print("{} valid sequences left".format(len(ok_data)))
+        return ok_data
+
+    def get_data_idx(self, idx):
+        # Create the datum to return
+        file_idx = np.argmax(idx < self.file_to_count)
+        data = h5py.File(self.data_files[file_idx], 'r')
+        if file_idx > 0:
+            # for lang2sym, idx is always 0
+            idx = idx - self.file_to_count[file_idx - 1]
+        return data, idx, file_idx
+
+    def add_noise_to_depth(self, depth_img):
+        """ add depth noise """
+        multiplicative_noise = np.random.gamma(self.gamma_shape, self.gamma_scale)
+        depth_img = multiplicative_noise * depth_img
+        return depth_img
+
+    def add_noise_to_xyz(self, xyz_img, depth_img):
+        """ TODO: remove this code or at least celean it up"""
+        xyz_img = xyz_img.copy()
+        H, W, C = xyz_img.shape
+        gp_rescale_factor = np.random.randint(self.gp_rescale_factor_range[0],
+                                              self.gp_rescale_factor_range[1])
+        gp_scale = np.random.uniform(self.gaussian_scale_range[0],
+                                     self.gaussian_scale_range[1])
+        small_H, small_W = (np.array([H, W]) / gp_rescale_factor).astype(int)
+        additive_noise = np.random.normal(loc=0.0, scale=gp_scale, size=(small_H, small_W, C))
+        additive_noise = cv2.resize(additive_noise, (W, H), interpolation=cv2.INTER_CUBIC)
+        xyz_img[depth_img > 0, :] += additive_noise[depth_img > 0, :]
+        return xyz_img
+
+    def random_index(self):
+        return self[np.random.randint(len(self))]
+
+    def _get_rgb(self, h5, idx, ee=True):
+        RGB = "ee_rgb" if ee else "rgb"
+        rgb1 = img.PNGToNumpy(h5[RGB][idx])[:, :, :3] / 255.  # remove alpha
+        return rgb1
+
+    def _get_depth(self, h5, idx, ee=True):
+        DEPTH = "ee_depth" if ee else "depth"
+
+    def _get_images(self, h5, idx, ee=True):
+        if ee:
+            RGB, DEPTH, SEG = "ee_rgb", "ee_depth", "ee_seg"
+            DMIN, DMAX = "ee_depth_min", "ee_depth_max"
+        else:
+            RGB, DEPTH, SEG = "rgb", "depth", "seg"
+            DMIN, DMAX = "depth_min", "depth_max"
+        dmin = h5[DMIN][idx]
+        dmax = h5[DMAX][idx]
+        rgb1 = img.PNGToNumpy(h5[RGB][idx])[:, :, :3] / 255.  # remove alpha
+        depth1 = h5[DEPTH][idx] / 20000. * (dmax - dmin) + dmin
+        seg1 = img.PNGToNumpy(h5[SEG][idx])
+
+        valid1 = np.logical_and(depth1 > 0.1, depth1 < 2.)
+
+        # proj_matrix = h5['proj_matrix'][()]
+        camera = cam.get_camera_from_h5(h5)
+        if self.data_augmentation:
+            depth1 = self.add_noise_to_depth(depth1)
+
+        xyz1 = cam.compute_xyz(depth1, camera)
+        if self.data_augmentation:
+            xyz1 = self.add_noise_to_xyz(xyz1, depth1)
+
+        # Transform the point cloud
+        # Here it is...
+        # CAM_POSE = "ee_cam_pose" if ee else "cam_pose"
+        CAM_POSE = "ee_camera_view" if ee else "camera_view"
+        cam_pose = h5[CAM_POSE][idx]
+        if ee:
+            # ee_camera_view has 0s for x, y, z
+            cam_pos = h5["ee_cam_pose"][:][:3, 3]
+            cam_pose[:3, 3] = cam_pos
+
+        # Get transformed point cloud
+        h, w, d = xyz1.shape
+        xyz1 = xyz1.reshape(h * w, -1)
+        xyz1 = trimesh.transform_points(xyz1, cam_pose)
+        xyz1 = xyz1.reshape(h, w, -1)
+
+        scene1 = rgb1, depth1, seg1, valid1, xyz1
+
+        return scene1
+
+    def __len__(self):
+        return len(self.arrangement_data)
+
+    def _get_ids(self, h5):
+        """
+        get object ids
+
+        @param h5:
+        @return:
+        """
+        ids = {}
+        for k in h5.keys():
+            if k.startswith("id_"):
+                ids[k[3:]] = h5[k][()]
+        return ids
+
+    def get_positive_ratio(self):
+        num_pos = 0
+        for d in self.arrangement_data:
+            filename, step_t = d
+            if step_t == 0:
+                num_pos += 1
+        return (len(self.arrangement_data) - num_pos) * 1.0 / num_pos
+
+    def get_object_position_vocab_sizes(self):
+        return self.tokenizer.get_object_position_vocab_sizes()
+
+    def get_vocab_size(self):
+        return self.tokenizer.get_vocab_size()
+
+    def get_data_index(self, idx):
+        filename = self.arrangement_data[idx]
+        return filename
+
+    def get_raw_data(self, idx, inference_mode=False, shuffle_object_index=False):
+        """
+
+        :param idx:
+        :param inference_mode:
+        :param shuffle_object_index: used to test different orders of objects
+        :return:
+        """
+
+        filename, _ = self.arrangement_data[idx]
+
+        h5 = h5py.File(filename, 'r')
+        ids = self._get_ids(h5)
+        all_objs = sorted([o for o in ids.keys() if "object_" in o])
+        goal_specification = json.loads(str(np.array(h5["goal_specification"])))
+        num_rearrange_objs = len(goal_specification["rearrange"]["objects"])
+        num_other_objs = len(goal_specification["anchor"]["objects"] + goal_specification["distract"]["objects"])
+        assert len(all_objs) == num_rearrange_objs + num_other_objs, "{}, {}".format(len(all_objs), num_rearrange_objs + num_other_objs)
+        assert num_rearrange_objs <= self.max_num_objects
+        assert num_other_objs <= self.max_num_other_objects
+
+        # important: only using the last step
+        step_t = num_rearrange_objs
+
+        target_objs = all_objs[:num_rearrange_objs]
+        other_objs = all_objs[num_rearrange_objs:]
+
+        structure_parameters = goal_specification["shape"]
+
+        # Important: ensure the order is correct
+        if structure_parameters["type"] == "circle" or structure_parameters["type"] == "line":
+            target_objs = target_objs[::-1]
+        elif structure_parameters["type"] == "tower" or structure_parameters["type"] == "dinner":
+            target_objs = target_objs
+        else:
+            raise KeyError("{} structure is not recognized".format(structure_parameters["type"]))
+        all_objs = target_objs + other_objs
+
+        ###################################
+        # getting scene images and point clouds
+        scene = self._get_images(h5, step_t, ee=True)
+        rgb, depth, seg, valid, xyz = scene
+        if inference_mode:
+            initial_scene = scene
+
+        # getting object point clouds
+        obj_pcs = []
+        obj_pad_mask = []
+        current_pc_poses = []
+        other_obj_pcs = []
+        other_obj_pad_mask = []
+        for obj in all_objs:
+            obj_mask = np.logical_and(seg == ids[obj], valid)
+            if np.sum(obj_mask) <= 0:
+                raise Exception
+            ok, obj_xyz, obj_rgb, _ = get_pts(xyz, rgb, obj_mask, num_pts=self.num_pts)
+            if not ok:
+                raise Exception
+
+            if obj in target_objs:
+                if self.ignore_rgb:
+                    obj_pcs.append(obj_xyz)
+                else:
+                    obj_pcs.append(torch.concat([obj_xyz, obj_rgb], dim=-1))
+                obj_pad_mask.append(0)
+                pc_pose = np.eye(4)
+                pc_pose[:3, 3] = torch.mean(obj_xyz, dim=0).numpy()
+                current_pc_poses.append(pc_pose)
+            elif obj in other_objs:
+                if self.ignore_rgb:
+                    other_obj_pcs.append(obj_xyz)
+                else:
+                    other_obj_pcs.append(torch.concat([obj_xyz, obj_rgb], dim=-1))
+                other_obj_pad_mask.append(0)
+            else:
+                raise Exception
+
+        ###################################
+        # computes goal positions for objects
+        # Important: because of the noises we added to point clouds, the rearranged point clouds will not be perfect
+        if self.use_virtual_structure_frame:
+            goal_structure_pose = tra.euler_matrix(structure_parameters["rotation"][0], structure_parameters["rotation"][1],
+                                              structure_parameters["rotation"][2])
+            goal_structure_pose[:3, 3] = [structure_parameters["position"][0], structure_parameters["position"][1],
+                                     structure_parameters["position"][2]]
+            goal_structure_pose_inv = np.linalg.inv(goal_structure_pose)
+
+        goal_obj_poses = []
+        current_obj_poses = []
+        goal_pc_poses = []
+        for obj, current_pc_pose in zip(target_objs, current_pc_poses):
+            goal_pose = h5[obj][0]
+            current_pose = h5[obj][step_t]
+            if inference_mode:
+                goal_obj_poses.append(goal_pose)
+                current_obj_poses.append(current_pose)
+
+            goal_pc_pose = goal_pose @ np.linalg.inv(current_pose) @ current_pc_pose
+            if self.use_virtual_structure_frame:
+                goal_pc_pose = goal_structure_pose_inv @ goal_pc_pose
+            goal_pc_poses.append(goal_pc_pose)
+
+        # transform current object point cloud to the goal point cloud in the world frame
+        if self.debug:
+            new_obj_pcs = [copy.deepcopy(pc.numpy()) for pc in obj_pcs]
+            for i, obj_pc in enumerate(new_obj_pcs):
+
+                current_pc_pose = current_pc_poses[i]
+                goal_pc_pose = goal_pc_poses[i]
+                if self.use_virtual_structure_frame:
+                    goal_pc_pose = goal_structure_pose @ goal_pc_pose
+                print("current pc pose", current_pc_pose)
+                print("goal pc pose", goal_pc_pose)
+
+                goal_pc_transform = goal_pc_pose @ np.linalg.inv(current_pc_pose)
+                print("transform", goal_pc_transform)
+                new_obj_pc = copy.deepcopy(obj_pc)
+                new_obj_pc[:, :3] = trimesh.transform_points(obj_pc[:, :3], goal_pc_transform)
+                print(new_obj_pc.shape)
+
+                # visualize rearrangement sequence (new_obj_xyzs), the current object before moving (obj_xyz), and other objects
+                new_obj_pcs[i] = new_obj_pc
+                new_obj_pcs[i][:, 3:] = np.tile(np.array([1, 0, 0], dtype=np.float), (new_obj_pc.shape[0], 1))
+                new_obj_rgb_current = np.tile(np.array([0, 1, 0], dtype=np.float), (new_obj_pc.shape[0], 1))
+                show_pcs([pc[:, :3] for pc in new_obj_pcs] + [pc[:, :3] for pc in other_obj_pcs] + [obj_pc[:, :3]],
+                         [pc[:, 3:] for pc in new_obj_pcs] + [pc[:, 3:] for pc in other_obj_pcs] + [new_obj_rgb_current],
+                         add_coordinate_frame=True)
+            show_pcs([pc[:, :3] for pc in new_obj_pcs], [pc[:, 3:] for pc in new_obj_pcs], add_coordinate_frame=True)
+
+        # pad data
+        for i in range(self.max_num_objects - len(target_objs)):
+            obj_pcs.append(torch.zeros_like(obj_pcs[0], dtype=torch.float32))
+            obj_pad_mask.append(1)
+        for i in range(self.max_num_other_objects - len(other_objs)):
+            other_obj_pcs.append(torch.zeros_like(obj_pcs[0], dtype=torch.float32))
+            other_obj_pad_mask.append(1)
+
+        ###################################
+        # preparing sentence
+        sentence = []
+        sentence_pad_mask = []
+
+        # structure parameters
+        # 5 parameters
+        structure_parameters = goal_specification["shape"]
+        if structure_parameters["type"] == "circle" or structure_parameters["type"] == "line":
+            sentence.append((structure_parameters["type"], "shape"))
+            sentence.append((structure_parameters["rotation"][2], "rotation"))
+            sentence.append((structure_parameters["position"][0], "position_x"))
+            sentence.append((structure_parameters["position"][1], "position_y"))
+            if structure_parameters["type"] == "circle":
+                sentence.append((structure_parameters["radius"], "radius"))
+            elif structure_parameters["type"] == "line":
+                sentence.append((structure_parameters["length"] / 2.0, "radius"))
+            if not self.use_sentence_embedding:
+                for _ in range(5):
+                    sentence_pad_mask.append(0)
+        else:
+            sentence.append((structure_parameters["type"], "shape"))
+            sentence.append((structure_parameters["rotation"][2], "rotation"))
+            sentence.append((structure_parameters["position"][0], "position_x"))
+            sentence.append((structure_parameters["position"][1], "position_y"))
+            if not self.use_sentence_embedding:
+                for _ in range(4):
+                    sentence_pad_mask.append(0)
+                sentence.append(("PAD", None))
+                sentence_pad_mask.append(1)
+
+        if self.use_sentence_embedding:
+
+            if self.use_incomplete_sentence:
+                token_idxs = np.random.permutation(len(sentence))
+                token_idxs = token_idxs[:np.random.randint(1, len(sentence) + 1)]
+                token_idxs = sorted(token_idxs)
+                incomplete_sentence = [sentence[ti] for ti in token_idxs]
+            else:
+                incomplete_sentence = sentence
+
+            type_value_tuple = self.tokenizer.convert_structure_params_to_type_value_tuple(incomplete_sentence)
+            template_sentence = np.random.choice(self.type_value_tuple_to_template_sentences[type_value_tuple])
+            sentence_embedding = self.template_sentence_to_embedding[template_sentence]
+            sentence_pad_mask = [0]
+
+        ###################################
+        # paddings
+        for i in range(self.max_num_objects - len(target_objs)):
+            goal_pc_poses.append(np.eye(4))
+
+        ###################################
+        if self.debug:
+            print("---")
+            print("all objects:", all_objs)
+            print("target objects:", target_objs)
+            print("other objects:", other_objs)
+            print("goal specification:", goal_specification)
+            print("sentence:", sentence)
+            if self.use_sentence_embedding:
+                print("use sentence embedding")
+                if self.use_incomplete_sentence:
+                    print("incomplete_sentence:", incomplete_sentence)
+                print("template sentence:", template_sentence)
+            show_pcs([pc[:, :3] for pc in obj_pcs + other_obj_pcs], [pc[:, 3:] for pc in obj_pcs + other_obj_pcs], add_coordinate_frame=True)
+
+        assert len(obj_pcs) == len(goal_pc_poses)
+        ###################################
+
+        # shuffle the position of objects
+        # important: only shuffle for dinner
+        if shuffle_object_index and structure_parameters["type"] == "dinner":
+            num_target_objs = len(target_objs)
+            shuffle_target_object_indices = list(range(num_target_objs))
+            random.shuffle(shuffle_target_object_indices)
+            shuffle_object_indices = shuffle_target_object_indices + list(range(num_target_objs, self.max_num_objects))
+            obj_pcs = [obj_pcs[i] for i in shuffle_object_indices]
+            goal_pc_poses = [goal_pc_poses[i] for i in shuffle_object_indices]
+            if inference_mode:
+                goal_obj_poses = [goal_obj_poses[i] for i in shuffle_object_indices[:num_target_objs]]
+                current_obj_poses = [current_obj_poses[i] for i in shuffle_object_indices[:num_target_objs]]
+                target_objs = [target_objs[i] for i in shuffle_target_object_indices[:num_target_objs]]
+                current_pc_poses = [current_pc_poses[i] for i in shuffle_object_indices[:num_target_objs]]
+
+        ###################################
+        if self.use_virtual_structure_frame:
+            if self.ignore_distractor_objects:
+                # language, structure virtual frame, target objects
+                pcs = obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [2] + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + [0] + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + [0] + obj_pad_mask
+            else:
+                # language, distractor objects, structure virtual frame, target objects
+                pcs = other_obj_pcs + obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [1] * self.max_num_other_objects + [2] + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + list(range(self.max_num_other_objects)) + [0] + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + other_obj_pad_mask + [0] + obj_pad_mask
+            goal_poses = [goal_structure_pose] + goal_pc_poses
+        else:
+            if self.ignore_distractor_objects:
+                # language, target objects
+                pcs = obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + obj_pad_mask
+            else:
+                # language, distractor objects, target objects
+                pcs = other_obj_pcs + obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [1] * self.max_num_other_objects + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + list(range(self.max_num_other_objects)) + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + other_obj_pad_mask + obj_pad_mask
+            goal_poses = goal_pc_poses
+
+        datum = {
+            "pcs": pcs,
+            "goal_poses": goal_poses,
+            "type_index": type_index,
+            "position_index": position_index,
+            "pad_mask": pad_mask,
+            "t": step_t,
+            "filename": filename
+        }
+        if self.use_sentence_embedding:
+            datum["sentence"] = sentence_embedding
+        else:
+            datum["sentence"] = sentence
+
+        if inference_mode:
+            datum["rgb"] = rgb
+            datum["goal_obj_poses"] = goal_obj_poses
+            datum["current_obj_poses"] = current_obj_poses
+            datum["target_objs"] = target_objs
+            datum["initial_scene"] = initial_scene
+            datum["ids"] = ids
+            datum["goal_specification"] = goal_specification
+            datum["current_pc_poses"] = current_pc_poses
+            if self.use_sentence_embedding:
+                datum["template_sentence"] = template_sentence
+
+        return datum
+
+    def build_data_from_xyzs(self, obj_xyzs, sentence_embedding, shuffle_object_index=True):
+
+        ## objects
+        obj_pcs = []
+        obj_pad_mask = []
+        current_pc_poses = []
+        other_obj_pcs = []
+        other_obj_pad_mask = []
+        for obj_xyz in obj_xyzs:
+            obj_pcs.append(torch.from_numpy(obj_xyz.astype(np.float32)))
+            obj_pad_mask.append(0)
+
+        # pad data
+        num_target_objs = len(obj_pcs)
+        for i in range(self.max_num_objects - num_target_objs):
+            obj_pcs.append(torch.zeros_like(obj_pcs[0], dtype=torch.float32))
+            obj_pad_mask.append(1)
+        for i in range(self.max_num_other_objects):
+            other_obj_pcs.append(torch.zeros_like(obj_pcs[0], dtype=torch.float32))
+            other_obj_pad_mask.append(1)
+
+        ## sentence
+        sentence_pad_mask = [0]
+
+        if shuffle_object_index:
+            num_target_objs = num_target_objs
+            shuffle_target_object_indices = list(range(num_target_objs))
+            random.shuffle(shuffle_target_object_indices)
+            shuffle_object_indices = shuffle_target_object_indices + list(range(num_target_objs, self.max_num_objects))
+            obj_pcs = [obj_pcs[i] for i in shuffle_object_indices]
+
+        ###################################
+        if self.use_virtual_structure_frame:
+            if self.ignore_distractor_objects:
+                # language, structure virtual frame, target objects
+                pcs = obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [2] + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + [0] + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + [0] + obj_pad_mask
+            else:
+                # language, distractor objects, structure virtual frame, target objects
+                pcs = other_obj_pcs + obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [1] * self.max_num_other_objects + [2] + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + list(range(self.max_num_other_objects)) + [0] + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + other_obj_pad_mask + [0] + obj_pad_mask
+            num_goal_poses = self.max_num_objects + 1
+        else:
+            if self.ignore_distractor_objects:
+                # language, target objects
+                pcs = obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + obj_pad_mask
+            else:
+                # language, distractor objects, target objects
+                pcs = other_obj_pcs + obj_pcs
+                type_index = [0] * self.max_num_shape_parameters + [1] * self.max_num_other_objects + [3] * self.max_num_objects
+                position_index = list(range(self.max_num_shape_parameters)) + list(range(self.max_num_other_objects)) + list(range(self.max_num_objects))
+                pad_mask = sentence_pad_mask + other_obj_pad_mask + obj_pad_mask
+            num_goal_poses = self.max_num_objects
+
+        datum = {
+            "pcs": pcs,
+            "type_index": type_index,
+            "position_index": position_index,
+            "pad_mask": pad_mask,
+            "sentence": sentence_embedding,
+            "num_goal_poses": num_goal_poses,
+            "t": 0,
+            "filename": "inference"
+        }
+
+        return datum
+
+    @staticmethod
+    def convert_to_tensors(datum, tokenizer, use_sentence_embedding=False):
+        tensors = {
+            "pcs": torch.stack(datum["pcs"], dim=0),
+            "type_index": torch.LongTensor(np.array(datum["type_index"])),
+            "position_index": torch.LongTensor(np.array(datum["position_index"])),
+            "pad_mask": torch.LongTensor(np.array(datum["pad_mask"])),
+            "t": datum["t"],
+            "filename": datum["filename"]
+        }
+        if "goal_poses" in datum:
+            tensors["goal_poses"] = torch.FloatTensor(np.array(datum["goal_poses"])),
+
+        if use_sentence_embedding:
+            tensors["sentence"] = torch.FloatTensor(datum["sentence"])  # after batching, B x sentence embed dim
+        else:
+            tensors["sentence"] = torch.LongTensor(np.array([tokenizer.tokenize(*i) for i in datum["sentence"]]))
+        return tensors
+
+    def __getitem__(self, idx):
+
+        datum = self.convert_to_tensors(self.get_raw_data(idx, shuffle_object_index=self.shuffle_object_index),
+                                        self.tokenizer,
+                                        self.use_sentence_embedding)
+
+        return datum
+
+    def single_datum_to_batch(self, x, num_samples, device, inference_mode=True):
+        tensor_x = {}
+
+        tensor_x["pcs"] = x["pcs"].to(device)[None, :, :, :].repeat(num_samples, 1, 1, 1)
+        tensor_x["sentence"] = x["sentence"].to(device)[None, :].repeat(num_samples, 1)
+        if not inference_mode:
+            tensor_x["goal_poses"] = x["goal_poses"].to(device)[None, :, :, :].repeat(num_samples, 1, 1, 1)
+
+        tensor_x["type_index"] = x["type_index"].to(device)[None, :].repeat(num_samples, 1)
+        tensor_x["position_index"] = x["position_index"].to(device)[None, :].repeat(num_samples, 1)
+        tensor_x["pad_mask"] = x["pad_mask"].to(device)[None, :].repeat(num_samples, 1)
+
+        return tensor_x
+
+
+def compute_min_max(dataloader):
+
+    # tensor([-0.3557, -0.3847,  0.0000, -1.0000, -1.0000, -0.4759, -1.0000, -1.0000,
+    #         -0.9079, -0.8668, -0.9105, -0.4186])
+    # tensor([0.3915, 0.3494, 0.3267, 1.0000, 1.0000, 0.8961, 1.0000, 1.0000, 0.8194,
+    #         0.4787, 0.6421, 1.0000])
+    # tensor([0.0918, -0.3758, 0.0000, -1.0000, -1.0000, 0.0000, -1.0000, -1.0000,
+    #         -0.0000, 0.0000, 0.0000, 1.0000])
+    # tensor([0.9199, 0.3710, 0.0000, 1.0000, 1.0000, 0.0000, 1.0000, 1.0000, -0.0000,
+    #         0.0000, 0.0000, 1.0000])
+
+    min_value = torch.ones(16) * 10000
+    max_value = torch.ones(16) * -10000
+    for d in tqdm(dataloader):
+        goal_poses = d["goal_poses"]
+        goal_poses = goal_poses.reshape(-1, 16)
+        current_max, _ = torch.max(goal_poses, dim=0)
+        current_min, _ = torch.min(goal_poses, dim=0)
+        max_value[max_value < current_max] = current_max[max_value < current_max]
+        max_value[max_value > current_min] = current_min[max_value > current_min]
+    print(f"{min_value} - {max_value}")
+
+
+if __name__ == "__main__":
+
+    tokenizer = Tokenizer("/home/weiyu/data_drive/data_new_objects/type_vocabs_coarse.json")
+
+    data_roots = []
+    index_roots = []
+    for shape, index in [("circle", "index_10k"), ("line", "index_10k"), ("stacking", "index_10k"), ("dinner", "index_10k")]:
+        data_roots.append("/home/weiyu/data_drive/data_new_objects/examples_{}_new_objects/result".format(shape))
+        index_roots.append(index)
+
+    dataset = SemanticArrangementDataset(data_roots=data_roots,
+                                         index_roots=index_roots,
+                                         split="valid", tokenizer=tokenizer,
+                                         max_num_target_objects=7,
+                                         max_num_distractor_objects=5,
+                                         max_num_shape_parameters=1,
+                                         max_num_rearrange_features=0,
+                                         max_num_anchor_features=0,
+                                         num_pts=1024,
+                                         use_virtual_structure_frame=True,
+                                         ignore_distractor_objects=True,
+                                         ignore_rgb=True,
+                                         filter_num_moved_objects_range=None,  # [5, 5]
+                                         data_augmentation=False,
+                                         shuffle_object_index=True,
+                                         sentence_embedding_file="/home/weiyu/Research/StructDiffusion/old/StructDiffusion/src/StructDiffusion/language/template_sentence_data.pkl",
+                                         use_incomplete_sentence=True,
+                                         debug=False)
+
+    # print(len(dataset))
+    # for d in dataset:
+    #     print("\n\n" + "="*100)
+
+    dataloader = DataLoader(dataset, batch_size=64, shuffle=False, num_workers=8)
+    for i, d in enumerate(tqdm(dataloader)):
+        for k in d:
+            if isinstance(d[k], torch.Tensor):
+                print("--size", k, d[k].shape)
+        for k in d:
+            print(k, d[k])
+
+        input("next?")

src/StructDiffusion/diffusion/sampler.py

@@ -1,6 +1,11 @@
 import torch
 from tqdm import tqdm
+
 from StructDiffusion.diffusion.noise_schedule import extract
+from StructDiffusion.diffusion.pose_conversion import get_struct_objs_poses
+from StructDiffusion.utils.batch_inference import move_pc_and_create_scene_new
+import StructDiffusion.utils.tra3d as tra3d
+
 
 class Sampler:
 
@@ -14,7 +19,7 @@ class Sampler:
         self.backbone.to(device)
         self.backbone.eval()
 
-    def sample(self, batch, num_poses, progress):
+    def sample(self, batch, num_poses):
 
         noise_schedule = self.model.noise_schedule
 
@@ -23,7 +28,7 @@ class Sampler:
         x_noisy = torch.randn((B, num_poses, 9), device=self.device)
 
         xs = []
-        for t_index in progress.tqdm(reversed(range(0, noise_schedule.timesteps)),
+        for t_index in tqdm(reversed(range(0, noise_schedule.timesteps)),
                             desc='sampling loop time step', total=noise_schedule.timesteps):
 
             t = torch.full((B,), t_index, device=self.device, dtype=torch.long)
@@ -57,236 +62,239 @@ class Sampler:
         xs = list(reversed(xs))
         return xs
 
-# class SamplerV2:
-#
-#     def __init__(self, diffusion_model_class, diffusion_checkpoint_path,
-#                  collision_model_class, collision_checkpoint_path,
-#                  device, debug=False):
-#
-#         self.debug = debug
-#         self.device = device
-#
-#         self.diffusion_model = diffusion_model_class.load_from_checkpoint(diffusion_checkpoint_path)
-#         self.diffusion_backbone = self.diffusion_model.model
-#         self.diffusion_backbone.to(device)
-#         self.diffusion_backbone.eval()
-#
-#         self.collision_model = collision_model_class.load_from_checkpoint(collision_checkpoint_path)
-#         self.collision_backbone = self.collision_model.model
-#         self.collision_backbone.to(device)
-#         self.collision_backbone.eval()
-#
-#     def sample(self, batch, num_poses):
-#
-#         noise_schedule = self.diffusion_model.noise_schedule
-#
-#         B = batch["pcs"].shape[0]
-#
-#         x_noisy = torch.randn((B, num_poses, 9), device=self.device)
-#
-#         xs = []
-#         for t_index in tqdm(reversed(range(0, noise_schedule.timesteps)),
-#                             desc='sampling loop time step', total=noise_schedule.timesteps):
-#
-#             t = torch.full((B,), t_index, device=self.device, dtype=torch.long)
-#
-#             # noise schedule
-#             betas_t = extract(noise_schedule.betas, t, x_noisy.shape)
-#             sqrt_one_minus_alphas_cumprod_t = extract(noise_schedule.sqrt_one_minus_alphas_cumprod, t, x_noisy.shape)
-#             sqrt_recip_alphas_t = extract(noise_schedule.sqrt_recip_alphas, t, x_noisy.shape)
-#
-#             # predict noise
-#             pcs = batch["pcs"]
-#             sentence = batch["sentence"]
-#             type_index = batch["type_index"]
-#             position_index = batch["position_index"]
-#             pad_mask = batch["pad_mask"]
-#             # calling the backbone instead of the pytorch-lightning model
-#             with torch.no_grad():
-#                 predicted_noise = self.diffusion_backbone.forward(t, pcs, sentence, x_noisy, type_index, position_index, pad_mask)
-#
-#             # compute noisy x at t
-#             model_mean = sqrt_recip_alphas_t * (x_noisy - betas_t * predicted_noise / sqrt_one_minus_alphas_cumprod_t)
-#             if t_index == 0:
-#                 x_noisy = model_mean
-#             else:
-#                 posterior_variance_t = extract(noise_schedule.posterior_variance, t, x_noisy.shape)
-#                 noise = torch.randn_like(x_noisy)
-#                 x_noisy = model_mean + torch.sqrt(posterior_variance_t) * noise
-#
-#             xs.append(x_noisy)
-#
-#         xs = list(reversed(xs))
-#
-#         visualize = True
-#
-#         struct_pose, pc_poses_in_struct = get_struct_objs_poses(xs[0])
-#         # struct_pose: B, 1, 4, 4
-#         # pc_poses_in_struct: B, N, 4, 4
-#
-#         S = B
-#         num_elite = 10
-#         ####################################################
-#         # only keep one copy
-#
-#         # N, P, 3
-#         obj_xyzs = batch["pcs"][0][:, :, :3]
-#         print("obj_xyzs shape", obj_xyzs.shape)
-#
-#         # 1, N
-#         # object_pad_mask: padding location has 1
-#         num_target_objs = num_poses
-#         if self.diffusion_backbone.use_virtual_structure_frame:
-#             num_target_objs -= 1
-#         object_pad_mask = batch["pad_mask"][0][-num_target_objs:].unsqueeze(0)
-#         target_object_inds = 1 - object_pad_mask
-#         print("target_object_inds shape", target_object_inds.shape)
-#         print("target_object_inds", target_object_inds)
-#
-#         N, P, _ = obj_xyzs.shape
-#         print("S, N, P: {}, {}, {}".format(S, N, P))
-#
-#         ####################################################
-#         # S, N, ...
-#
-#         struct_pose = struct_pose.repeat(1, N, 1, 1)  # S, N, 4, 4
-#         struct_pose = struct_pose.reshape(S * N, 4, 4)  # S x N, 4, 4
-#
-#         new_obj_xyzs = obj_xyzs.repeat(S, 1, 1, 1)  # S, N, P, 3
-#         current_pc_pose = torch.eye(4).repeat(S, N, 1, 1).to(self.device)  # S, N, 4, 4
-#         current_pc_pose[:, :, :3, 3] = torch.mean(new_obj_xyzs, dim=2)  # S, N, 4, 4
-#         current_pc_pose = current_pc_pose.reshape(S * N, 4, 4)  # S x N, 4, 4
-#
-#         # optimize xyzrpy
-#         obj_params = torch.zeros((S, N, 6)).to(self.device)
-#         obj_params[:, :, :3] = pc_poses_in_struct[:, :, :3, 3]
-#         obj_params[:, :, 3:] = tra3d.matrix_to_euler_angles(pc_poses_in_struct[:, :, :3, :3], "XYZ")  # S, N, 6
-#         #
-#         # new_obj_xyzs_before_cem, goal_pc_pose_before_cem = move_pc(obj_xyzs, obj_params, struct_pose, current_pc_pose, device)
-#         #
-#         # if visualize:
-#         #     print("visualizing rearrangements predicted by the generator")
-#         #     visualize_batch_pcs(new_obj_xyzs_before_cem, S, N, P, limit_B=5)
-#
-#         ####################################################
-#         # rank
-#
-#         # evaluate in batches
-#         scores = torch.zeros(S).to(self.device)
-#         no_intersection_scores = torch.zeros(S).to(self.device)  # the higher the better
-#         num_batches = int(S / B)
-#         if S % B != 0:
-#             num_batches += 1
-#         for b in range(num_batches):
-#             if b + 1 == num_batches:
-#                 cur_batch_idxs_start = b * B
-#                 cur_batch_idxs_end = S
-#             else:
-#                 cur_batch_idxs_start = b * B
-#                 cur_batch_idxs_end = (b + 1) * B
-#             cur_batch_size = cur_batch_idxs_end - cur_batch_idxs_start
-#
-#             # print("current batch idxs start", cur_batch_idxs_start)
-#             # print("current batch idxs end", cur_batch_idxs_end)
-#             # print("size of the current batch", cur_batch_size)
-#
-#             batch_obj_params = obj_params[cur_batch_idxs_start: cur_batch_idxs_end]
-#             batch_struct_pose = struct_pose[cur_batch_idxs_start * N: cur_batch_idxs_end * N]
-#             batch_current_pc_pose = current_pc_pose[cur_batch_idxs_start * N:cur_batch_idxs_end * N]
-#
-#             new_obj_xyzs, _, subsampled_scene_xyz, _, obj_pair_xyzs = \
-#                 move_pc_and_create_scene_new(obj_xyzs, batch_obj_params, batch_struct_pose, batch_current_pc_pose,
-#                                              target_object_inds, self.device,
-#                                              return_scene_pts=False,
-#                                              return_scene_pts_and_pc_idxs=False,
-#                                              num_scene_pts=False,
-#                                              normalize_pc=False,
-#                                              return_pair_pc=True,
-#                                              num_pair_pc_pts=self.collision_model.data_cfg.num_scene_pts,
-#                                              normalize_pair_pc=self.collision_model.data_cfg.normalize_pc)
-#
-#             #######################################
-#             # predict whether there are pairwise collisions
-#             # if collision_score_weight > 0:
-#             with torch.no_grad():
-#                 _, num_comb, num_pair_pc_pts, _ = obj_pair_xyzs.shape
-#                 # obj_pair_xyzs = obj_pair_xyzs.reshape(cur_batch_size * num_comb, num_pair_pc_pts, -1)
-#                 collision_logits = self.collision_backbone.forward(obj_pair_xyzs.reshape(cur_batch_size * num_comb, num_pair_pc_pts, -1))
-#                 collision_scores = self.collision_backbone.convert_logits(collision_logits).reshape(cur_batch_size, num_comb)  # cur_batch_size, num_comb
-#
-#                 # debug
-#                 # for bi, this_obj_pair_xyzs in enumerate(obj_pair_xyzs):
-#                 #     print("batch id", bi)
-#                 #     for pi, obj_pair_xyz in enumerate(this_obj_pair_xyzs):
-#                 #         print("pair", pi)
-#                 #         # obj_pair_xyzs: 2 * P, 5
-#                 #         print("collision score", collision_scores[bi, pi])
-#                 #         trimesh.PointCloud(obj_pair_xyz[:, :3].cpu()).show()
-#
-#                 # 1 - mean() since the collision model predicts 1 if there is a collision
-#                 no_intersection_scores[cur_batch_idxs_start:cur_batch_idxs_end] = 1 - torch.mean(collision_scores, dim=1)
-#             if visualize:
-#                 print("no intersection scores", no_intersection_scores)
-#             # #######################################
-#             # if discriminator_score_weight > 0:
-#             #     # # debug:
-#             #     # print(subsampled_scene_xyz.shape)
-#             #     # print(subsampled_scene_xyz[0])
-#             #     # trimesh.PointCloud(subsampled_scene_xyz[0, :, :3].cpu().numpy()).show()
-#             #     #
-#             #     with torch.no_grad():
-#             #
-#             #         # Important: since this discriminator only uses local structure param, takes sentence from the first and last position
-#             #         # local_sentence = sentence[:, [0, 4]]
-#             #         # local_sentence_pad_mask = sentence_pad_mask[:, [0, 4]]
-#             #         # sentence_disc, sentence_pad_mask_disc, position_index_dic = discriminator_inference.dataset.tensorfy_sentence(raw_sentence_discriminator, raw_sentence_pad_mask_discriminator, raw_position_index_discriminator)
-#             #
-#             #         sentence_disc = torch.LongTensor(
-#             #             [discriminator_tokenizer.tokenize(*i) for i in raw_sentence_discriminator])
-#             #         sentence_pad_mask_disc = torch.LongTensor(raw_sentence_pad_mask_discriminator)
-#             #         position_index_dic = torch.LongTensor(raw_position_index_discriminator)
-#             #
-#             #         preds = discriminator_model.forward(subsampled_scene_xyz,
-#             #                                             sentence_disc.unsqueeze(0).repeat(cur_batch_size, 1).to(device),
-#             #                                             sentence_pad_mask_disc.unsqueeze(0).repeat(cur_batch_size,
-#             #                                                                                        1).to(device),
-#             #                                             position_index_dic.unsqueeze(0).repeat(cur_batch_size, 1).to(
-#             #                                                 device))
-#             #         # preds = discriminator_model.forward(subsampled_scene_xyz)
-#             #         preds = discriminator_model.convert_logits(preds)
-#             #         preds = preds["is_circle"]  # cur_batch_size,
-#             #         scores[cur_batch_idxs_start:cur_batch_idxs_end] = preds
-#             #     if visualize:
-#             #         print("discriminator scores", scores)
-#
-#         # scores = scores * discriminator_score_weight + no_intersection_scores * collision_score_weight
-#         scores = no_intersection_scores
-#         sort_idx = torch.argsort(scores).flip(dims=[0])[:num_elite]
-#         elite_obj_params = obj_params[sort_idx]  # num_elite, N, 6
-#         elite_struct_poses = struct_pose.reshape(S, N, 4, 4)[sort_idx]  # num_elite, N, 4, 4
-#         elite_struct_poses = elite_struct_poses.reshape(num_elite * N, 4, 4)  # num_elite x N, 4, 4
-#         elite_scores = scores[sort_idx]
-#         print("elite scores:", elite_scores)
-#
-#         ####################################################
-#         # # visualize best samples
-#         # num_scene_pts = 4096 # if discriminator_num_scene_pts is None else discriminator_num_scene_pts
-#         # batch_current_pc_pose = current_pc_pose[0: num_elite * N]
-#         # best_new_obj_xyzs, best_goal_pc_pose, best_subsampled_scene_xyz, _, _ = \
-#         #     move_pc_and_create_scene_new(obj_xyzs, elite_obj_params, elite_struct_poses, batch_current_pc_pose,
-#         #                                  target_object_inds, self.device,
-#         #                                  return_scene_pts=True, num_scene_pts=num_scene_pts, normalize_pc=True)
-#         # if visualize:
-#         #     print("visualizing elite rearrangements ranked by collision model/discriminator")
-#         #     visualize_batch_pcs(best_new_obj_xyzs, num_elite, limit_B=num_elite)
-#
-#         # num_elite, N, 6
-#         elite_obj_params = elite_obj_params.reshape(num_elite * N, -1)
-#         pc_poses_in_struct = torch.eye(4).repeat(num_elite * N, 1, 1).to(self.device)
-#         pc_poses_in_struct[:, :3, :3] = tra3d.euler_angles_to_matrix(elite_obj_params[:, 3:], "XYZ")
-#         pc_poses_in_struct[:, :3, 3] = elite_obj_params[:, :3]
-#         pc_poses_in_struct = pc_poses_in_struct.reshape(num_elite, N, 4, 4)  # num_elite, N, 4, 4
-#
-#         struct_pose = elite_struct_poses.reshape(num_elite, N, 4, 4)[:, 0,].unsqueeze(1)  # num_elite, 1, 4, 4
-#
-#         return struct_pose, pc_poses_in_struct
+class SamplerV2:
+
+    def __init__(self, diffusion_model_class, diffusion_checkpoint_path,
+                 collision_model_class, collision_checkpoint_path,
+                 device, debug=False):
+
+        self.debug = debug
+        self.device = device
+
+        self.diffusion_model = diffusion_model_class.load_from_checkpoint(diffusion_checkpoint_path)
+        self.diffusion_backbone = self.diffusion_model.model
+        self.diffusion_backbone.to(device)
+        self.diffusion_backbone.eval()
+
+        self.collision_model = collision_model_class.load_from_checkpoint(collision_checkpoint_path)
+        self.collision_backbone = self.collision_model.model
+        self.collision_backbone.to(device)
+        self.collision_backbone.eval()
+
+    def sample(self, batch, num_poses, num_elite, discriminator_batch_size):
+
+        noise_schedule = self.diffusion_model.noise_schedule
+
+        B = batch["pcs"].shape[0]
+
+        x_noisy = torch.randn((B, num_poses, 9), device=self.device)
+
+        xs = []
+        for t_index in tqdm(reversed(range(0, noise_schedule.timesteps)),
+                            desc='sampling loop time step', total=noise_schedule.timesteps):
+
+            t = torch.full((B,), t_index, device=self.device, dtype=torch.long)
+
+            # noise schedule
+            betas_t = extract(noise_schedule.betas, t, x_noisy.shape)
+            sqrt_one_minus_alphas_cumprod_t = extract(noise_schedule.sqrt_one_minus_alphas_cumprod, t, x_noisy.shape)
+            sqrt_recip_alphas_t = extract(noise_schedule.sqrt_recip_alphas, t, x_noisy.shape)
+
+            # predict noise
+            pcs = batch["pcs"]
+            sentence = batch["sentence"]
+            type_index = batch["type_index"]
+            position_index = batch["position_index"]
+            pad_mask = batch["pad_mask"]
+            # calling the backbone instead of the pytorch-lightning model
+            with torch.no_grad():
+                predicted_noise = self.diffusion_backbone.forward(t, pcs, sentence, x_noisy, type_index, position_index, pad_mask)
+
+            # compute noisy x at t
+            model_mean = sqrt_recip_alphas_t * (x_noisy - betas_t * predicted_noise / sqrt_one_minus_alphas_cumprod_t)
+            if t_index == 0:
+                x_noisy = model_mean
+            else:
+                posterior_variance_t = extract(noise_schedule.posterior_variance, t, x_noisy.shape)
+                noise = torch.randn_like(x_noisy)
+                x_noisy = model_mean + torch.sqrt(posterior_variance_t) * noise
+
+            xs.append(x_noisy)
+
+        xs = list(reversed(xs))
+
+        visualize = True
+
+        struct_pose, pc_poses_in_struct = get_struct_objs_poses(xs[0])
+        # struct_pose: B, 1, 4, 4
+        # pc_poses_in_struct: B, N, 4, 4
+
+        S = B
+        B_discriminator = discriminator_batch_size
+        ####################################################
+        # only keep one copy
+
+        # N, P, 3
+        obj_xyzs = batch["pcs"][0][:, :, :3]
+        print("obj_xyzs shape", obj_xyzs.shape)
+
+        # 1, N
+        # object_pad_mask: padding location has 1
+        num_target_objs = num_poses
+        if self.diffusion_backbone.use_virtual_structure_frame:
+            num_target_objs -= 1
+        object_pad_mask = batch["pad_mask"][0][-num_target_objs:].unsqueeze(0)
+        target_object_inds = 1 - object_pad_mask
+        print("target_object_inds shape", target_object_inds.shape)
+        print("target_object_inds", target_object_inds)
+
+        N, P, _ = obj_xyzs.shape
+        print("S, N, P: {}, {}, {}".format(S, N, P))
+
+        ####################################################
+        # S, N, ...
+
+        struct_pose = struct_pose.repeat(1, N, 1, 1)  # S, N, 4, 4
+        struct_pose = struct_pose.reshape(S * N, 4, 4)  # S x N, 4, 4
+
+        new_obj_xyzs = obj_xyzs.repeat(S, 1, 1, 1)  # S, N, P, 3
+        current_pc_pose = torch.eye(4).repeat(S, N, 1, 1).to(self.device)  # S, N, 4, 4
+        current_pc_pose[:, :, :3, 3] = torch.mean(new_obj_xyzs, dim=2)  # S, N, 4, 4
+        current_pc_pose = current_pc_pose.reshape(S * N, 4, 4)  # S x N, 4, 4
+
+        # optimize xyzrpy
+        obj_params = torch.zeros((S, N, 6)).to(self.device)
+        obj_params[:, :, :3] = pc_poses_in_struct[:, :, :3, 3]
+        obj_params[:, :, 3:] = tra3d.matrix_to_euler_angles(pc_poses_in_struct[:, :, :3, :3], "XYZ")  # S, N, 6
+        #
+        # new_obj_xyzs_before_cem, goal_pc_pose_before_cem = move_pc(obj_xyzs, obj_params, struct_pose, current_pc_pose, device)
+        #
+        # if visualize:
+        #     print("visualizing rearrangements predicted by the generator")
+        #     visualize_batch_pcs(new_obj_xyzs_before_cem, S, N, P, limit_B=5)
+
+        ####################################################
+        # rank
+
+        # evaluate in batches
+        scores = torch.zeros(S).to(self.device)
+        no_intersection_scores = torch.zeros(S).to(self.device)  # the higher the better
+        num_batches = int(S / B_discriminator)
+        if S % B_discriminator != 0:
+            num_batches += 1
+        for b in range(num_batches):
+            if b + 1 == num_batches:
+                cur_batch_idxs_start = b * B_discriminator
+                cur_batch_idxs_end = S
+            else:
+                cur_batch_idxs_start = b * B_discriminator
+                cur_batch_idxs_end = (b + 1) * B_discriminator
+            cur_batch_size = cur_batch_idxs_end - cur_batch_idxs_start
+
+            # print("current batch idxs start", cur_batch_idxs_start)
+            # print("current batch idxs end", cur_batch_idxs_end)
+            # print("size of the current batch", cur_batch_size)
+
+            batch_obj_params = obj_params[cur_batch_idxs_start: cur_batch_idxs_end]
+            batch_struct_pose = struct_pose[cur_batch_idxs_start * N: cur_batch_idxs_end * N]
+            batch_current_pc_pose = current_pc_pose[cur_batch_idxs_start * N:cur_batch_idxs_end * N]
+
+            new_obj_xyzs, _, subsampled_scene_xyz, _, obj_pair_xyzs = \
+                move_pc_and_create_scene_new(obj_xyzs, batch_obj_params, batch_struct_pose, batch_current_pc_pose,
+                                             target_object_inds, self.device,
+                                             return_scene_pts=False,
+                                             return_scene_pts_and_pc_idxs=False,
+                                             num_scene_pts=False,
+                                             normalize_pc=False,
+                                             return_pair_pc=True,
+                                             num_pair_pc_pts=self.collision_model.data_cfg.num_scene_pts,
+                                             normalize_pair_pc=self.collision_model.data_cfg.normalize_pc)
+
+            #######################################
+            # predict whether there are pairwise collisions
+            # if collision_score_weight > 0:
+            with torch.no_grad():
+                _, num_comb, num_pair_pc_pts, _ = obj_pair_xyzs.shape
+                # obj_pair_xyzs = obj_pair_xyzs.reshape(cur_batch_size * num_comb, num_pair_pc_pts, -1)
+                collision_logits = self.collision_backbone.forward(obj_pair_xyzs.reshape(cur_batch_size * num_comb, num_pair_pc_pts, -1))
+                collision_scores = self.collision_backbone.convert_logits(collision_logits).reshape(cur_batch_size, num_comb)  # cur_batch_size, num_comb
+
+                # debug
+                # for bi, this_obj_pair_xyzs in enumerate(obj_pair_xyzs):
+                #     print("batch id", bi)
+                #     for pi, obj_pair_xyz in enumerate(this_obj_pair_xyzs):
+                #         print("pair", pi)
+                #         # obj_pair_xyzs: 2 * P, 5
+                #         print("collision score", collision_scores[bi, pi])
+                #         trimesh.PointCloud(obj_pair_xyz[:, :3].cpu()).show()
+
+                # 1 - mean() since the collision model predicts 1 if there is a collision
+                no_intersection_scores[cur_batch_idxs_start:cur_batch_idxs_end] = 1 - torch.mean(collision_scores, dim=1)
+            if visualize:
+                print("no intersection scores", no_intersection_scores)
+            # #######################################
+            # if discriminator_score_weight > 0:
+            #     # # debug:
+            #     # print(subsampled_scene_xyz.shape)
+            #     # print(subsampled_scene_xyz[0])
+            #     # trimesh.PointCloud(subsampled_scene_xyz[0, :, :3].cpu().numpy()).show()
+            #     #
+            #     with torch.no_grad():
+            #
+            #         # Important: since this discriminator only uses local structure param, takes sentence from the first and last position
+            #         # local_sentence = sentence[:, [0, 4]]
+            #         # local_sentence_pad_mask = sentence_pad_mask[:, [0, 4]]
+            #         # sentence_disc, sentence_pad_mask_disc, position_index_dic = discriminator_inference.dataset.tensorfy_sentence(raw_sentence_discriminator, raw_sentence_pad_mask_discriminator, raw_position_index_discriminator)
+            #
+            #         sentence_disc = torch.LongTensor(
+            #             [discriminator_tokenizer.tokenize(*i) for i in raw_sentence_discriminator])
+            #         sentence_pad_mask_disc = torch.LongTensor(raw_sentence_pad_mask_discriminator)
+            #         position_index_dic = torch.LongTensor(raw_position_index_discriminator)
+            #
+            #         preds = discriminator_model.forward(subsampled_scene_xyz,
+            #                                             sentence_disc.unsqueeze(0).repeat(cur_batch_size, 1).to(device),
+            #                                             sentence_pad_mask_disc.unsqueeze(0).repeat(cur_batch_size,
+            #                                                                                        1).to(device),
+            #                                             position_index_dic.unsqueeze(0).repeat(cur_batch_size, 1).to(
+            #                                                 device))
+            #         # preds = discriminator_model.forward(subsampled_scene_xyz)
+            #         preds = discriminator_model.convert_logits(preds)
+            #         preds = preds["is_circle"]  # cur_batch_size,
+            #         scores[cur_batch_idxs_start:cur_batch_idxs_end] = preds
+            #     if visualize:
+            #         print("discriminator scores", scores)
+
+        # scores = scores * discriminator_score_weight + no_intersection_scores * collision_score_weight
+        scores = no_intersection_scores
+        sort_idx = torch.argsort(scores).flip(dims=[0])[:num_elite]
+        elite_obj_params = obj_params[sort_idx]  # num_elite, N, 6
+        elite_struct_poses = struct_pose.reshape(S, N, 4, 4)[sort_idx]  # num_elite, N, 4, 4
+        elite_struct_poses = elite_struct_poses.reshape(num_elite * N, 4, 4)  # num_elite x N, 4, 4
+        elite_scores = scores[sort_idx]
+        print("elite scores:", elite_scores)
+
+        ####################################################
+        # # visualize best samples
+        # num_scene_pts = 4096 # if discriminator_num_scene_pts is None else discriminator_num_scene_pts
+        # batch_current_pc_pose = current_pc_pose[0: num_elite * N]
+        # best_new_obj_xyzs, best_goal_pc_pose, best_subsampled_scene_xyz, _, _ = \
+        #     move_pc_and_create_scene_new(obj_xyzs, elite_obj_params, elite_struct_poses, batch_current_pc_pose,
+        #                                  target_object_inds, self.device,
+        #                                  return_scene_pts=True, num_scene_pts=num_scene_pts, normalize_pc=True)
+        # if visualize:
+        #     print("visualizing elite rearrangements ranked by collision model/discriminator")
+        #     visualize_batch_pcs(best_new_obj_xyzs, num_elite, limit_B=num_elite)
+
+        # num_elite, N, 6
+        elite_obj_params = elite_obj_params.reshape(num_elite * N, -1)
+        pc_poses_in_struct = torch.eye(4).repeat(num_elite * N, 1, 1).to(self.device)
+        pc_poses_in_struct[:, :3, :3] = tra3d.euler_angles_to_matrix(elite_obj_params[:, 3:], "XYZ")
+        pc_poses_in_struct[:, :3, 3] = elite_obj_params[:, :3]
+        pc_poses_in_struct = pc_poses_in_struct.reshape(num_elite, N, 4, 4)  # num_elite, N, 4, 4
+
+        struct_pose = elite_struct_poses.reshape(num_elite, N, 4, 4)[:, 0,].unsqueeze(1)  # num_elite, 1, 4, 4
+
+        print(struct_pose.shape)
+        print(pc_poses_in_struct.shape)
+
+        return struct_pose, pc_poses_in_struct

src/StructDiffusion/language/convert_to_natural_language.ipynb

@@ -0,0 +1,773 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 51,
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import h5py\n",
+    "import json\n",
+    "import numpy as np\n",
+    "import tqdm\n",
+    "import itertools\n",
+    "import copy\n",
+    "from collections import defaultdict\n",
+    "\n",
+    "from StructDiffuser.tokenizer import Tokenizer"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {
+    "collapsed": true,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "class SemanticArrangementDataset:\n",
+    "\n",
+    "    def __init__(self, data_roots, index_roots, splits, tokenizer):\n",
+    "\n",
+    "        self.data_roots = data_roots\n",
+    "        print(\"data dirs:\", self.data_roots)\n",
+    "\n",
+    "        self.tokenizer = tokenizer\n",
+    "\n",
+    "        self.arrangement_data = []\n",
+    "        arrangement_steps = []\n",
+    "        for split in splits:\n",
+    "            for data_root, index_root in zip(data_roots, index_roots):\n",
+    "                arrangement_indices_file = os.path.join(data_root, index_root, \"{}_arrangement_indices_file_all.txt\".format(split))\n",
+    "                if os.path.exists(arrangement_indices_file):\n",
+    "                    with open(arrangement_indices_file, \"r\") as fh:\n",
+    "                        arrangement_steps.extend([(os.path.join(data_root, f[0]), f[1]) for f in eval(fh.readline().strip())])\n",
+    "                else:\n",
+    "                    print(\"{} does not exist\".format(arrangement_indices_file))\n",
+    "\n",
+    "        # only keep one dummy step for each rearrangement\n",
+    "        for filename, step_t in arrangement_steps:\n",
+    "            if step_t == 0:\n",
+    "                self.arrangement_data.append(filename)\n",
+    "        print(\"{} valid sequences\".format(len(self.arrangement_data)))\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.arrangement_data)\n",
+    "\n",
+    "    def get_raw_data(self, idx):\n",
+    "\n",
+    "        filename = self.arrangement_data[idx]\n",
+    "        h5 = h5py.File(filename, 'r')\n",
+    "        goal_specification = json.loads(str(np.array(h5[\"goal_specification\"])))\n",
+    "\n",
+    "        ###################################\n",
+    "        # preparing sentence\n",
+    "        struct_spec = []\n",
+    "\n",
+    "        # structure parameters\n",
+    "        # 5 parameters\n",
+    "        structure_parameters = goal_specification[\"shape\"]\n",
+    "        if structure_parameters[\"type\"] == \"circle\" or structure_parameters[\"type\"] == \"line\":\n",
+    "            struct_spec.append((structure_parameters[\"type\"], \"shape\"))\n",
+    "            struct_spec.append((structure_parameters[\"rotation\"][2], \"rotation\"))\n",
+    "            struct_spec.append((structure_parameters[\"position\"][0], \"position_x\"))\n",
+    "            struct_spec.append((structure_parameters[\"position\"][1], \"position_y\"))\n",
+    "            if structure_parameters[\"type\"] == \"circle\":\n",
+    "                struct_spec.append((structure_parameters[\"radius\"], \"radius\"))\n",
+    "            elif structure_parameters[\"type\"] == \"line\":\n",
+    "                struct_spec.append((structure_parameters[\"length\"] / 2.0, \"radius\"))\n",
+    "        else:\n",
+    "            struct_spec.append((structure_parameters[\"type\"], \"shape\"))\n",
+    "            struct_spec.append((structure_parameters[\"rotation\"][2], \"rotation\"))\n",
+    "            struct_spec.append((structure_parameters[\"position\"][0], \"position_x\"))\n",
+    "            struct_spec.append((structure_parameters[\"position\"][1], \"position_y\"))\n",
+    "\n",
+    "        return struct_spec"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Build one vacab for everything...\n",
+      "The vocab has 124 tokens: {'PAD': 0, 'CLS': 1, 'class:MASK': 2, 'class:Basket': 3, 'class:BeerBottle': 4, 'class:Book': 5, 'class:Bottle': 6, 'class:Bowl': 7, 'class:Calculator': 8, 'class:Candle': 9, 'class:CellPhone': 10, 'class:ComputerMouse': 11, 'class:Controller': 12, 'class:Cup': 13, 'class:Donut': 14, 'class:Fork': 15, 'class:Hammer': 16, 'class:Knife': 17, 'class:Marker': 18, 'class:MilkCarton': 19, 'class:Mug': 20, 'class:Pan': 21, 'class:Pen': 22, 'class:PillBottle': 23, 'class:Plate': 24, 'class:PowerStrip': 25, 'class:Scissors': 26, 'class:SoapBottle': 27, 'class:SodaCan': 28, 'class:Spoon': 29, 'class:Stapler': 30, 'class:Teapot': 31, 'class:VideoGameController': 32, 'class:WineBottle': 33, 'class:CanOpener': 34, 'class:Fruit': 35, 'scene:MASK': 36, 'scene:dinner': 37, 'size:MASK': 38, 'size:L': 39, 'size:M': 40, 'size:S': 41, 'color:MASK': 42, 'color:blue': 43, 'color:cyan': 44, 'color:green': 45, 'color:magenta': 46, 'color:red': 47, 'color:yellow': 48, 'material:MASK': 49, 'material:glass': 50, 'material:metal': 51, 'material:plastic': 52, 'radius:MASK': 53, 'radius:less': 54, 'radius:greater': 55, 'radius:equal': 56, 'radius:0': 57, 'radius:1': 58, 'radius:2': 59, 'position_x:MASK': 60, 'position_x:less': 61, 'position_x:greater': 62, 'position_x:equal': 63, 'position_x:0': 64, 'position_x:1': 65, 'position_x:2': 66, 'position_y:MASK': 67, 'position_y:less': 68, 'position_y:greater': 69, 'position_y:equal': 70, 'position_y:0': 71, 'position_y:1': 72, 'position_y:2': 73, 'rotation:MASK': 74, 'rotation:less': 75, 'rotation:greater': 76, 'rotation:equal': 77, 'rotation:0': 78, 'rotation:1': 79, 'rotation:2': 80, 'rotation:3': 81, 'height:MASK': 82, 'height:less': 83, 'height:greater': 84, 'height:equal': 85, 'height:0': 86, 'height:1': 87, 'height:2': 88, 'height:3': 89, 'height:4': 90, 'height:5': 91, 'height:6': 92, 'height:7': 93, 'height:8': 94, 'height:9': 95, 'volumn:MASK': 96, 'volumn:less': 97, 'volumn:greater': 98, 'volumn:equal': 99, 'volumn:0': 100, 'volumn:1': 101, 'volumn:2': 102, 'volumn:3': 103, 'volumn:4': 104, 'volumn:5': 105, 'volumn:6': 106, 'volumn:7': 107, 'volumn:8': 108, 'volumn:9': 109, 'uniform_angle:MASK': 110, 'uniform_angle:False': 111, 'uniform_angle:True': 112, 'face_center:MASK': 113, 'face_center:False': 114, 'face_center:True': 115, 'angle_ratio:MASK': 116, 'angle_ratio:0.5': 117, 'angle_ratio:1.0': 118, 'shape:MASK': 119, 'shape:circle': 120, 'shape:line': 121, 'shape:tower': 122, 'shape:dinner': 123}\n",
+      "\n",
+      "Build vocabs for object position\n",
+      "The obj_x vocab has 202 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201}\n",
+      "The obj_y vocab has 202 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201}\n",
+      "The obj_z vocab has 202 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201}\n",
+      "The obj_rr vocab has 362 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201, '200': 202, '201': 203, '202': 204, '203': 205, '204': 206, '205': 207, '206': 208, '207': 209, '208': 210, '209': 211, '210': 212, '211': 213, '212': 214, '213': 215, '214': 216, '215': 217, '216': 218, '217': 219, '218': 220, '219': 221, '220': 222, '221': 223, '222': 224, '223': 225, '224': 226, '225': 227, '226': 228, '227': 229, '228': 230, '229': 231, '230': 232, '231': 233, '232': 234, '233': 235, '234': 236, '235': 237, '236': 238, '237': 239, '238': 240, '239': 241, '240': 242, '241': 243, '242': 244, '243': 245, '244': 246, '245': 247, '246': 248, '247': 249, '248': 250, '249': 251, '250': 252, '251': 253, '252': 254, '253': 255, '254': 256, '255': 257, '256': 258, '257': 259, '258': 260, '259': 261, '260': 262, '261': 263, '262': 264, '263': 265, '264': 266, '265': 267, '266': 268, '267': 269, '268': 270, '269': 271, '270': 272, '271': 273, '272': 274, '273': 275, '274': 276, '275': 277, '276': 278, '277': 279, '278': 280, '279': 281, '280': 282, '281': 283, '282': 284, '283': 285, '284': 286, '285': 287, '286': 288, '287': 289, '288': 290, '289': 291, '290': 292, '291': 293, '292': 294, '293': 295, '294': 296, '295': 297, '296': 298, '297': 299, '298': 300, '299': 301, '300': 302, '301': 303, '302': 304, '303': 305, '304': 306, '305': 307, '306': 308, '307': 309, '308': 310, '309': 311, '310': 312, '311': 313, '312': 314, '313': 315, '314': 316, '315': 317, '316': 318, '317': 319, '318': 320, '319': 321, '320': 322, '321': 323, '322': 324, '323': 325, '324': 326, '325': 327, '326': 328, '327': 329, '328': 330, '329': 331, '330': 332, '331': 333, '332': 334, '333': 335, '334': 336, '335': 337, '336': 338, '337': 339, '338': 340, '339': 341, '340': 342, '341': 343, '342': 344, '343': 345, '344': 346, '345': 347, '346': 348, '347': 349, '348': 350, '349': 351, '350': 352, '351': 353, '352': 354, '353': 355, '354': 356, '355': 357, '356': 358, '357': 359, '358': 360, '359': 361}\n",
+      "The obj_rp vocab has 362 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201, '200': 202, '201': 203, '202': 204, '203': 205, '204': 206, '205': 207, '206': 208, '207': 209, '208': 210, '209': 211, '210': 212, '211': 213, '212': 214, '213': 215, '214': 216, '215': 217, '216': 218, '217': 219, '218': 220, '219': 221, '220': 222, '221': 223, '222': 224, '223': 225, '224': 226, '225': 227, '226': 228, '227': 229, '228': 230, '229': 231, '230': 232, '231': 233, '232': 234, '233': 235, '234': 236, '235': 237, '236': 238, '237': 239, '238': 240, '239': 241, '240': 242, '241': 243, '242': 244, '243': 245, '244': 246, '245': 247, '246': 248, '247': 249, '248': 250, '249': 251, '250': 252, '251': 253, '252': 254, '253': 255, '254': 256, '255': 257, '256': 258, '257': 259, '258': 260, '259': 261, '260': 262, '261': 263, '262': 264, '263': 265, '264': 266, '265': 267, '266': 268, '267': 269, '268': 270, '269': 271, '270': 272, '271': 273, '272': 274, '273': 275, '274': 276, '275': 277, '276': 278, '277': 279, '278': 280, '279': 281, '280': 282, '281': 283, '282': 284, '283': 285, '284': 286, '285': 287, '286': 288, '287': 289, '288': 290, '289': 291, '290': 292, '291': 293, '292': 294, '293': 295, '294': 296, '295': 297, '296': 298, '297': 299, '298': 300, '299': 301, '300': 302, '301': 303, '302': 304, '303': 305, '304': 306, '305': 307, '306': 308, '307': 309, '308': 310, '309': 311, '310': 312, '311': 313, '312': 314, '313': 315, '314': 316, '315': 317, '316': 318, '317': 319, '318': 320, '319': 321, '320': 322, '321': 323, '322': 324, '323': 325, '324': 326, '325': 327, '326': 328, '327': 329, '328': 330, '329': 331, '330': 332, '331': 333, '332': 334, '333': 335, '334': 336, '335': 337, '336': 338, '337': 339, '338': 340, '339': 341, '340': 342, '341': 343, '342': 344, '343': 345, '344': 346, '345': 347, '346': 348, '347': 349, '348': 350, '349': 351, '350': 352, '351': 353, '352': 354, '353': 355, '354': 356, '355': 357, '356': 358, '357': 359, '358': 360, '359': 361}\n",
+      "The obj_ry vocab has 362 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201, '200': 202, '201': 203, '202': 204, '203': 205, '204': 206, '205': 207, '206': 208, '207': 209, '208': 210, '209': 211, '210': 212, '211': 213, '212': 214, '213': 215, '214': 216, '215': 217, '216': 218, '217': 219, '218': 220, '219': 221, '220': 222, '221': 223, '222': 224, '223': 225, '224': 226, '225': 227, '226': 228, '227': 229, '228': 230, '229': 231, '230': 232, '231': 233, '232': 234, '233': 235, '234': 236, '235': 237, '236': 238, '237': 239, '238': 240, '239': 241, '240': 242, '241': 243, '242': 244, '243': 245, '244': 246, '245': 247, '246': 248, '247': 249, '248': 250, '249': 251, '250': 252, '251': 253, '252': 254, '253': 255, '254': 256, '255': 257, '256': 258, '257': 259, '258': 260, '259': 261, '260': 262, '261': 263, '262': 264, '263': 265, '264': 266, '265': 267, '266': 268, '267': 269, '268': 270, '269': 271, '270': 272, '271': 273, '272': 274, '273': 275, '274': 276, '275': 277, '276': 278, '277': 279, '278': 280, '279': 281, '280': 282, '281': 283, '282': 284, '283': 285, '284': 286, '285': 287, '286': 288, '287': 289, '288': 290, '289': 291, '290': 292, '291': 293, '292': 294, '293': 295, '294': 296, '295': 297, '296': 298, '297': 299, '298': 300, '299': 301, '300': 302, '301': 303, '302': 304, '303': 305, '304': 306, '305': 307, '306': 308, '307': 309, '308': 310, '309': 311, '310': 312, '311': 313, '312': 314, '313': 315, '314': 316, '315': 317, '316': 318, '317': 319, '318': 320, '319': 321, '320': 322, '321': 323, '322': 324, '323': 325, '324': 326, '325': 327, '326': 328, '327': 329, '328': 330, '329': 331, '330': 332, '331': 333, '332': 334, '333': 335, '334': 336, '335': 337, '336': 338, '337': 339, '338': 340, '339': 341, '340': 342, '341': 343, '342': 344, '343': 345, '344': 346, '345': 347, '346': 348, '347': 349, '348': 350, '349': 351, '350': 352, '351': 353, '352': 354, '353': 355, '354': 356, '355': 357, '356': 358, '357': 359, '358': 360, '359': 361}\n",
+      "The struct_x vocab has 202 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201}\n",
+      "The struct_y vocab has 202 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201}\n",
+      "The struct_z vocab has 202 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201}\n",
+      "The struct_rr vocab has 362 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201, '200': 202, '201': 203, '202': 204, '203': 205, '204': 206, '205': 207, '206': 208, '207': 209, '208': 210, '209': 211, '210': 212, '211': 213, '212': 214, '213': 215, '214': 216, '215': 217, '216': 218, '217': 219, '218': 220, '219': 221, '220': 222, '221': 223, '222': 224, '223': 225, '224': 226, '225': 227, '226': 228, '227': 229, '228': 230, '229': 231, '230': 232, '231': 233, '232': 234, '233': 235, '234': 236, '235': 237, '236': 238, '237': 239, '238': 240, '239': 241, '240': 242, '241': 243, '242': 244, '243': 245, '244': 246, '245': 247, '246': 248, '247': 249, '248': 250, '249': 251, '250': 252, '251': 253, '252': 254, '253': 255, '254': 256, '255': 257, '256': 258, '257': 259, '258': 260, '259': 261, '260': 262, '261': 263, '262': 264, '263': 265, '264': 266, '265': 267, '266': 268, '267': 269, '268': 270, '269': 271, '270': 272, '271': 273, '272': 274, '273': 275, '274': 276, '275': 277, '276': 278, '277': 279, '278': 280, '279': 281, '280': 282, '281': 283, '282': 284, '283': 285, '284': 286, '285': 287, '286': 288, '287': 289, '288': 290, '289': 291, '290': 292, '291': 293, '292': 294, '293': 295, '294': 296, '295': 297, '296': 298, '297': 299, '298': 300, '299': 301, '300': 302, '301': 303, '302': 304, '303': 305, '304': 306, '305': 307, '306': 308, '307': 309, '308': 310, '309': 311, '310': 312, '311': 313, '312': 314, '313': 315, '314': 316, '315': 317, '316': 318, '317': 319, '318': 320, '319': 321, '320': 322, '321': 323, '322': 324, '323': 325, '324': 326, '325': 327, '326': 328, '327': 329, '328': 330, '329': 331, '330': 332, '331': 333, '332': 334, '333': 335, '334': 336, '335': 337, '336': 338, '337': 339, '338': 340, '339': 341, '340': 342, '341': 343, '342': 344, '343': 345, '344': 346, '345': 347, '346': 348, '347': 349, '348': 350, '349': 351, '350': 352, '351': 353, '352': 354, '353': 355, '354': 356, '355': 357, '356': 358, '357': 359, '358': 360, '359': 361}\n",
+      "The struct_rp vocab has 362 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201, '200': 202, '201': 203, '202': 204, '203': 205, '204': 206, '205': 207, '206': 208, '207': 209, '208': 210, '209': 211, '210': 212, '211': 213, '212': 214, '213': 215, '214': 216, '215': 217, '216': 218, '217': 219, '218': 220, '219': 221, '220': 222, '221': 223, '222': 224, '223': 225, '224': 226, '225': 227, '226': 228, '227': 229, '228': 230, '229': 231, '230': 232, '231': 233, '232': 234, '233': 235, '234': 236, '235': 237, '236': 238, '237': 239, '238': 240, '239': 241, '240': 242, '241': 243, '242': 244, '243': 245, '244': 246, '245': 247, '246': 248, '247': 249, '248': 250, '249': 251, '250': 252, '251': 253, '252': 254, '253': 255, '254': 256, '255': 257, '256': 258, '257': 259, '258': 260, '259': 261, '260': 262, '261': 263, '262': 264, '263': 265, '264': 266, '265': 267, '266': 268, '267': 269, '268': 270, '269': 271, '270': 272, '271': 273, '272': 274, '273': 275, '274': 276, '275': 277, '276': 278, '277': 279, '278': 280, '279': 281, '280': 282, '281': 283, '282': 284, '283': 285, '284': 286, '285': 287, '286': 288, '287': 289, '288': 290, '289': 291, '290': 292, '291': 293, '292': 294, '293': 295, '294': 296, '295': 297, '296': 298, '297': 299, '298': 300, '299': 301, '300': 302, '301': 303, '302': 304, '303': 305, '304': 306, '305': 307, '306': 308, '307': 309, '308': 310, '309': 311, '310': 312, '311': 313, '312': 314, '313': 315, '314': 316, '315': 317, '316': 318, '317': 319, '318': 320, '319': 321, '320': 322, '321': 323, '322': 324, '323': 325, '324': 326, '325': 327, '326': 328, '327': 329, '328': 330, '329': 331, '330': 332, '331': 333, '332': 334, '333': 335, '334': 336, '335': 337, '336': 338, '337': 339, '338': 340, '339': 341, '340': 342, '341': 343, '342': 344, '343': 345, '344': 346, '345': 347, '346': 348, '347': 349, '348': 350, '349': 351, '350': 352, '351': 353, '352': 354, '353': 355, '354': 356, '355': 357, '356': 358, '357': 359, '358': 360, '359': 361}\n",
+      "The struct_ry vocab has 362 tokens: {'PAD': 0, 'MASK': 1, '0': 2, '1': 3, '2': 4, '3': 5, '4': 6, '5': 7, '6': 8, '7': 9, '8': 10, '9': 11, '10': 12, '11': 13, '12': 14, '13': 15, '14': 16, '15': 17, '16': 18, '17': 19, '18': 20, '19': 21, '20': 22, '21': 23, '22': 24, '23': 25, '24': 26, '25': 27, '26': 28, '27': 29, '28': 30, '29': 31, '30': 32, '31': 33, '32': 34, '33': 35, '34': 36, '35': 37, '36': 38, '37': 39, '38': 40, '39': 41, '40': 42, '41': 43, '42': 44, '43': 45, '44': 46, '45': 47, '46': 48, '47': 49, '48': 50, '49': 51, '50': 52, '51': 53, '52': 54, '53': 55, '54': 56, '55': 57, '56': 58, '57': 59, '58': 60, '59': 61, '60': 62, '61': 63, '62': 64, '63': 65, '64': 66, '65': 67, '66': 68, '67': 69, '68': 70, '69': 71, '70': 72, '71': 73, '72': 74, '73': 75, '74': 76, '75': 77, '76': 78, '77': 79, '78': 80, '79': 81, '80': 82, '81': 83, '82': 84, '83': 85, '84': 86, '85': 87, '86': 88, '87': 89, '88': 90, '89': 91, '90': 92, '91': 93, '92': 94, '93': 95, '94': 96, '95': 97, '96': 98, '97': 99, '98': 100, '99': 101, '100': 102, '101': 103, '102': 104, '103': 105, '104': 106, '105': 107, '106': 108, '107': 109, '108': 110, '109': 111, '110': 112, '111': 113, '112': 114, '113': 115, '114': 116, '115': 117, '116': 118, '117': 119, '118': 120, '119': 121, '120': 122, '121': 123, '122': 124, '123': 125, '124': 126, '125': 127, '126': 128, '127': 129, '128': 130, '129': 131, '130': 132, '131': 133, '132': 134, '133': 135, '134': 136, '135': 137, '136': 138, '137': 139, '138': 140, '139': 141, '140': 142, '141': 143, '142': 144, '143': 145, '144': 146, '145': 147, '146': 148, '147': 149, '148': 150, '149': 151, '150': 152, '151': 153, '152': 154, '153': 155, '154': 156, '155': 157, '156': 158, '157': 159, '158': 160, '159': 161, '160': 162, '161': 163, '162': 164, '163': 165, '164': 166, '165': 167, '166': 168, '167': 169, '168': 170, '169': 171, '170': 172, '171': 173, '172': 174, '173': 175, '174': 176, '175': 177, '176': 178, '177': 179, '178': 180, '179': 181, '180': 182, '181': 183, '182': 184, '183': 185, '184': 186, '185': 187, '186': 188, '187': 189, '188': 190, '189': 191, '190': 192, '191': 193, '192': 194, '193': 195, '194': 196, '195': 197, '196': 198, '197': 199, '198': 200, '199': 201, '200': 202, '201': 203, '202': 204, '203': 205, '204': 206, '205': 207, '206': 208, '207': 209, '208': 210, '209': 211, '210': 212, '211': 213, '212': 214, '213': 215, '214': 216, '215': 217, '216': 218, '217': 219, '218': 220, '219': 221, '220': 222, '221': 223, '222': 224, '223': 225, '224': 226, '225': 227, '226': 228, '227': 229, '228': 230, '229': 231, '230': 232, '231': 233, '232': 234, '233': 235, '234': 236, '235': 237, '236': 238, '237': 239, '238': 240, '239': 241, '240': 242, '241': 243, '242': 244, '243': 245, '244': 246, '245': 247, '246': 248, '247': 249, '248': 250, '249': 251, '250': 252, '251': 253, '252': 254, '253': 255, '254': 256, '255': 257, '256': 258, '257': 259, '258': 260, '259': 261, '260': 262, '261': 263, '262': 264, '263': 265, '264': 266, '265': 267, '266': 268, '267': 269, '268': 270, '269': 271, '270': 272, '271': 273, '272': 274, '273': 275, '274': 276, '275': 277, '276': 278, '277': 279, '278': 280, '279': 281, '280': 282, '281': 283, '282': 284, '283': 285, '284': 286, '285': 287, '286': 288, '287': 289, '288': 290, '289': 291, '290': 292, '291': 293, '292': 294, '293': 295, '294': 296, '295': 297, '296': 298, '297': 299, '298': 300, '299': 301, '300': 302, '301': 303, '302': 304, '303': 305, '304': 306, '305': 307, '306': 308, '307': 309, '308': 310, '309': 311, '310': 312, '311': 313, '312': 314, '313': 315, '314': 316, '315': 317, '316': 318, '317': 319, '318': 320, '319': 321, '320': 322, '321': 323, '322': 324, '323': 325, '324': 326, '325': 327, '326': 328, '327': 329, '328': 330, '329': 331, '330': 332, '331': 333, '332': 334, '333': 335, '334': 336, '335': 337, '336': 338, '337': 339, '338': 340, '339': 341, '340': 342, '341': 343, '342': 344, '343': 345, '344': 346, '345': 347, '346': 348, '347': 349, '348': 350, '349': 351, '350': 352, '351': 353, '352': 354, '353': 355, '354': 356, '355': 357, '356': 358, '357': 359, '358': 360, '359': 361}\n",
+      "data dirs: ['/home/weiyu/data_drive/data_new_objects/examples_circle_new_objects/result', '/home/weiyu/data_drive/data_new_objects/examples_line_new_objects/result', '/home/weiyu/data_drive/data_new_objects/examples_tower_new_objects/result', '/home/weiyu/data_drive/data_new_objects/examples_dinner_new_objects/result']\n",
+      "40000 valid sequences\n"
+     ]
+    }
+   ],
+   "source": [
+    "tokenizer = Tokenizer(\"/home/weiyu/data_drive/data_new_objects/type_vocabs_coarse.json\")\n",
+    "\n",
+    "data_roots = []\n",
+    "index_roots = []\n",
+    "for shape, index in [(\"circle\", \"index_10k\"), (\"line\", \"index_10k\"), (\"tower\", \"index_10k\"), (\"dinner\", \"index_10k\")]:\n",
+    "    data_roots.append(\"/home/weiyu/data_drive/data_new_objects/examples_{}_new_objects/result\".format(shape))\n",
+    "    index_roots.append(index)\n",
+    "\n",
+    "dataset = SemanticArrangementDataset(data_roots=data_roots, index_roots=index_roots, splits=[\"train\", \"valid\", \"test\"], tokenizer=tokenizer)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "{'place_at_once': 'False', 'position': [0.4530459674902468, 0.2866384076623889, 0.011194709806729462], 'rotation': [5.101818936729106e-05, 1.362746309147995e-06, 2.145504341444197], 'type': 'tower'}\n",
+      "[('tower', 'shape'), (2.145504341444197, 'rotation'), (0.4530459674902468, 'position_x'), (0.2866384076623889, 'position_y')]\n",
+      "tower in the middle left of the table facing west\n",
+      "[('tower', 'shape'), (2.145504341444197, 'rotation'), (0.4530459674902468, 'position_x'), (0.2866384076623889, 'position_y')]\n",
+      "tower in the middle left of the table facing west\n",
+      "(('rotation', 'west'), ('shape', 'tower'), ('x', 'middle'), ('y', 'left'))\n",
+      "\n",
+      "\n",
+      "{'length': 0.15789473684210525, 'length_increment': 0.05, 'max_length': 1.0, 'min_length': 0.0, 'place_at_once': 'True', 'position': [0.5744088910421017, 0.0, 0.0], 'rotation': [0.0, -0.0, 0.0], 'type': 'dinner', 'uniform_space': 'False'}\n",
+      "[('dinner', 'shape'), (0.0, 'rotation'), (0.5744088910421017, 'position_x'), (0.0, 'position_y')]\n",
+      "dinner in the middle center of the table facing south\n",
+      "[('dinner', 'shape'), (0.0, 'rotation'), (0.5744088910421017, 'position_x'), (0.0, 'position_y')]\n",
+      "dinner in the middle center of the table facing south\n",
+      "(('rotation', 'south'), ('shape', 'dinner'), ('x', 'middle'), ('y', 'center'))\n",
+      "\n",
+      "\n",
+      "{'place_at_once': 'False', 'position': [0.5300184865230677, -0.11749143967722209, 0.043775766459831195], 'rotation': [8.311828443210225e-05, 2.8403995850279114e-05, -1.9831750137833084], 'type': 'tower'}\n",
+      "[('tower', 'shape'), (-1.9831750137833084, 'rotation'), (0.5300184865230677, 'position_x'), (-0.11749143967722209, 'position_y')]\n",
+      "tower in the middle center of the table facing north\n",
+      "[('tower', 'shape')]\n",
+      "tower\n",
+      "(('shape', 'tower'),)\n",
+      "\n",
+      "\n",
+      "{'length': 0.3157894736842105, 'length_increment': 0.05, 'max_length': 1.0, 'min_length': 0.0, 'place_at_once': 'True', 'position': [0.6482385523146229, 0.0, 0.0], 'rotation': [0.0, -0.0, 0.0], 'type': 'dinner', 'uniform_space': 'False'}\n",
+      "[('dinner', 'shape'), (0.0, 'rotation'), (0.6482385523146229, 'position_x'), (0.0, 'position_y')]\n",
+      "dinner in the top center of the table facing south\n",
+      "[('dinner', 'shape')]\n",
+      "dinner\n",
+      "(('shape', 'dinner'),)\n",
+      "\n",
+      "\n",
+      "{'angle_ratio': 1.0, 'face_center': 'True', 'max_radius': 0.5, 'min_radius': 0.050687861718942046, 'place_at_once': 'True', 'position': [0.2998438437491998, -0.03599718247376027, 0.0], 'radius': 0.0966402394976866, 'radius_increment': 0.005, 'rotation': [0.0, -0.0, 2.053106459668934], 'type': 'circle', 'uniform_angle': 'True'}\n",
+      "[('circle', 'shape'), (2.053106459668934, 'rotation'), (0.2998438437491998, 'position_x'), (-0.03599718247376027, 'position_y'), (0.0966402394976866, 'radius')]\n",
+      "small circle in the middle center of the table facing west\n",
+      "[('circle', 'shape'), (2.053106459668934, 'rotation'), (0.2998438437491998, 'position_x'), (-0.03599718247376027, 'position_y'), (0.0966402394976866, 'radius')]\n",
+      "small circle in the middle center of the table facing west\n",
+      "(('rotation', 'west'), ('shape', 'circle'), ('size', 'small'), ('x', 'middle'), ('y', 'center'))\n",
+      "\n",
+      "\n",
+      "{'length': 0.4245597103515504, 'length_increment': 0.005, 'max_length': 1.0, 'min_length': 0.21760311495166934, 'place_at_once': 'True', 'position': [0.6672547106460816, 0.0, 0.0], 'rotation': [0.0, -0.0, 0.0], 'type': 'line', 'uniform_space': 'True'}\n",
+      "[('line', 'shape'), (0.0, 'rotation'), (0.6672547106460816, 'position_x'), (0.0, 'position_y'), (0.2122798551757752, 'radius')]\n",
+      "medium line in the top center of the table facing south\n",
+      "[('line', 'shape'), (0.0, 'rotation'), (0.6672547106460816, 'position_x'), (0.2122798551757752, 'radius')]\n",
+      "medium line in the top facing south\n",
+      "(('rotation', 'south'), ('shape', 'line'), ('size', 'medium'), ('x', 'top'))\n",
+      "\n",
+      "\n",
+      "{'place_at_once': 'False', 'position': [0.6555576184899171, 0.22241488561049588, 0.006522659915853506], 'rotation': [-0.000139418832574769, -7.243860660016997e-05, 2.2437880740062814], 'type': 'tower'}\n",
+      "[('tower', 'shape'), (2.2437880740062814, 'rotation'), (0.6555576184899171, 'position_x'), (0.22241488561049588, 'position_y')]\n",
+      "tower in the top left of the table facing west\n",
+      "[(2.2437880740062814, 'rotation'), (0.6555576184899171, 'position_x')]\n",
+      "in the top facing west\n",
+      "(('rotation', 'west'), ('x', 'top'))\n",
+      "\n",
+      "\n",
+      "{'length': 0.4925060249864075, 'length_increment': 0.005, 'max_length': 1.0, 'min_length': 0.4925060249864075, 'place_at_once': 'True', 'position': [0.7754676784901477, 0.0, 0.0], 'rotation': [0.0, -0.0, 0.0], 'type': 'line', 'uniform_space': 'False'}\n",
+      "[('line', 'shape'), (0.0, 'rotation'), (0.7754676784901477, 'position_x'), (0.0, 'position_y'), (0.24625301249320375, 'radius')]\n",
+      "medium line in the top center of the table facing south\n",
+      "[(0.0, 'rotation'), (0.7754676784901477, 'position_x')]\n",
+      "in the top facing south\n",
+      "(('rotation', 'south'), ('x', 'top'))\n",
+      "\n",
+      "\n",
+      "{'angle_ratio': 1.0, 'face_center': 'True', 'max_radius': 0.5, 'min_radius': 0.2260219063147572, 'place_at_once': 'True', 'position': [0.6256453430245876, 0.1131426073908803, 0.0], 'radius': 0.2260219063147572, 'radius_increment': 0.005, 'rotation': [0.0, -0.0, 1.6063513593439724], 'type': 'circle', 'uniform_angle': 'True'}\n",
+      "[('circle', 'shape'), (1.6063513593439724, 'rotation'), (0.6256453430245876, 'position_x'), (0.1131426073908803, 'position_y'), (0.2260219063147572, 'radius')]\n",
+      "medium circle in the middle center of the table facing west\n",
+      "[(1.6063513593439724, 'rotation'), (0.6256453430245876, 'position_x')]\n",
+      "in the middle facing west\n",
+      "(('rotation', 'west'), ('x', 'middle'))\n",
+      "\n",
+      "\n",
+      "{'angle_ratio': 1.0, 'face_center': 'True', 'max_radius': 0.5, 'min_radius': 0.14976631196286583, 'place_at_once': 'True', 'position': [0.5157008668336853, 0.11005531020590054, 0.0], 'radius': 0.15991801306539147, 'radius_increment': 0.005, 'rotation': [0.0, -0.0, -2.2145659262893918], 'type': 'circle', 'uniform_angle': 'True'}\n",
+      "[('circle', 'shape'), (-2.2145659262893918, 'rotation'), (0.5157008668336853, 'position_x'), (0.11005531020590054, 'position_y'), (0.15991801306539147, 'radius')]\n",
+      "small circle in the middle center of the table facing north\n",
+      "[('circle', 'shape'), (0.5157008668336853, 'position_x'), (0.15991801306539147, 'radius')]\n",
+      "small circle in the middle\n",
+      "(('shape', 'circle'), ('size', 'small'), ('x', 'middle'))\n"
+     ]
+    }
+   ],
+   "source": [
+    "idxs = np.random.permutation(len(dataset))\n",
+    "for i in idxs[:10]:\n",
+    "    print(\"\\n\")\n",
+    "    struct_spec = dataset.get_raw_data(i)\n",
+    "    print(struct_spec)\n",
+    "    struct_word_spec = tokenizer.convert_structure_params_to_natural_language(struct_spec)\n",
+    "    print(struct_word_spec)\n",
+    "\n",
+    "    token_idxs = np.random.permutation(len(struct_spec))\n",
+    "    token_idxs = token_idxs[:np.random.randint(1, len(struct_spec) + 1)]\n",
+    "    token_idxs = sorted(token_idxs)\n",
+    "    incomplete_struct_spec = [struct_spec[ti] for ti in token_idxs]\n",
+    "\n",
+    "    print(incomplete_struct_spec)\n",
+    "    print(tokenizer.convert_structure_params_to_natural_language(incomplete_struct_spec))\n",
+    "\n",
+    "    type_value_tuple = tokenizer.convert_structure_params_to_type_value_tuple(incomplete_struct_spec)\n",
+    "    print(type_value_tuple)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 49,
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 40000/40000 [00:23<00:00, 1699.94it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "669\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "unique_type_value_tuples = set()\n",
+    "for i in tqdm.tqdm(idxs):\n",
+    "    struct_spec = dataset.get_raw_data(i)\n",
+    "\n",
+    "    incomplete_struct_specs = []\n",
+    "    for L in range(1, len(struct_spec) + 1):\n",
+    "        for subset in itertools.combinations(struct_spec, L):\n",
+    "            incomplete_struct_specs.append(subset)\n",
+    "\n",
+    "    # print(incomplete_struct_specs)\n",
+    "\n",
+    "    type_value_tuples = []\n",
+    "    for incomplete_struct_spec in incomplete_struct_specs:\n",
+    "        type_value_tuples.append(tokenizer.convert_structure_params_to_type_value_tuple(incomplete_struct_spec))\n",
+    "\n",
+    "    unique_type_value_tuples.update(type_value_tuples)\n",
+    "\n",
+    "print(len(unique_type_value_tuples))"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "sentence_template = [\n",
+    "    \"Put the objects {in a [size][shape]} on the {[x][y] of} the table {facing [rotation]}.\",\n",
+    "    \"Build a [size][shape] of the [objects] on the [x][y] of the table facing [rotation].\",\n",
+    "    \"Put the [objects] on the [x][y] of the table and make a [shape] facing [rotation].\",\n",
+    "    \"Rearrange the [objects] into a [shape], and put the structure on the [x][y] of the table facing [rotation].\",\n",
+    "    \"Could you ...\",\n",
+    "    \"Please ...\",\n",
+    "    \"Pick up the objects, put them into a [size][shape], place the [shape] on the [x][y] of table, make sure the [shape] is facing [rotation].\"]\n",
+    "\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "Enumerate all possible combinations of types"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "31\n",
+      "[('size',), ('shape',), ('x',), ('y',), ('rotation',), ('shape', 'size'), ('size', 'x'), ('size', 'y'), ('rotation', 'size'), ('shape', 'x'), ('shape', 'y'), ('rotation', 'shape'), ('x', 'y'), ('rotation', 'x'), ('rotation', 'y'), ('shape', 'size', 'x'), ('shape', 'size', 'y'), ('rotation', 'shape', 'size'), ('size', 'x', 'y'), ('rotation', 'size', 'x'), ('rotation', 'size', 'y'), ('shape', 'x', 'y'), ('rotation', 'shape', 'x'), ('rotation', 'shape', 'y'), ('rotation', 'x', 'y'), ('shape', 'size', 'x', 'y'), ('rotation', 'shape', 'size', 'x'), ('rotation', 'shape', 'size', 'y'), ('rotation', 'size', 'x', 'y'), ('rotation', 'shape', 'x', 'y'), ('rotation', 'shape', 'size', 'x', 'y')]\n"
+     ]
+    }
+   ],
+   "source": [
+    "import itertools\n",
+    "types = [\"size\", \"shape\", \"x\", \"y\", \"rotation\"]\n",
+    "\n",
+    "type_combs = []\n",
+    "for L in range(1, len(types) + 1):\n",
+    "    for subset in itertools.combinations(types, L):\n",
+    "        type_combs.append(tuple(sorted(subset)))\n",
+    "\n",
+    "print(len(type_combs))\n",
+    "print(type_combs)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 46,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "build a [size] shape from the objects ('size',)\n",
+      "put the objects in to a [size] shape ('size',)\n",
+      "place the objects as a [size] shape ('size',)\n",
+      "make a [size] shape from the objects ('size',)\n",
+      "rearrange the objects into a [size] structure ('size',)\n",
+      "build a [shape] ('shape',)\n",
+      "make a [shape] ('shape',)\n",
+      "put the objects into a [shape] ('shape',)\n",
+      "place the objects as a [shape] ('shape',)\n",
+      "pick up the objects, and place them as a [shape] ('shape',)\n",
+      "place the objects on the [x] of the table ('x',)\n",
+      "put the objects on [x] ('x',)\n",
+      "make a structure from the objects and place it on [x] ('x',)\n",
+      "on the [x] of the table, place the objects ('x',)\n",
+      "move the objects to the [x] ('x',)\n",
+      "place the objects on the [y] of the table ('y',)\n",
+      "put the objects on [y] ('y',)\n",
+      "make a structure from the objects and place it on [y] ('y',)\n",
+      "on the [y] of the table, place the objects ('y',)\n",
+      "move the objects to the [y] ('y',)\n",
+      "build a structure facing [rotation] ('rotation',)\n",
+      "make a structure from the objects and make sure it is pointing [rotation] ('rotation',)\n",
+      "put the objects in a structure that faces [rotation] ('rotation',)\n",
+      "rotate the object structure so that it points [rotation] ('rotation',)\n",
+      "[rotation] is the direction the structure built from the objects should be facing ('rotation',)\n",
+      "build a [size] [shape] ('shape', 'size')\n",
+      "make a [size] [shape] ('shape', 'size')\n",
+      "put the objects into a [size] [shape] ('shape', 'size')\n",
+      "place the objects as a [size] [shape] ('shape', 'size')\n",
+      "pick up the objects, and place them as a [size] [shape] ('shape', 'size')\n",
+      "build a [size] shape from the objects on the [x] of the table ('size', 'x')\n",
+      "put the objects in to a [size] shape and place it on [x] ('size', 'x')\n",
+      "on the [x] of the table, place the objects as a [size] shape ('size', 'x')\n",
+      "make a [size] shape from the objects and move it to [x] ('size', 'x')\n",
+      "rearrange the objects into a [size] structure on [x] ('size', 'x')\n",
+      "build a [size] shape from the objects on the [y] of the table ('size', 'y')\n",
+      "put the objects in to a [size] shape and place it on [y] ('size', 'y')\n",
+      "on the [y] of the table, place the objects as a [size] shape ('size', 'y')\n",
+      "make a [size] shape from the objects and move it to [y] ('size', 'y')\n",
+      "rearrange the objects into a [size] structure on [y] ('size', 'y')\n",
+      "build a [size] shape from the objects facing [rotation] ('rotation', 'size')\n",
+      "put the objects in to a [size] shape and place it so that it faces [rotation] ('rotation', 'size')\n",
+      "place the objects as a [size] shape and [rotation] is the direction the shape built from the objects should be facing ('rotation', 'size')\n",
+      "make a [size] structure from the objects and rotate the object structure so that it points [rotation] ('rotation', 'size')\n",
+      "rearrange the objects into a [size] structure that points to [rotation] ('rotation', 'size')\n",
+      "build a [shape] from the objects on the [x] of the table ('shape', 'x')\n",
+      "put the objects in to a [shape] and place it on [x] ('shape', 'x')\n",
+      "on the [x] of the table, place the objects as a [shape] ('shape', 'x')\n",
+      "make a [shape] from the objects and move it to [x] ('shape', 'x')\n",
+      "rearrange the objects into a [shape] on [x] ('shape', 'x')\n",
+      "build a [shape] from the objects on the [y] of the table ('shape', 'y')\n",
+      "put the objects in to a [shape] and place it on [y] ('shape', 'y')\n",
+      "on the [y] of the table, place the objects as a [shape] ('shape', 'y')\n",
+      "make a [shape] from the objects and move it to [y] ('shape', 'y')\n",
+      "rearrange the objects into a [shape] on [y] ('shape', 'y')\n",
+      "build a [shape] from the objects facing [rotation] ('rotation', 'shape')\n",
+      "put the objects in to a [shape] and place it so that it faces [rotation] ('rotation', 'shape')\n",
+      "place the objects as a [shape] and [rotation] is the direction the shape built from the objects should be facing ('rotation', 'shape')\n",
+      "make a [shape] from the objects and rotate the shape so that it points [rotation] ('rotation', 'shape')\n",
+      "rearrange the objects into a [shape] that points to [rotation] ('rotation', 'shape')\n",
+      "place the objects on the [x] and [y] of the table ('x', 'y')\n",
+      "put the objects on [x] [y] of the table ('x', 'y')\n",
+      "make a structure from the objects and place it on [x] [y] ('x', 'y')\n",
+      "on the [x] [y] of the table, place the objects ('x', 'y')\n",
+      "move the objects to the [x] [y] ('x', 'y')\n",
+      "build a structure on the [x] of the table facing [rotation] ('rotation', 'x')\n",
+      "make a structure from the objects and make sure it is pointing [rotation] and on [x] ('rotation', 'x')\n",
+      "rearrange the objects in a structure that faces [rotation] and place it on [x] ('rotation', 'x')\n",
+      "move and rotate the object structure so that it is on [x] and points [rotation] ('rotation', 'x')\n",
+      "[rotation] is the direction the structure built from the objects should be facing, [x] is the location ('rotation', 'x')\n",
+      "build a structure on the [y] of the table facing [rotation] ('rotation', 'y')\n",
+      "make a structure from the objects and make sure it is pointing [rotation] and on [y] ('rotation', 'y')\n",
+      "rearrange the objects in a structure that faces [rotation] and place it on [y] ('rotation', 'y')\n",
+      "move and rotate the object structure so that it is on [y] and points [rotation] ('rotation', 'y')\n",
+      "[rotation] is the direction the structure built from the objects should be facing, [y] is the location ('rotation', 'y')\n",
+      "build a [size] [shape] from the objects on the [x] of the table ('shape', 'size', 'x')\n",
+      "put the objects in to a [size] [shape] and place it on [x] ('shape', 'size', 'x')\n",
+      "on the [x] of the table, place the objects as a [shape], make the shape [size] ('shape', 'size', 'x')\n",
+      "make a [size] [shape] from the objects and move it to [x] ('shape', 'size', 'x')\n",
+      "rearrange the objects into a [size] [shape] on [x] ('shape', 'size', 'x')\n",
+      "build a [size] [shape] from the objects on the [y] of the table ('shape', 'size', 'y')\n",
+      "put the objects in to a [size] [shape] and place it on [y] ('shape', 'size', 'y')\n",
+      "on the [y] of the table, place the objects as a [shape], make the shape [size] ('shape', 'size', 'y')\n",
+      "make a [size] [shape] from the objects and move it to [y] ('shape', 'size', 'y')\n",
+      "rearrange the objects into a [size] [shape] on [y] ('shape', 'size', 'y')\n",
+      "build a [size] [shape] from the objects facing [rotation] ('rotation', 'shape', 'size')\n",
+      "put the objects in to a [size] [shape] and place it so that it faces [rotation] ('rotation', 'shape', 'size')\n",
+      "place the objects as a [size] [shape] and [rotation] is the direction the shape built from the objects should be facing ('rotation', 'shape', 'size')\n",
+      "make a [size] [shape] from the objects and rotate the shape so that it points [rotation] ('rotation', 'shape', 'size')\n",
+      "rearrange the objects into a [size] [shape] that points to [rotation] ('rotation', 'shape', 'size')\n",
+      "build a [size] shape from the objects on the [x] [y] of the table ('size', 'x', 'y')\n",
+      "put the objects in to a [size] shape and place it on [x] and [y] ('size', 'x', 'y')\n",
+      "on the [x] [y] of the table, place the objects as a [size] shape ('size', 'x', 'y')\n",
+      "make a [size] shape from the objects and move it to [x] [y] ('size', 'x', 'y')\n",
+      "rearrange the objects into a [size] structure on [x] and on [y] ('size', 'x', 'y')\n",
+      "build a [size] structure on the [x] of the table facing [rotation] ('rotation', 'size', 'x')\n",
+      "make a [size] structure from the objects and make sure it is pointing [rotation] and on [x] ('rotation', 'size', 'x')\n",
+      "rearrange the objects in a [size] structure that faces [rotation] and place it on [x] ('rotation', 'size', 'x')\n",
+      "move and rotate the [size] object structure so that it is on [x] and points [rotation] ('rotation', 'size', 'x')\n",
+      "[rotation] is the direction the [size] structure built from the objects should be facing, [x] is the location ('rotation', 'size', 'x')\n",
+      "build a [size] structure on the [y] of the table facing [rotation] ('rotation', 'size', 'y')\n",
+      "make a [size] structure from the objects and make sure it is pointing [rotation] and on [y] ('rotation', 'size', 'y')\n",
+      "rearrange the objects in a [size] structure that faces [rotation] and place it on [y] ('rotation', 'size', 'y')\n",
+      "move and rotate the [size] object structure so that it is on [y] and points [rotation] ('rotation', 'size', 'y')\n",
+      "[rotation] is the direction the [size] structure built from the objects should be facing, [y] is the location ('rotation', 'size', 'y')\n",
+      "build a [shape] from the objects on the [x] [y] of the table ('shape', 'x', 'y')\n",
+      "put the objects in to a [shape] and place it on [x] and [y] ('shape', 'x', 'y')\n",
+      "on the [x] [y] of the table, place the objects as a [shape] ('shape', 'x', 'y')\n",
+      "make a [shape] from the objects and move it to [x] [y] ('shape', 'x', 'y')\n",
+      "rearrange the objects into a [shape] on [x] and on [y] ('shape', 'x', 'y')\n",
+      "build a [shape] on the [x] of the table facing [rotation] ('rotation', 'shape', 'x')\n",
+      "make a [shape] from the objects and make sure it is pointing [rotation] and on [x] ('rotation', 'shape', 'x')\n",
+      "rearrange the objects in a [shape] that faces [rotation] and place it on [x] ('rotation', 'shape', 'x')\n",
+      "move and rotate the [shape] so that it is on [x] and points [rotation] ('rotation', 'shape', 'x')\n",
+      "[rotation] is the direction the [shape] built from the objects should be facing, [x] is the location ('rotation', 'shape', 'x')\n",
+      "build a [shape] on the [y] of the table facing [rotation] ('rotation', 'shape', 'y')\n",
+      "make a [shape] from the objects and make sure it is pointing [rotation] and on [y] ('rotation', 'shape', 'y')\n",
+      "rearrange the objects in a [shape] that faces [rotation] and place it on [y] ('rotation', 'shape', 'y')\n",
+      "move and rotate the [shape] so that it is on [y] and points [rotation] ('rotation', 'shape', 'y')\n",
+      "[rotation] is the direction the [shape] built from the objects should be facing, [y] is the location ('rotation', 'shape', 'y')\n",
+      "build a structure on the [x] [y] of the table facing [rotation] ('rotation', 'x', 'y')\n",
+      "make a structure from the objects and make sure it is pointing [rotation] and on [x] [y] ('rotation', 'x', 'y')\n",
+      "rearrange the objects in a structure that faces [rotation] and place it on [x] [y] ('rotation', 'x', 'y')\n",
+      "move and rotate the object structure so that it is on [x] [y] and points [rotation] ('rotation', 'x', 'y')\n",
+      "[rotation] is the direction the structure built from the objects should be facing, [x] [y] is the location ('rotation', 'x', 'y')\n",
+      "build a [shape] from the objects on the [x] [y] of the table, make the [shape] [size] ('shape', 'size', 'x', 'y')\n",
+      "put the objects in to a [size] [shape] and place it on [x] and [y] ('shape', 'size', 'x', 'y')\n",
+      "on the [x] [y] of the table, place the objects as a [size] [shape] ('shape', 'size', 'x', 'y')\n",
+      "make a [size] [shape] from the objects and move it to [x] [y] ('shape', 'size', 'x', 'y')\n",
+      "rearrange the objects into a [size] [shape] on [x] and on [y] ('shape', 'size', 'x', 'y')\n",
+      "build a [size] [shape] on the [x] of the table facing [rotation] ('rotation', 'shape', 'size', 'x')\n",
+      "make a [size] [shape] from the objects and make sure it is pointing [rotation] and on [x] ('rotation', 'shape', 'size', 'x')\n",
+      "rearrange the objects in a [size] [shape] that faces [rotation] and place it on [x] ('rotation', 'shape', 'size', 'x')\n",
+      "move and rotate the [size] [shape] so that it is on [x] and points [rotation] ('rotation', 'shape', 'size', 'x')\n",
+      "[rotation] is the direction the [size] [shape] built from the objects should be facing, [x] is the location ('rotation', 'shape', 'size', 'x')\n",
+      "build a [size] [shape] on the [y] of the table facing [rotation] ('rotation', 'shape', 'size', 'y')\n",
+      "make a [size] [shape] from the objects and make sure it is pointing [rotation] and on [y] ('rotation', 'shape', 'size', 'y')\n",
+      "rearrange the objects in a [size] [shape] that faces [rotation] and place it on [y] ('rotation', 'shape', 'size', 'y')\n",
+      "move and rotate the [size] [shape] so that it is on [y] and points [rotation] ('rotation', 'shape', 'size', 'y')\n",
+      "[rotation] is the direction the [size] [shape] built from the objects should be facing, [y] is the location ('rotation', 'shape', 'size', 'y')\n",
+      "build a [size] structure on the [x] [y] of the table facing [rotation] ('rotation', 'size', 'x', 'y')\n",
+      "make a [size] structure from the objects and make sure it is pointing [rotation] and on [x] [y] ('rotation', 'size', 'x', 'y')\n",
+      "rearrange the objects in a [size] structure that faces [rotation] and place it on [x] [y] ('rotation', 'size', 'x', 'y')\n",
+      "move and rotate the [size] object structure so that it is on [x] [y] and points [rotation] ('rotation', 'size', 'x', 'y')\n",
+      "[rotation] is the direction the [size] structure built from the objects should be facing, [x] [y] is the location ('rotation', 'size', 'x', 'y')\n",
+      "build a [shape] on the [x] [y] of the table facing [rotation] ('rotation', 'shape', 'x', 'y')\n",
+      "make a [shape] from the objects and make sure it is pointing [rotation] and on [x] [y] ('rotation', 'shape', 'x', 'y')\n",
+      "rearrange the objects as a [shape] that faces [rotation] and place it on [x] [y] ('rotation', 'shape', 'x', 'y')\n",
+      "move and rotate the [shape] so that it is on [x] [y] and points [rotation] ('rotation', 'shape', 'x', 'y')\n",
+      "[rotation] is the direction the [shape] built from the objects should be facing, [x] [y] is the location ('rotation', 'shape', 'x', 'y')\n",
+      "build a [size] [shape] on the [x] [y] of the table facing [rotation] ('rotation', 'shape', 'size', 'x', 'y')\n",
+      "make a [size] [shape] from the objects and make sure it is pointing [rotation] and on [x] [y] ('rotation', 'shape', 'size', 'x', 'y')\n",
+      "rearrange the objects as a [size] [shape] that faces [rotation] and place it on [x] [y] ('rotation', 'shape', 'size', 'x', 'y')\n",
+      "move and rotate the [size] [shape] so that it is on [x] [y] and points [rotation] ('rotation', 'shape', 'size', 'x', 'y')\n",
+      "[rotation] is the direction the [size] [shape] built from the objects should be facing, [x] [y] is the location ('rotation', 'shape', 'size', 'x', 'y')\n"
+     ]
+    }
+   ],
+   "source": [
+    "sentence_template_file = \"/home/weiyu/Research/intern/StructDiffuser/src/StructDiffuser/language/sentence_template.txt\"\n",
+    "\n",
+    "import re\n",
+    "\n",
+    "type_comb_to_templates = {}\n",
+    "for type_comb in type_combs:\n",
+    "    type_comb_to_templates[type_comb] = []\n",
+    "\n",
+    "with open(sentence_template_file, \"r\") as fh:\n",
+    "    for line in fh:\n",
+    "        line = line.strip()\n",
+    "        if line:\n",
+    "            if line[0] == \"#\":\n",
+    "                continue\n",
+    "            type_list = re.findall('\\[[^\\]]*\\]', line)\n",
+    "            type_comb = tuple(sorted(list(set([t[1:-1] for t in type_list]))))\n",
+    "            print(line, type_comb)\n",
+    "\n",
+    "            type_comb_to_templates[type_comb].append(line)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 47,
+   "outputs": [],
+   "source": [
+    "for type_comb in type_comb_to_templates:\n",
+    "    if len(type_comb_to_templates[type_comb]) != 5:\n",
+    "        print(\"{} does not have 5 templates\".format(type_comb))"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 58,
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 669/669 [00:00<00:00, 60546.98it/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "template_sentences = []\n",
+    "type_value_tuple_to_template_sentences = defaultdict(set)\n",
+    "for type_value_tuple in tqdm.tqdm(list(unique_type_value_tuples)):\n",
+    "    type_comb = tuple(sorted([tv[0] for tv in type_value_tuple]))\n",
+    "    template_sentences = copy.deepcopy(type_comb_to_templates[type_comb])\n",
+    "\n",
+    "    # print(type_value_tuple)\n",
+    "    for template_sentence in template_sentences:\n",
+    "        for t, v in type_value_tuple:\n",
+    "            template_sentence = template_sentence.replace(\"[{}]\".format(t), v)\n",
+    "        # print(template_sentence)\n",
+    "\n",
+    "        type_value_tuple_to_template_sentences[type_value_tuple].add(template_sentence)\n",
+    "\n",
+    "# convert to list\n",
+    "for type_value_tuple in type_value_tuple_to_template_sentences:\n",
+    "    type_value_tuple_to_template_sentences[type_value_tuple] = list(type_value_tuple_to_template_sentences[type_value_tuple])"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 73,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3345 unique template sentences\n"
+     ]
+    }
+   ],
+   "source": [
+    "unique_template_sentences = set()\n",
+    "\n",
+    "for type_value_tuple in type_value_tuple_to_template_sentences:\n",
+    "    # print(\"\\n\")\n",
+    "    # print(type_value_tuple)\n",
+    "    for template_sentence in type_value_tuple_to_template_sentences[type_value_tuple]:\n",
+    "        # print(template_sentence)\n",
+    "        unique_template_sentences.add(template_sentence)\n",
+    "\n",
+    "unique_template_sentences = list(unique_template_sentences)\n",
+    "print(\"{} unique template sentences\".format(len(unique_template_sentences)))"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 72,
+   "outputs": [],
+   "source": [
+    "from sentence_transformers import SentenceTransformer\n",
+    "model = SentenceTransformer('all-MiniLM-L6-v2')"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 76,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(3345, 384)\n"
+     ]
+    }
+   ],
+   "source": [
+    "#Our sentences we like to encode\n",
+    "# sentences = ['This framework generates embeddings for each input sentence',\n",
+    "#     'Sentences are passed as a list of string.',\n",
+    "#     'The quick brown fox jumps over the lazy dog.']\n",
+    "#Sentences are encoded by calling model.encode()\n",
+    "\n",
+    "\n",
+    "embeddings = model.encode(unique_template_sentences)\n",
+    "print(embeddings.shape)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 80,
+   "outputs": [],
+   "source": [
+    "template_sentence_to_embedding = {}\n",
+    "for embedding, template_sentence in zip(embeddings, unique_template_sentences):\n",
+    "    template_sentence_to_embedding[template_sentence] = embedding"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 82,
+   "outputs": [],
+   "source": [
+    "import pickle\n",
+    "template_sentence_data = {\"template_sentence_to_embedding\": template_sentence_to_embedding,\n",
+    "                          \"type_value_tuple_to_template_sentences\": type_value_tuple_to_template_sentences}\n",
+    "with open(\"/home/weiyu/Research/intern/StructDiffuser/src/StructDiffuser/language/template_sentence_data.pkl\", \"wb\") as fh:\n",
+    "    pickle.dump(template_sentence_data, fh)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

src/StructDiffusion/language/sentence_encoder.py

@@ -0,0 +1,23 @@
+from sentence_transformers import SentenceTransformer
+
+class SentenceBertEncoder:
+
+    def __init__(self):
+        self.model = SentenceTransformer('all-MiniLM-L6-v2')
+
+    def encode(self, sentences):
+        #Our sentences we like to encode
+        # sentences = ['This framework generates embeddings for each input sentence',
+        #     'Sentences are passed as a list of string.',
+        #     'The quick brown fox jumps over the lazy dog.']
+        #Sentences are encoded by calling model.encode()
+
+        embeddings = self.model.encode(sentences)
+        # print(embeddings.shape)
+        return embeddings
+
+
+if __name__ == "__main__":
+    sentence_encoder = SentenceBertEncoder()
+    embedding = sentence_encoder.encode(["this is cool!"])
+    print(embedding.shape)

src/StructDiffusion/language/test_parrot_paraphrase.py

@@ -0,0 +1,38 @@
+from parrot import Parrot
+import torch
+import warnings
+warnings.filterwarnings("ignore")
+
+# [top]
+
+# Put the [objects] in a [size][shape] on the [x][y] of the table facing [rotation].
+# Build a [size][shape] of the [objects] on the [x][y] of the table facing [rotation].
+# Put the [objects] on the [x][y] of the table and make a [shape] facing [rotation].
+# Rearrange the [objects] into a [shape], and put the structure on the [x][y] of the table facing [rotation].
+# Could you ...
+# Please ...
+# Pick up the objects, put them into a [size][shape], place the [shape] on the [x][y] of table, make sure the [shape] is facing [rotation].
+
+if __name__ == "__main__":
+    ''' 
+    uncomment to get reproducable paraphrase generations
+    def random_state(seed):
+      torch.manual_seed(seed)
+      if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    
+    random_state(1234)
+    '''
+
+    #Init models (make sure you init ONLY once if you integrate this to your code)
+    parrot = Parrot(model_tag="prithivida/parrot_paraphraser_on_T5")
+
+    phrases = ["Rearrange the mugs in a circle on the top left of the table."]
+
+    for phrase in phrases:
+      print("-"*100)
+      print("Input_phrase: ", phrase)
+      print("-"*100)
+      para_phrases = parrot.augment(input_phrase=phrase, use_gpu=False, max_return_phrases=100, do_diverse=True)
+      for para_phrase in para_phrases:
+       print(para_phrase)

src/StructDiffusion/language/tokenizer.py

@@ -517,25 +517,4 @@ class ContinuousTokenizer:
                 idx = value
             else:
                 raise KeyError("Do not recognize the type {} of the given token: {}".format(typ, value))
-        return idx
-
-
-if __name__ == "__main__":
-    tokenizer = Tokenizer("/home/weiyu/data_drive/data_new_objects/type_vocabs_coarse.json")
-    # print(tokenizer.get_all_values_of_type("class"))
-    # print(tokenizer.get_all_values_of_type("color"))
-    # print(tokenizer.get_all_values_of_type("material"))
-    #
-    # for type in tokenizer.type_vocabs:
-    #     print(type, tokenizer.type_vocabs[type])
-
-    tokenizer.prepare_grounding_reference()
-
-    # for i in range(100):
-    #     types = list(tokenizer.continuous_types) + list(tokenizer.discrete_types)
-    #     for t in types:
-    #         v = tokenizer.get_valid_random_value(t)
-    #         print(v)
-    #         print(tokenizer.tokenize(v, t))
-
-    # build_vocab("/home/weiyu/data_drive/examples_v4/leonardo/vocab.json", "/home/weiyu/data_drive/examples_v4/leonardo/type_vocabs.json")
+        return idx

src/StructDiffusion/models/models.py

@@ -26,6 +26,8 @@ class TransformerDiffusionModel(torch.nn.Module):
                  word_emb_dim=160,
                  time_emb_dim=80,
                  use_virtual_structure_frame=True,
+                 use_sentence_embedding=False,
+                 sentence_embedding_dim=None,
                  ):
         super(TransformerDiffusionModel, self).__init__()
 
@@ -53,7 +55,12 @@ class TransformerDiffusionModel(torch.nn.Module):
             self.virtual_frame_embed = nn.Parameter(torch.randn(1, 1, posed_pc_emb_dim))  # B, 1, posed_pc_emb_dim
 
         # for language
-        self.word_embeddings = torch.nn.Embedding(vocab_size, word_emb_dim, padding_idx=0)
+        self.sentence_embedding_dim = sentence_embedding_dim
+        self.use_sentence_embedding = use_sentence_embedding
+        if use_sentence_embedding:
+            self.sentence_embedding_down_sample = torch.nn.Linear(sentence_embedding_dim, word_emb_dim)
+        else:
+            self.word_embeddings = torch.nn.Embedding(vocab_size, word_emb_dim, padding_idx=0)
 
         # for diffusion
         self.pose_encoder = nn.Sequential(nn.Linear(action_dim, pose_emb_dim))
@@ -88,7 +95,10 @@ class TransformerDiffusionModel(torch.nn.Module):
 
         batch_size, num_objects, num_pts, _ = pcs.shape
         _, num_poses, _ = poses.shape
-        _, sentence_len = sentence.shape
+        if self.use_sentence_embedding:
+            assert sentence.shape == (batch_size, self.sentence_embedding_dim), sentence.shape
+        else:
+            _, sentence_len = sentence.shape
         _, total_len = type_index.shape
 
         pcs = pcs.reshape(batch_size * num_objects, num_pts, -1)
@@ -102,7 +112,11 @@ class TransformerDiffusionModel(torch.nn.Module):
         tgt_obj_embed = torch.cat([pose_embed, posed_pc_embed], dim=-1)
 
         #########################
-        sentence_embed = self.word_embeddings(sentence)
+        if self.use_sentence_embedding:
+            # sentence: B, sentence_embedding_dim
+            sentence_embed = self.sentence_embedding_down_sample(sentence).unsqueeze(1)  # B, 1, word_emb_dim
+        else:
+            sentence_embed = self.word_embeddings(sentence)
 
         #########################
 

src/StructDiffusion/utils/batch_inference.py

@@ -1,175 +1,10 @@
 import os
 import torch
 import numpy as np
-import pytorch3d.transforms as tra3d
 
 from StructDiffusion.utils.rearrangement import show_pcs_color_order, show_pcs_with_trimesh
 from StructDiffusion.utils.pointnet import random_point_sample, index_points
-
-
-def move_pc_and_create_scene(obj_xyzs, obj_params, struct_pose, current_pc_pose, target_object_inds,
-                             num_scene_pts, device, normalize_pc=False,
-                             return_pair_pc=False, normalize_pair_pc=False, num_pair_pc_pts=None,
-                             return_scene_pts=True, return_scene_pts_and_pc_idxs=False):
-
-    # obj_xyzs: N, P, 3
-    # obj_params: B, N, 6
-    # struct_pose: B x N, 4, 4
-    # current_pc_pose: B x N, 4, 4
-    # target_object_inds: 1, N
-
-    B, N, _ = obj_params.shape
-    _, P, _ = obj_xyzs.shape
-
-    # B, N, 6
-    flat_obj_params = obj_params.reshape(B * N, -1)
-    goal_pc_pose_in_struct = torch.eye(4).repeat(B * N, 1, 1).to(device)
-    goal_pc_pose_in_struct[:, :3, :3] = tra3d.euler_angles_to_matrix(flat_obj_params[:, 3:], "XYZ")
-    goal_pc_pose_in_struct[:, :3, 3] = flat_obj_params[:, :3]  # B x N, 4, 4
-
-    goal_pc_pose = struct_pose @ goal_pc_pose_in_struct
-    goal_pc_transform = goal_pc_pose @ torch.inverse(current_pc_pose)  # cur_batch_size x N, 4, 4
-
-    # important: pytorch3d uses row-major ordering, need to transpose each transformation matrix
-    transpose = tra3d.Transform3d(matrix=goal_pc_transform.transpose(1, 2))
-
-    # obj_xyzs: N, P, 3
-    new_obj_xyzs = obj_xyzs.repeat(B, 1, 1)
-    new_obj_xyzs = transpose.transform_points(new_obj_xyzs)
-
-    # put it back to B, N, P, 3
-    new_obj_xyzs = new_obj_xyzs.reshape(B, N, P, -1)
-    # visualize_batch_pcs(new_obj_xyzs, S, N, P)
-
-    # ===================================
-    # Pass to discriminator
-    subsampled_scene_xyz = None
-    if return_scene_pts:
-
-        num_indicator = N
-
-        # add one hot
-        indicator_variables = torch.eye(num_indicator).repeat(B, 1, 1, P).reshape(B, num_indicator, P, num_indicator).to(device)  # B, N, P, N
-        # print(indicator_variables.shape)
-        # print(new_obj_xyzs.shape)
-        new_obj_xyzs = torch.cat([new_obj_xyzs, indicator_variables], dim=-1)  # B, N, P, 3 + N
-
-        # combine pcs in each scene
-        scene_xyzs = new_obj_xyzs.reshape(B, N * P, 3 + N)
-
-        # ToDo: maybe convert this to a batch operation
-        subsampled_scene_xyz = torch.FloatTensor(B, num_scene_pts, 3 + N).to(device)
-        for si, scene_xyz in enumerate(scene_xyzs):
-            # scene_xyz: N*P, 3+N
-            # target_object_inds: 1, N
-            subsample_idx = torch.randint(0, torch.sum(target_object_inds[0]) * P, (num_scene_pts,)).to(device)
-            subsampled_scene_xyz[si] = scene_xyz[subsample_idx]
-
-            # # debug:
-            # print("-"*50)
-            # if si < 10:
-            #     trimesh.PointCloud(scene_xyz[:, :3].cpu().numpy(), colors=[255, 0, 0, 255]).show()
-            #     trimesh.PointCloud(subsampled_scene_xyz[si, :, :3].cpu().numpy(), colors=[0, 255, 0, 255]).show()
-
-        # subsampled_scene_xyz: B, num_scene_pts, 3+N
-        # new_obj_xyzs: B, N, P, 3
-        # goal_pc_pose: B, N, 4, 4
-
-        # important:
-        if normalize_pc:
-            subsampled_scene_xyz[:, :, 0:3] = pc_normalize_batch(subsampled_scene_xyz[:, :, 0:3])
-
-            # # debug:
-            # for si in range(10):
-            #     trimesh.PointCloud(subsampled_scene_xyz[si, :, :3].cpu().numpy(), colors=[0, 0, 255, 255]).show()
-
-    if return_scene_pts_and_pc_idxs:
-        num_indicator = N
-        pc_idxs = torch.arange(0, num_indicator)[:, None].repeat(B, 1, P).reshape(B, num_indicator, P).to(device)  # B, N, P
-        # new_obj_xyzs: B, N, P, 3 + 1
-
-        # combine pcs in each scene
-        scene_xyzs = new_obj_xyzs.reshape(B, N * P, 3)
-        pc_idxs = pc_idxs.reshape(B, N*P)
-
-        subsampled_scene_xyz = torch.FloatTensor(B, num_scene_pts, 3).to(device)
-        subsampled_pc_idxs = torch.LongTensor(B, num_scene_pts).to(device)
-        for si, (scene_xyz, pc_idx) in enumerate(zip(scene_xyzs, pc_idxs)):
-            # scene_xyz: N*P, 3+1
-            # target_object_inds: 1, N
-            subsample_idx = torch.randint(0, torch.sum(target_object_inds[0]) * P, (num_scene_pts,)).to(device)
-            subsampled_scene_xyz[si] = scene_xyz[subsample_idx]
-            subsampled_pc_idxs[si] = pc_idx[subsample_idx]
-
-        # subsampled_scene_xyz: B, num_scene_pts, 3
-        # subsampled_pc_idxs: B, num_scene_pts
-        # new_obj_xyzs: B, N, P, 3
-        # goal_pc_pose: B, N, 4, 4
-
-        # important:
-        if normalize_pc:
-            subsampled_scene_xyz[:, :, 0:3] = pc_normalize_batch(subsampled_scene_xyz[:, :, 0:3])
-
-        # TODO: visualize each individual object
-        # debug
-        # print(subsampled_scene_xyz.shape)
-        # print(subsampled_pc_idxs.shape)
-        # print("visualize subsampled scene")
-        # for si in range(5):
-        #     trimesh.PointCloud(subsampled_scene_xyz[si, :, :3].cpu().numpy(), colors=[0, 0, 255, 255]).show()
-
-    ###############################################
-    # Create input for pairwise collision detector
-    if return_pair_pc:
-
-        assert num_pair_pc_pts is not None
-
-        # new_obj_xyzs: B, N, P, 3 + N
-        # target_object_inds: 1, N
-        # ignore paddings
-        num_objs = torch.sum(target_object_inds[0])
-        obj_pair_idxs = torch.combinations(torch.arange(num_objs), r=2)  # num_comb, 2
-
-        # use [:, :, :, :3] to get obj_xyzs without object-wise indicator
-        obj_pair_xyzs = new_obj_xyzs[:, :, :, :3][:, obj_pair_idxs]  # B, num_comb, 2 (obj 1 and obj 2), P, 3
-        num_comb = obj_pair_xyzs.shape[1]
-        pair_indicator_variables = torch.eye(2).repeat(B, num_comb, 1, 1, P).reshape(B, num_comb, 2, P, 2).to(device)  # B, num_comb, 2, P, 2
-        obj_pair_xyzs = torch.cat([obj_pair_xyzs, pair_indicator_variables], dim=-1)  # B, num_comb, 2, P, 3 (pc channels) + 2 (indicator for obj 1 and obj 2)
-        obj_pair_xyzs = obj_pair_xyzs.reshape(B, num_comb, P * 2, 5)
-
-        # random sample: idx = np.random.randint(0, scene_xyz.shape[0], self.num_scene_pts)
-        obj_pair_xyzs = obj_pair_xyzs.reshape(B * num_comb, P * 2, 5)
-        # random_point_sample() input dim: B, N, C
-        rand_idxs = random_point_sample(obj_pair_xyzs, num_pair_pc_pts)  # B * num_comb, num_pair_pc_pts
-        obj_pair_xyzs = index_points(obj_pair_xyzs, rand_idxs)  # B * num_comb, num_pair_pc_pts, 5
-
-        if normalize_pair_pc:
-            # pc_normalize_batch() input dim: pc: B, num_scene_pts, 3
-            # obj_pair_xyzs = obj_pair_xyzs.reshape(B * num_comb, num_pair_pc_pts, 5)
-            obj_pair_xyzs[:, :, 0:3] = pc_normalize_batch(obj_pair_xyzs[:, :, 0:3])
-            obj_pair_xyzs = obj_pair_xyzs.reshape(B, num_comb, num_pair_pc_pts, 5)
-
-            # # debug
-            # for bi, this_obj_pair_xyzs in enumerate(obj_pair_xyzs):
-            #     print("batch id", bi)
-            #     for pi, obj_pair_xyz in enumerate(this_obj_pair_xyzs):
-            #         print("pair", pi)
-            #         # obj_pair_xyzs: 2 * P, 5
-            #         print(obj_pair_xyz[:, :3].shape)
-            #         trimesh.PointCloud(obj_pair_xyz[:, :3].cpu()).show()
-
-    # obj_pair_xyzs: B, num_comb, num_pair_pc_pts, 3 + 2
-    goal_pc_pose = goal_pc_pose.reshape(B, N, 4, 4)
-
-    # TODO: update the return logic, a mess right now
-    if return_scene_pts_and_pc_idxs:
-        return subsampled_scene_xyz, subsampled_pc_idxs, new_obj_xyzs, goal_pc_pose
-
-    if return_pair_pc:
-        return subsampled_scene_xyz, new_obj_xyzs, goal_pc_pose, obj_pair_xyzs
-    else:
-        return subsampled_scene_xyz, new_obj_xyzs, goal_pc_pose
-
+import StructDiffusion.utils.tra3d as tra3d
 
 def move_pc_and_create_scene_new(obj_xyzs, obj_params, struct_pose, current_pc_pose, target_object_inds, device,
                                  return_scene_pts=False, return_scene_pts_and_pc_idxs=False, num_scene_pts=None, normalize_pc=False,
@@ -193,12 +28,17 @@ def move_pc_and_create_scene_new(obj_xyzs, obj_params, struct_pose, current_pc_p
     goal_pc_pose = struct_pose @ goal_pc_pose_in_struct
     goal_pc_transform = goal_pc_pose @ torch.inverse(current_pc_pose)  # cur_batch_size x N, 4, 4
 
-    # important: pytorch3d uses row-major ordering, need to transpose each transformation matrix
-    transpose = tra3d.Transform3d(matrix=goal_pc_transform.transpose(1, 2))
+    # # important: pytorch3d uses row-major ordering, need to transpose each transformation matrix
+    # transpose = tra3d.Transform3d(matrix=goal_pc_transform.transpose(1, 2))
+    # # obj_xyzs: N, P, 3
+    # new_obj_xyzs = obj_xyzs.repeat(B, 1, 1)
+    # new_obj_xyzs = transpose.transform_points(new_obj_xyzs)
+
+    # a verision that does not rely on pytorch3d
+    new_obj_xyzs = obj_xyzs.repeat(B, 1, 1)  # B x N, P, 3
+    new_obj_xyzs = torch.concat([new_obj_xyzs, torch.ones(B * N, P, 1).to(device)], dim=-1)  # B x N, P, 4
+    new_obj_xyzs = torch.einsum('bij,bkj->bki', goal_pc_transform, new_obj_xyzs)[:, :, :3]  # # B x N, P, 3
 
-    # obj_xyzs: N, P, 3
-    new_obj_xyzs = obj_xyzs.repeat(B, 1, 1)
-    new_obj_xyzs = transpose.transform_points(new_obj_xyzs)
 
     # put it back to B, N, P, 3
     new_obj_xyzs = new_obj_xyzs.reshape(B, N, P, -1)
@@ -332,82 +172,6 @@ def move_pc_and_create_scene_new(obj_xyzs, obj_params, struct_pose, current_pc_p
     return new_obj_xyzs, goal_pc_pose, subsampled_scene_xyz, subsampled_pc_idxs, obj_pair_xyzs
 
 
-def move_pc(obj_xyzs, obj_params, struct_pose, current_pc_pose, device):
-
-    # obj_xyzs: N, P, 3
-    # obj_params: B, N, 6
-    # struct_pose: B x N, 4, 4
-    # current_pc_pose: B x N, 4, 4
-    # target_object_inds: 1, N
-
-    B, N, _ = obj_params.shape
-    _, P, _ = obj_xyzs.shape
-
-    # B, N, 6
-    flat_obj_params = obj_params.reshape(B * N, -1)
-    goal_pc_pose_in_struct = torch.eye(4).repeat(B * N, 1, 1).to(device)
-    goal_pc_pose_in_struct[:, :3, :3] = tra3d.euler_angles_to_matrix(flat_obj_params[:, 3:], "XYZ")
-    goal_pc_pose_in_struct[:, :3, 3] = flat_obj_params[:, :3]  # B x N, 4, 4
-
-    goal_pc_pose = struct_pose @ goal_pc_pose_in_struct
-    goal_pc_transform = goal_pc_pose @ torch.inverse(current_pc_pose)  # cur_batch_size x N, 4, 4
-
-    # important: pytorch3d uses row-major ordering, need to transpose each transformation matrix
-    transpose = tra3d.Transform3d(matrix=goal_pc_transform.transpose(1, 2))
-
-    # obj_xyzs: N, P, 3
-    new_obj_xyzs = obj_xyzs.repeat(B, 1, 1)
-    new_obj_xyzs = transpose.transform_points(new_obj_xyzs)
-
-    # put it back to B, N, P, 3
-    new_obj_xyzs = new_obj_xyzs.reshape(B, N, P, -1)
-    # visualize_batch_pcs(new_obj_xyzs, S, N, P)
-
-    # subsampled_scene_xyz: B, num_scene_pts, 3+N
-    # new_obj_xyzs: B, N, P, 3
-    # goal_pc_pose: B, N, 4, 4
-
-    goal_pc_pose = goal_pc_pose.reshape(B, N, 4, 4)
-    return new_obj_xyzs, goal_pc_pose
-
-
-def move_pc_and_create_scene_simple(obj_xyzs, struct_pose, pc_poses_in_struct):
-
-    device = obj_xyzs.device
-
-    # obj_xyzs: B, N, P, 3 or 6
-    # struct_pose: B, 1, 4, 4
-    # pc_poses_in_struct: B, N, 4, 4
-
-    B, N, _, _ = pc_poses_in_struct.shape
-    _, _, P, _ = obj_xyzs.shape
-
-    current_pc_poses = torch.eye(4).repeat(B, N, 1, 1).to(device)  # B, N, 4, 4
-    # print(torch.mean(obj_xyzs, dim=2).shape)
-    current_pc_poses[:, :, :3, 3] = torch.mean(obj_xyzs[:, :, :, :3], dim=2)  # B, N, 4, 4
-    current_pc_poses = current_pc_poses.reshape(B * N, 4, 4)  # B x N, 4, 4
-
-    struct_pose = struct_pose.repeat(1, N, 1, 1) # B, N, 4, 4
-    struct_pose = struct_pose.reshape(B * N, 4, 4)  # B x 1, 4, 4
-    pc_poses_in_struct = pc_poses_in_struct.reshape(B * N, 4, 4)  # B x N, 4, 4
-
-    goal_pc_pose = struct_pose @ pc_poses_in_struct  # B x N, 4, 4
-    # print("goal pc poses")
-    # print(goal_pc_pose)
-    goal_pc_transform = goal_pc_pose @ torch.inverse(current_pc_poses)  # B x N, 4, 4
-
-    # important: pytorch3d uses row-major ordering, need to transpose each transformation matrix
-    transpose = tra3d.Transform3d(matrix=goal_pc_transform.transpose(1, 2))
-
-    new_obj_xyzs = obj_xyzs.reshape(B * N, P, -1)  # B x N, P, 3
-    new_obj_xyzs[:, :, :3] = transpose.transform_points(new_obj_xyzs[:, :, :3])
-
-    # put it back to B, N, P, 3
-    new_obj_xyzs = new_obj_xyzs.reshape(B, N, P, -1)
-
-    return new_obj_xyzs
-
-
 def compute_current_and_goal_pc_poses(obj_xyzs, struct_pose, pc_poses_in_struct):
 
     device = obj_xyzs.device

src/StructDiffusion/utils/files.py

@@ -1,9 +1,17 @@
 import os
 
+
 def get_checkpoint_path_from_dir(checkpoint_dir):
     checkpoint_path = None
     for file in os.listdir(checkpoint_dir):
         if "ckpt" in file:
             checkpoint_path = os.path.join(checkpoint_dir, file)
     assert checkpoint_path is not None
-    return checkpoint_path
+    return checkpoint_path
+
+
+def replace_config_for_testing_data(cfg, testing_data_cfg):
+    cfg.DATASET.data_roots = testing_data_cfg.DATASET.data_roots
+    cfg.DATASET.index_roots = testing_data_cfg.DATASET.index_roots
+    cfg.DATASET.vocab_dir = testing_data_cfg.DATASET.vocab_dir
+

src/StructDiffusion/utils/np_speed_test.py

@@ -0,0 +1,41 @@
+import numpy as np
+import time
+
+def numpy_test(w, h, r):
+    v, u = np.indices((h, w))
+    theta = (u - w / 2.) * 2 * np.pi / w
+    phi = (v - h / 2.) * np.pi / h
+    cos_phi = np.cos(phi)
+    x = cos_phi * np.sin(theta)
+    y = np.sin(phi)
+    z = cos_phi * np.cos(theta)
+    
+    ray = np.dstack((x, y, z))
+    ray = ray.reshape(-1, 3).dot(r.T)
+    ray.shape = (h, w, 3)
+    
+    x, y, z = np.dsplit(ray, 3)
+    theta = np.arctan2(x, z)
+    phi = np.arcsin(y)
+    u = theta * w / 2 / np.pi + w / 2.
+    v = phi * h / np.pi + h / 2.
+    xymap = np.dstack((u, v)).astype(np.float32)
+    
+    return xymap
+
+def matrix_multiplication():
+    for i in range(100):
+        np.random.random((1000, 1000)) @ np.random.random((1000, 1000))
+
+
+if __name__ == "__main__":
+    w = 3584
+    h = int(w / 2)
+    r = np.array([
+    [0.61566148, -0.78369395, 0.08236955],
+    [0.78801075, 0.61228882, -0.06435415],
+    [0, 0.10452846, 0.9945219],])
+    begin_time = time.time()
+    # numpy_test(w, h, r)
+    matrix_multiplication()
+    print(time.time() - begin_time)

src/StructDiffusion/utils/rearrangement.py

@@ -558,9 +558,12 @@ def fit_gaussians(samples, sigma_eps=0.01):
     return mus, sigmas
 
 
-def show_pcs_with_trimesh(obj_xyzs, obj_rgbs, return_scene=False):
-    vis_pcs = [trimesh.PointCloud(obj_xyz, colors=np.concatenate([obj_rgb * 255, np.ones([obj_rgb.shape[0], 1]) * 255], axis=-1)) for
-               obj_xyz, obj_rgb in zip(obj_xyzs, obj_rgbs)]
+def show_pcs_with_trimesh(obj_xyzs, obj_rgbs=None, return_scene=False):
+    if obj_rgbs is not None:
+        vis_pcs = [trimesh.PointCloud(obj_xyz, colors=np.concatenate([obj_rgb * 255, np.ones([obj_rgb.shape[0], 1]) * 255], axis=-1)) for
+                   obj_xyz, obj_rgb in zip(obj_xyzs, obj_rgbs)]
+    else:
+        vis_pcs = [trimesh.PointCloud(obj_xyz) for obj_xyz in obj_xyzs]
     scene = trimesh.Scene()
     # add the coordinate frame first
     geom = trimesh.creation.axis(0.01)
@@ -582,13 +585,35 @@ def show_pcs_with_trimesh(obj_xyzs, obj_rgbs, return_scene=False):
     RT_4x4 = np.linalg.inv(RT_4x4)
     RT_4x4 = RT_4x4 @ np.diag([1, -1, -1, 1])
     scene.camera_transform = RT_4x4
-
     if return_scene:
         return scene
     else:
         scene.show()
 
 
+def get_trimesh_scene_with_table():
+    scene = trimesh.Scene()
+    # add the coordinate frame first
+    geom = trimesh.creation.axis(0.01)
+    scene.add_geometry(geom)
+    table = trimesh.creation.box(extents=[1.0, 1.0, 0.02])
+    table.apply_translation([0.5, 0, -0.01])
+    table.visual.vertex_colors = [150, 111, 87, 125]
+    scene.add_geometry(table)
+    # bounds = trimesh.creation.box(extents=[4.0, 4.0, 4.0])
+    bounds = trimesh.creation.icosphere(subdivisions=3, radius=3.1)
+    bounds.apply_translation([0, 0, 0])
+    bounds.visual.vertex_colors = [30, 30, 30, 30]
+    # scene.add_geometry(bounds)
+    RT_4x4 = np.array([[-0.39560353822208355, -0.9183993826406329, 0.006357240869497738, 0.2651463080169481],
+                       [-0.797630370081598, 0.3401340617616391, -0.4980909683511864, 0.2225696480721997],
+                       [0.45528412367406523, -0.2021172778236285, -0.8671014777611122, 0.9449050652025951],
+                       [0.0, 0.0, 0.0, 1.0]])
+    RT_4x4 = np.linalg.inv(RT_4x4)
+    RT_4x4 = RT_4x4 @ np.diag([1, -1, -1, 1])
+    scene.camera_transform = RT_4x4
+    return scene
+
 def show_pcs_with_predictions(xyz, rgb, gts, predictions, add_coordinate_frame=False, return_buffer=False, add_table=True, side_view=True):
     """ Display point clouds """
 

src/StructDiffusion/utils/tra3d.py

@@ -0,0 +1,148 @@
+import torch
+
+
+# source: https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html#matrix_to_euler_angles
+# we don't want to build pytorch3d, so only pick functions we need to use
+
+def _index_from_letter(letter: str) -> int:
+    if letter == "X":
+        return 0
+    if letter == "Y":
+        return 1
+    if letter == "Z":
+        return 2
+    raise ValueError("letter must be either X, Y or Z.")
+
+
+def _angle_from_tan(
+    axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool
+) -> torch.Tensor:
+    """
+    Extract the first or third Euler angle from the two members of
+    the matrix which are positive constant times its sine and cosine.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z" for the angle we are finding.
+        other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
+            convention.
+        data: Rotation matrices as tensor of shape (..., 3, 3).
+        horizontal: Whether we are looking for the angle for the third axis,
+            which means the relevant entries are in the same row of the
+            rotation matrix. If not, they are in the same column.
+        tait_bryan: Whether the first and third axes in the convention differ.
+
+    Returns:
+        Euler Angles in radians for each matrix in data as a tensor
+        of shape (...).
+    """
+
+    i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
+    if horizontal:
+        i2, i1 = i1, i2
+    even = (axis + other_axis) in ["XY", "YZ", "ZX"]
+    if horizontal == even:
+        return torch.atan2(data[..., i1], data[..., i2])
+    if tait_bryan:
+        return torch.atan2(-data[..., i2], data[..., i1])
+    return torch.atan2(data[..., i2], -data[..., i1])
+
+
+def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor:
+    """
+    Return the rotation matrices for one of the rotations about an axis
+    of which Euler angles describe, for each value of the angle given.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z".
+        angle: any shape tensor of Euler angles in radians
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+
+    cos = torch.cos(angle)
+    sin = torch.sin(angle)
+    one = torch.ones_like(angle)
+    zero = torch.zeros_like(angle)
+
+    if axis == "X":
+        R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
+    elif axis == "Y":
+        R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
+    elif axis == "Z":
+        R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
+    else:
+        raise ValueError("letter must be either X, Y or Z.")
+
+    return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))
+
+
+def matrix_to_euler_angles(matrix: torch.Tensor, convention: str) -> torch.Tensor:
+    """
+    Convert rotations given as rotation matrices to Euler angles in radians.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+        convention: Convention string of three uppercase letters.
+
+    Returns:
+        Euler angles in radians as tensor of shape (..., 3).
+    """
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
+    i0 = _index_from_letter(convention[0])
+    i2 = _index_from_letter(convention[2])
+    tait_bryan = i0 != i2
+    if tait_bryan:
+        central_angle = torch.asin(
+            matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0)
+        )
+    else:
+        central_angle = torch.acos(matrix[..., i0, i0])
+
+    o = (
+        _angle_from_tan(
+            convention[0], convention[1], matrix[..., i2], False, tait_bryan
+        ),
+        central_angle,
+        _angle_from_tan(
+            convention[2], convention[1], matrix[..., i0, :], True, tait_bryan
+        ),
+    )
+    return torch.stack(o, -1)
+
+
+def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) -> torch.Tensor:
+    """
+    Convert rotations given as Euler angles in radians to rotation matrices.
+
+    Args:
+        euler_angles: Euler angles in radians as tensor of shape (..., 3).
+        convention: Convention string of three uppercase letters from
+            {"X", "Y", and "Z"}.
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
+        raise ValueError("Invalid input euler angles.")
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    matrices = [
+        _axis_angle_rotation(c, e)
+        for c, e in zip(convention, torch.unbind(euler_angles, -1))
+    ]
+    # return functools.reduce(torch.matmul, matrices)
+    return torch.matmul(torch.matmul(matrices[0], matrices[1]), matrices[2])

wandb_logs/StructDiffusion/CollisionDiscriminator/checkpoints/epoch=199-step=653400.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed90c4976e69f96324d0245ecc635fa987d66d84c512d1e2105ce9f7f3df39ea
+size 34533413

wandb_logs/StructDiffusion/ConditionalPoseDiffusionLanguage/checkpoints/epoch=199-step=100000.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:edd6c08c5fafd129f365fd70fc1b6ad68e643a9111e522432279afb7ba387a89
+size 59947673