Revert some changes, apply black formatting, add more taxposed config…

…s, add taxposed launch commands to readme

configs/commands/rlbench/phone_on_base/train_taxposed_mlat_pzX_place.yaml

@@ -0,0 +1,19 @@
+# @package _global_
+
+defaults:
+  - /commands/rlbench/_train_taxposed@_here_
+  - override /model: taxposed_mlat_s100
+  - override /task:  phone_on_base
+  - override /phase: place
+  - _self_
+
+break_symmetry: False
+
+training:
+  init_cond_x: True
+  freeze_embnn: True
+  freeze_residual_flow: True
+  freeze_z_embnn: True
+
+  load_from_checkpoint: True
+  checkpoint_file: null

configs/commands/rlbench/phone_on_base/train_taxposed_mlat_pzY_place.yaml

@@ -0,0 +1,10 @@
+# @package _global_
+
+defaults:
+  - /commands/rlbench/_train_taxposed@_here_
+  - override /model: taxposed_mlat_s100
+  - override /task:  phone_on_base
+  - override /phase: place
+  - _self_
+
+break_symmetry: False

configs/commands/rlbench/place_hanger_on_rack/train_taxposed_mlat_pzX_place.yaml

@@ -0,0 +1,23 @@
+# @package _global_
+
+defaults:
+  - /commands/rlbench/_train_taxposed@_here_
+  - override /model: taxposed_mlat_s100
+  - override /task:  place_hanger_on_rack
+  - override /phase: place
+  - _self_
+
+break_symmetry: False
+
+training:
+  init_cond_x: True
+  freeze_embnn: True
+  freeze_residual_flow: True
+  freeze_z_embnn: True
+
+  load_from_checkpoint: True
+  checkpoint_file: null
+
+dm:
+  train_dset:
+    num_points: 256

configs/commands/rlbench/place_hanger_on_rack/train_taxposed_mlat_pzY_place.yaml

@@ -0,0 +1,14 @@
+# @package _global_
+
+defaults:
+  - /commands/rlbench/_train_taxposed@_here_
+  - override /model: taxposed_mlat_s100
+  - override /task:  place_hanger_on_rack
+  - override /phase: place
+  - _self_
+
+break_symmetry: False
+
+dm:
+  train_dset:
+    num_points: 256

configs/commands/rlbench/put_toilet_roll_on_stand/train_taxposed_mlat_pzX_place.yaml

@@ -0,0 +1,19 @@
+# @package _global_
+
+defaults:
+  - /commands/rlbench/_train_taxposed@_here_
+  - override /model: taxposed_mlat_s100
+  - override /task:  put_toilet_roll_on_stand
+  - override /phase: place
+  - _self_
+
+break_symmetry: False
+
+training:
+  init_cond_x: True
+  freeze_embnn: True
+  freeze_residual_flow: True
+  freeze_z_embnn: True
+
+  load_from_checkpoint: True
+  checkpoint_file: null

configs/commands/rlbench/put_toilet_roll_on_stand/train_taxposed_mlat_pzY_place.yaml

@@ -0,0 +1,10 @@
+# @package _global_
+
+defaults:
+  - /commands/rlbench/_train_taxposed@_here_
+  - override /model: taxposed_mlat_s100
+  - override /task:  put_toilet_roll_on_stand
+  - override /phase: place
+  - _self_
+
+break_symmetry: False

configs/commands/rlbench/solve_puzzle/train_taxposed_pzX_grasp.yaml → configs/commands/rlbench/solve_puzzle/train_taxposed_mlat_pzX_place.yaml

@@ -2,9 +2,9 @@
 
 defaults:
   - /commands/rlbench/_train_taxposed@_here_
-  - override /model: taxposed
+  - override /model: taxposed_mlat_s100
   - override /task:  solve_puzzle
-  - override /phase: grasp
+  - override /phase: place
   - _self_
 
 break_symmetry: False

configs/commands/rlbench/solve_puzzle/train_taxposed_pzY_grasp.yaml → configs/commands/rlbench/solve_puzzle/train_taxposed_mlat_pzY_place.yaml

@@ -2,9 +2,9 @@
 
 defaults:
   - /commands/rlbench/_train_taxposed@_here_
-  - override /model: taxposed
+  - override /model: taxposed_mlat_s100
   - override /task:  solve_puzzle
-  - override /phase: grasp
+  - override /phase: place
   - _self_
 
 break_symmetry: False

configs/commands/rlbench/stack_wine/train_taxposed_mlat_pzX_place.yaml

@@ -0,0 +1,19 @@
+# @package _global_
+
+defaults:
+  - /commands/rlbench/_train_taxposed@_here_
+  - override /model: taxposed_mlat_s100
+  - override /task:  stack_wine
+  - override /phase: place
+  - _self_
+
+break_symmetry: False
+
+training:
+  init_cond_x: True
+  freeze_embnn: True
+  freeze_residual_flow: True
+  freeze_z_embnn: True
+
+  load_from_checkpoint: True
+  checkpoint_file: null

