Add TAXPoseD #18

oadonca · 2023-12-15T23:46:09Z

Add TAXPoseD with support for base TAXPose decoder and RLBench training.

Only changes that should possibly affect non TAXPoseD code are the changes to scripts/train_residual_flow.py, taxpose/nets/transformer_flow.py, and third_party/dcp/model.

Currently using a separate base config file (configs/train_ndf_multimodal.yaml), let me know if you would prefer a solution that doesn't need this.

…split_next

beneisner

Looks good. try running the unit tests and formatter

pytest

and

black (forget how this works...)

beneisner · 2023-12-18T04:46:06Z

third_party/dcp/model.py

this file looks kind of weird, no? why is the diff so scattered.... any way to make the diff clearer?

Not sure, will move these changes into taxpose/nets which should also have a cleaner diff

Moved these changes out of third_party/... and into a new file in taxpose/nets/taxposed_dgcnn.py, maybe a better file name exists, could also probably merge the other DGCNN implementations (DGCNN_GC, VN_DGCNN) into a single custom dgcnns file in the future.

beneisner · 2023-12-18T04:55:19Z

third_party/dcp/model.py

@@ -280,18 +281,100 @@ def forward(self, x):


 class DGCNN(nn.Module):
-    def __init__(self, emb_dims=512):
+    def __init__(self, emb_dims=512, input_dims=3, num_heads=1, conditioning_size=0, last_relu=True):


two thoughts:

In general, I'm trying to keep this file as close as possible to the original code we took from the DCP paper. In its current state it's a true third-party import, and changing it kind of contaminates that. If you want to add additional models, put them in taxpose/nets, and mention at the top that it's derived from this file. I know it's a bit clunky, but I really want to make sure we keep track of various low-level modifications to networks or code we borrow.

What is the motivation behind teh changes to DGCNN here? is it to just make it conditional? in my mind that's a pretty distinct (custom) model that should live in taxpose/nets. Are there other logical changes? i.e. last_relu...

Makes complete sense, will move these changes into a new class in taxpose/nets

Motivation is primarily to handle conditioning, either through an extra channels in the input or added to the later conv layers. Additionally, supports having multiple heads when TAXPoseD vae encoder needs to predict a normal distribution.

beneisner · 2023-12-18T04:55:48Z

taxpose/utils/se3.py

-    )
-    constrained_axix_angle = rot_ratio * axis_angle_random  # max angle is rot_var
-    R = axis_angle_to_matrix(constrained_axix_angle)
+def random_se3(N, rot_var=np.pi/180 * 5, trans_var=0.1, 


nit: this won't pass black formatting

reverted this change as i realized it was unnecessary

beneisner · 2023-12-18T04:58:09Z

taxpose/nets/transformer_flow.py

        if emb_nn == "dgcnn":
-            self.emb_nn_action = DGCNN(emb_dims=self.emb_dims)
-            self.emb_nn_anchor = DGCNN(emb_dims=self.emb_dims)
+            self.emb_nn_action = DGCNN(emb_dims=self.emb_dims, input_dims=self.input_dims, conditioning_size=self.conditioning_size)


not sure if i'd prefer a different case here - like c-dgcnn or something.

Understood, will add a new case for this

added a case for "cond_dgcnn"

beneisner · 2023-12-18T04:58:59Z

taxpose/nets/transformer_flow.py

+        action_points = input[0].permute(0, 2, 1)[:, :self.input_dims]  # B,3,num_points
+        anchor_points = input[1].permute(0, 2, 1)[:, :self.input_dims]
+
+        if action_center is None:


what is action_center / anchor_center here?

taxposed vae encoder will output distributions over the action/anchor point clouds which we sample a point from each to encode the demo. Jenny's found some benefit to centering the point clouds about the sampled points instead of mean, so on the taxposed side we pass in either the mean or sampled points for action/anchor point cloud centering as action_center/anchor_center

beneisner · 2023-12-18T04:59:27Z

taxpose/nets/transformer_flow.py

@@ -606,8 +629,8 @@ def forward(self, *input):
        head_action_output = self.head_action(
            action_embedding_tf,
            anchor_embedding_tf,
-            action_points,
-            anchor_points,
+            action_points[:, :3, :],


what gets reshaped to make this necessary?

taxposed action_points/anchor_points will be [Batch size, XYZ + Conditioning size, Num. points], we want the full size inputs for the dgcnn embeddings but for the heads it seems we only need XYZ, additionally seems necessary for dimensions to align.

beneisner · 2023-12-18T05:00:07Z

taxpose/training/multimodal_flow_equivariance_training_module_nocentering.py

this is fine, totally isolated from other code in the repo.

beneisner · 2023-12-18T05:03:33Z

scripts/train_residual_flow.py

-        softmax_temperature=cfg.task.phase.softmax_temperature,
-        flow_supervision=cfg.training.flow_supervision,
-    )
+    if cfg.model.name in ["taxposed", "taxposed_mlat_s100", "taxposed_mlat_s256"]:


might want to extract this conditional into a function above. just so that it's clear what's actually being set up in the main function. (i.e. a network, a model, a dataset, and training).

created functions for network creation, model creation, and maybe loading weights so main function is clear

beneisner · 2023-12-18T05:09:10Z

Seems pretty cool - can you post links to some training runs?

Curious to see if the mlat stuff worked as well!

also in general, the more you can isolate changes for these codepaths the better (which is why i commented only on files that touch normal taxpose training). I'm a bit paranoid about logic/structural changes to existing code which is known to work w/o running full training + eval on a few tasks. So the less that happens, the easier it is to say that nothing meaningful has changed about the code.

But in general I'm really impressed with how quickly you got this going. The config complexity alone is annoying enough as it is, so good work figuring that out so fast!

beneisner · 2023-12-18T05:12:08Z

Also, can you put some example commands for how to train this thing in scripts/README.md?

You might even try to get this running on the cluster o.O

oadonca · 2023-12-19T17:59:31Z

Updated code with black formatting and ran pytests, only cases that fail are due to missing checkpoint configs, mainly for evaluations and loading pretrained weights.

Ran the first training stage for TAXPoseD on the RLBench tasks with and without mlat: https://wandb.ai/r-pad/taxposed
Will start second training stage for these shortly but the above should give an indication for a kind of upper bound on performance.

Adding back NDF dataset support with some feature additions

…s, add taxposed launch commands to readme

beneisner added 30 commits July 26, 2023 19:17

non-corl additions

25a8a42

corl additions

bb03794

refactor the point cloud dataset so it can be swapped out

5712e5c

partially complete dataset

0d044c1

replaced with cfg stuff

db16b37

enum

0f9fe29

first version of training rlbench

2b82315

changes to some downloads

a29e57b

small changes to pretraining

b3027e4

pretraining works for all objects now

7cf195a

initial attempt at ndf pretraining

a44f73d

add a bunch of commands and configs

9f213ed

launch script

09bce40

amend

57c63d2

additional stuff

7976d19

lots of extra goodies...

dd8e36c

got training and downloading going

88c2ed3

added bottle place

be99025

added remaining things

98456c7

add open3d dependency

2656e87

add pretraining options

3d99175

add in a config file

60ce4ee

add in a config file

b7789b5

fix bottle config

693a3ab

fix bottle config

3d3f65a

add taxpose training scripts

6c33b7d

fix logic bug

9a10a06

Merge remote-tracking branch 'origin/rlbench_dataset' into corl_2023_…

2048cad

…split_next

add in multilateration

e3c0d11

add new commands for cluster

3ec0905

beneisner added 14 commits November 20, 2023 16:04

update the config to use max points

65892ad

add puzzle training

522a7e8

added new tasks

245831b

add evaluations

4bc2f6a

ablation

3bb2227

iclr 2024 rebuttal

4033a25

pick and place works well

18a42aa

eval works

8f1c54d

pytest works

84fae64

move rlbench around

aeb0a73

move train

4392445

configs look good

d6ffabe

huge change to the config structure to enable better scaling to rlbench

7eef0ea

nice new training commands

48ca396

oadonca requested review from beneisner and himty December 15, 2023 23:46

oadonca changed the title ~~Add TAXPoseD with support for base TAXPose decoder and RLBench~~ Add TAXPoseD Dec 18, 2023

beneisner requested changes Dec 18, 2023

View reviewed changes

oadonca and others added 7 commits December 20, 2023 21:09

Adding back NDF dataset support with some feature additions

d114cc1

Use occlusion_cfg itself to check whether to apply occlusions

9b7332e

Merge pull request #19 from r-pad/add_back_ndf

96b0f0d

Adding back NDF dataset support with some feature additions

fix rlbench eval

9af0790

Add TAXPoseD with support for base TAXPose decoder and RLBench

3ec7a83

Add TAXPoseD support for RelDist decoder

a8e0291

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

062ffb0

…s, add taxposed launch commands to readme

oadonca force-pushed the rlbench_taxposed branch from f4e7f48 to 062ffb0 Compare December 27, 2023 00:07

Base automatically changed from back_to_rlbench to main May 14, 2024 21:05

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add TAXPoseD #18

Add TAXPoseD #18

oadonca commented Dec 15, 2023

beneisner left a comment

beneisner Dec 18, 2023

oadonca Dec 18, 2023

oadonca Dec 19, 2023

beneisner Dec 18, 2023

oadonca Dec 18, 2023

beneisner Dec 18, 2023

oadonca Dec 19, 2023

beneisner Dec 18, 2023

oadonca Dec 18, 2023

oadonca Dec 19, 2023

beneisner Dec 18, 2023

oadonca Dec 18, 2023

beneisner Dec 18, 2023

oadonca Dec 18, 2023

beneisner Dec 18, 2023

beneisner Dec 18, 2023

oadonca Dec 19, 2023

beneisner commented Dec 18, 2023

beneisner commented Dec 18, 2023

oadonca commented Dec 19, 2023

Add TAXPoseD #18

Are you sure you want to change the base?

Add TAXPoseD #18

Conversation

oadonca commented Dec 15, 2023

beneisner left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

beneisner commented Dec 18, 2023

beneisner commented Dec 18, 2023

oadonca commented Dec 19, 2023