configs/commands/rlbench/stack_wine/train_taxposed_mlat_pzY_place.yaml

@@ -0,0 +1,10 @@
+# @package _global_
+
+defaults:
+  - /commands/rlbench/_train_taxposed@_here_
+  - override /model: taxposed_mlat_s100
+  - override /task:  stack_wine
+  - override /phase: place
+  - _self_
+
+break_symmetry: False

configs/model/_taxposed.yaml

@@ -3,7 +3,7 @@ name: ???
 conditioning: pos_delta_l2norm
 
 emb_dims: 512
-emb_nn: dgcnn
+emb_nn: cond_dgcnn
 return_flow_component: True
 center_feature: True
 inital_sampling_ratio: 1
@@ -13,6 +13,7 @@ mlat_sample: null
 mlat_nkps: null
 break_symmetry: False
 latent_z_linear_size: 40
+num_points: 1024
 
 gumbel_temp: 1
 division_smooth_factor: 1

configs/train_ndf.yaml

@@ -45,11 +45,13 @@ dm:
 training:
   batch_size: 8
   max_epochs: 500
+  max_steps: -1
 
   sigmoid_on: True
 
   # Optimizer Settings
   lr: 1e-4
+  gradient_clipping: null
 
   # Loss Settings
   flow_supervision: both

configs/train_ndf_multimodal.yaml

@@ -44,7 +44,7 @@ dm:
 
 training:
   batch_size: 8
-  max_epochs: 20
+  max_epochs: 500
   max_steps: -1
 
   sigmoid_on: True
@@ -82,7 +82,7 @@ training:
   checkpoint_file: null
 
   # Visualization Settings
-  image_logging_period: 1001
+  image_logging_period: 100
   log_every_n_steps: 100
   check_val_every_n_epoch: 5
 

scripts/README.md

@@ -85,6 +85,63 @@ If you write some scripts which are meant to be run stand-alone, and not importe
 ./launch_autobot.sh 4 python scripts/train_residual_flow.py --config-name commands/rlbench/train_mlat_rlbench_solve_puzzle_place.yaml wandb.group=rlbench_mlat resources.num_workers=0
 ./launch_autobot.sh 5 python scripts/train_residual_flow.py --config-name commands/rlbench/train_mlat_rlbench_place_hanger_on_rack_place.yaml wandb.group=rlbench_mlat resources.num_workers=0
 
+### TAXPoseD 
+#### TAXPoseD p(z|Y)
+
+./launch.sh ${RPAD_PLATFORM} 0 python scripts/train_residual_flow.py --config-name commands/rlbench/stack_wine/train_taxposed_pzY_place.yaml wandb.group=rlbench_taxposed
+
+./launch.sh ${RPAD_PLATFORM} 1 python scripts/train_residual_flow.py --config-name commands/rlbench/insert_onto_square_peg/train_taxposed_pzY_place.yaml wandb.group=rlbench_taxposed
+
+./launch.sh ${RPAD_PLATFORM} 2 python scripts/train_residual_flow.py --config-name commands/rlbench/phone_on_base/train_taxposed_pzY_place.yaml wandb.group=rlbench_taxposed
+
+./launch.sh ${RPAD_PLATFORM} 3 python scripts/train_residual_flow.py --config-name commands/rlbench/put_toilet_roll_on_stand/train_taxposed_pzY_place.yaml wandb.group=rlbench_taxposed
+
+./launch.sh ${RPAD_PLATFORM} 4 python scripts/train_residual_flow.py --config-name commands/rlbench/place_hanger_on_rack/train_taxposed_pzY_place.yaml wandb.group=rlbench_taxposed
+
+./launch.sh ${RPAD_PLATFORM} 5 python scripts/train_residual_flow.py --config-name commands/rlbench/solve_puzzle/train_taxposed_pzY_place.yaml wandb.group=rlbench_taxposed
+
+#### TAXPoseD p(z|X)
+
+./launch.sh ${RPAD_PLATFORM} 0 python scripts/train_residual_flow.py --config-name commands/rlbench/stack_wine/train_taxposed_pzX_place.yaml wandb.group=rlbench_taxposed training.checkpoint_file=<stack_wine p(z|Y) checkpoint>
+
+./launch.sh ${RPAD_PLATFORM} 1 python scripts/train_residual_flow.py --config-name commands/rlbench/insert_onto_square_peg/train_taxposed_pzX_place.yaml wandb.group=rlbench_taxposed training.checkpoint_file=<insert_onto_square_peg p(z|Y) checkpoint>
+
+./launch.sh ${RPAD_PLATFORM} 2 python scripts/train_residual_flow.py --config-name commands/rlbench/phone_on_base/train_taxposed_pzX_place.yaml wandb.group=rlbench_taxposed training.checkpoint_file=<phone_on_base p(z|Y) checkpoint>
+
+./launch.sh ${RPAD_PLATFORM} 3 python scripts/train_residual_flow.py --config-name commands/rlbench/put_toilet_roll_on_stand/train_taxposed_pzX_place.yaml wandb.group=rlbench_taxposed training.checkpoint_file=<put_toilet_roll_on_stand p(z|Y) checkpoint>
+
+./launch.sh ${RPAD_PLATFORM} 4 python scripts/train_residual_flow.py --config-name commands/rlbench/place_hanger_on_rack/train_taxposed_pzX_place.yaml wandb.group=rlbench_taxposed training.checkpoint_file=<place_hanger_on_rack p(z|Y) checkpoint>
+
+./launch.sh ${RPAD_PLATFORM} 5 python scripts/train_residual_flow.py --config-name commands/rlbench/solve_puzzle/train_taxposed_pzX_place.yaml wandb.group=rlbench_taxposed training.checkpoint_file=<solve_puzzle p(z|Y) checkpoint>
+
+#### TAXPoseD Mlat p(z|Y)
+
+./launch.sh ${RPAD_PLATFORM} 0 python scripts/train_residual_flow.py --config-name commands/rlbench/stack_wine/train_taxposed_mlat_pzY_place.yaml wandb.group=rlbench_taxposed_mlat
+
+./launch.sh ${RPAD_PLATFORM} 1 python scripts/train_residual_flow.py --config-name commands/rlbench/insert_onto_square_peg/train_taxposed_mlat_pzY_place.yaml wandb.group=rlbench_taxposed_mlat
+
+./launch.sh ${RPAD_PLATFORM} 2 python scripts/train_residual_flow.py --config-name commands/rlbench/phone_on_base/train_taxposed_mlat_pzY_place.yaml wandb.group=rlbench_taxposed_mlat
+
+./launch.sh ${RPAD_PLATFORM} 3 python scripts/train_residual_flow.py --config-name commands/rlbench/put_toilet_roll_on_stand/train_taxposed_mlat_pzY_place.yaml wandb.group=rlbench_taxposed_mlat
+
+./launch.sh ${RPAD_PLATFORM} 4 python scripts/train_residual_flow.py --config-name commands/rlbench/place_hanger_on_rack/train_taxposed_mlat_pzY_place.yaml wandb.group=rlbench_taxposed_mlat
+
+./launch.sh ${RPAD_PLATFORM} 5 python scripts/train_residual_flow.py --config-name commands/rlbench/solve_puzzle/train_taxposed_mlat_pzY_place.yaml wandb.group=rlbench_taxposed_mlat
+
+#### TAXPoseD Mlat p(z|X)
+
+./launch.sh ${RPAD_PLATFORM} 0 python scripts/train_residual_flow.py --config-name commands/rlbench/stack_wine/train_taxposed_mlat_pzX_place.yaml wandb.group=rlbench_taxposed_mlat training.checkpoint_file=<stack_wine p(z|Y) checkpoint>
+
+./launch.sh ${RPAD_PLATFORM} 1 python scripts/train_residual_flow.py --config-name commands/rlbench/insert_onto_square_peg/train_taxposed_mlat_pzX_place.yaml wandb.group=rlbench_taxposed_mlat training.checkpoint_file=<insert_onto_square_peg p(z|Y) checkpoint>
+
+./launch.sh ${RPAD_PLATFORM} 2 python scripts/train_residual_flow.py --config-name commands/rlbench/phone_on_base/train_taxposed_mlat_pzX_place.yaml wandb.group=rlbench_taxposed_mlat training.checkpoint_file=<phone_on_base p(z|Y) checkpoint>
+
+./launch.sh ${RPAD_PLATFORM} 3 python scripts/train_residual_flow.py --config-name commands/rlbench/put_toilet_roll_on_stand/train_taxposed_mlat_pzX_place.yaml wandb.group=rlbench_taxposed_mlat training.checkpoint_file=<put_toilet_roll_on_stand p(z|Y) checkpoint>
+
+./launch.sh ${RPAD_PLATFORM} 4 python scripts/train_residual_flow.py --config-name commands/rlbench/place_hanger_on_rack/train_taxposed_mlat_pzX_place.yaml wandb.group=rlbench_taxposed_mlat training.checkpoint_file=<place_hanger_on_rack p(z|Y) checkpoint>
+
+./launch.sh ${RPAD_PLATFORM} 5 python scripts/train_residual_flow.py --config-name commands/rlbench/solve_puzzle/train_taxposed_mlat_pzX_place.yaml wandb.group=rlbench_taxposed_mlat training.checkpoint_file=<solve_puzzle p(z|Y) checkpoint>
+
 
 ### Ablations