├── configs ├── local │ └── .gitkeep ├── callbacks │ ├── none.yaml │ ├── device_stats_monitor.yaml │ ├── rich_progress_bar.yaml │ ├── lr_monitor.yaml │ ├── model_summary.yaml │ ├── stochastic_weight_averaging.yaml │ ├── default.yaml │ ├── early_stopping.yaml │ └── model_checkpoint.yaml ├── trainer │ ├── cpu.yaml │ ├── gpu.yaml │ ├── mps.yaml │ ├── ddp_sim.yaml │ ├── ddp.yaml │ └── default.yaml ├── __init__.py ├── debug │ ├── fdr.yaml │ ├── profiler.yaml │ ├── overfit.yaml │ ├── limit.yaml │ └── default.yaml ├── logger │ ├── csv.yaml │ ├── many_loggers.yaml │ ├── tensorboard.yaml │ ├── neptune.yaml │ ├── mlflow.yaml │ ├── comet.yaml │ ├── wandb.yaml │ └── aim.yaml ├── extras │ └── default.yaml ├── exp_rlbench_act_policy │ ├── rlbench_model │ │ ├── scratch_vit_rgb.yaml │ │ ├── pretrained_vc1_rgb.yaml │ │ ├── pretrained_r3m_rgb.yaml │ │ ├── scratch_vit_rgbd.yaml │ │ ├── scratch_resnet50_rgb.yaml │ │ ├── scratch_pointnet_pcd.yaml │ │ ├── pretrained_multimae_rgbd.yaml │ │ ├── scratch_resnet50_rgbd.yaml │ │ ├── scratch_spunet_pcd.yaml │ │ ├── scratch_multivit_rgbd.yaml │ │ └── pretrained_ponderv2_pcd.yaml │ └── base.yaml ├── data │ ├── maniskill2_act_rgbd_dataset.yaml │ ├── maniskill2_diffusion_policy_rgbd_dataset.yaml │ ├── rlbench_act_rgbd_dataset.yaml │ ├── rlbench_diffusion_policy_rgbd_dataset.yaml │ ├── maniskill2_act_pcd_dataset.yaml │ ├── maniskill2_diffusion_policy_pcd_dataset.yaml │ ├── rlbench_act_pcd_dataset.yaml │ └── rlbench_diffusion_policy_pcd_dataset.yaml ├── exp_maniskill2_act_policy │ ├── maniskill2_model │ │ ├── scratch_vit_pointmap.yaml │ │ ├── scratch_vit_rgbd.yaml │ │ ├── scratch_resnet50_rgbd.yaml │ │ ├── scratch_resnet50_pointmap.yaml │ │ ├── scratch_vit_depth_only.yaml │ │ ├── scratch_resnet50_depth_only.yaml │ │ ├── scratch_pointnet_pcd.yaml │ │ ├── scratch_spunet_pcd.yaml │ │ ├── scratch_pointnet_pcd_presample.yaml │ │ ├── scratch_spunet_pcd_presample.yaml │ │ ├── scratch_pointnet_pcd_wo_rgb.yaml │ │ ├── scratch_pointnet_pcd_wo_xyz.yaml │ │ ├── scratch_spunet_pcd_wo_rgb.yaml │ │ ├── scratch_spunet_pcd_wo_xyz.yaml │ │ ├── scratch_pointnet_pcd_presample_wo_rgb.yaml │ │ ├── scratch_pointnet_pcd_presample_wo_xyz.yaml │ │ ├── scratch_spunet_pcd_presample_wo_rgb.yaml │ │ └── scratch_spunet_pcd_presample_wo_xyz.yaml │ ├── maniskill2_pcd_task │ │ ├── Excavate-v0.yaml │ │ ├── Fill-v0.yaml │ │ ├── Hang-v0.yaml │ │ ├── Pour-v0.yaml │ │ ├── StackCube-v0.yaml │ │ ├── PickCube-v0.yaml │ │ ├── TurnFaucet-v0.yaml │ │ └── PegInsertionSide-v0.yaml │ ├── maniskill2_task │ │ ├── Fill-v0.yaml │ │ ├── Hang-v0.yaml │ │ ├── Pour-v0.yaml │ │ ├── StackCube-v0.yaml │ │ ├── PickCube-v0.yaml │ │ ├── Excavate-v0.yaml │ │ ├── TurnFaucet-v0.yaml │ │ └── PegInsertionSide-v0.yaml │ └── base.yaml ├── exp_maniskill2_diffusion_policy │ ├── maniskill2_model │ │ ├── scratch_vit_rgb.yaml │ │ ├── pretrained_vc1_rgb.yaml │ │ ├── pretrained_r3m_rgb.yaml │ │ ├── scratch_resnet50_rgb.yaml │ │ ├── scratch_vit_pointmap.yaml │ │ ├── scratch_resnet50_pointmap.yaml │ │ ├── scratch_vit_rgbd.yaml │ │ ├── pretrained_multimae_rgbd.yaml │ │ ├── scratch_vit_depth_only.yaml │ │ ├── scratch_multivit_rgbd.yaml │ │ ├── scratch_resnet50_rgbd.yaml │ │ ├── scratch_resnet50_depth_only.yaml │ │ ├── scratch_spunet_pcd.yaml │ │ ├── scratch_pointnet_pcd.yaml │ │ ├── pretrained_ponderv2_pcd.yaml │ │ ├── scratch_spunet_pcd_presample.yaml │ │ ├── scratch_pointnet_pcd_presample.yaml │ │ ├── scratch_spunet_pcd_wo_rgb.yaml │ │ ├── scratch_spunet_pcd_wo_xyz.yaml │ │ ├── scratch_pointnet_pcd_wo_rgb.yaml │ │ ├── scratch_pointnet_pcd_wo_xyz.yaml │ │ ├── scratch_pointnet_pcd_presample_wo_rgb.yaml │ │ ├── scratch_pointnet_pcd_presample_wo_xyz.yaml │ │ ├── scratch_spunet_pcd_presample_wo_rgb.yaml │ │ └── scratch_spunet_pcd_presample_wo_xyz.yaml │ ├── maniskill2_pcd_task │ │ ├── Fill-v0.yaml │ │ ├── Hang-v0.yaml │ │ ├── Pour-v0.yaml │ │ ├── Excavate-v0.yaml │ │ ├── PickCube-v0.yaml │ │ ├── StackCube-v0.yaml │ │ ├── TurnFaucet-v0.yaml │ │ └── PegInsertionSide-v0.yaml │ ├── maniskill2_task │ │ ├── Fill-v0.yaml │ │ ├── Hang-v0.yaml │ │ ├── Pour-v0.yaml │ │ ├── Excavate-v0.yaml │ │ ├── StackCube-v0.yaml │ │ ├── PickCube-v0.yaml │ │ ├── TurnFaucet-v0.yaml │ │ └── PegInsertionSide-v0.yaml │ └── base.yaml ├── hydra │ └── default.yaml ├── exp_rlbench_diffusion_policy │ ├── rlbench_model │ │ ├── pretrained_vc1_rgb.yaml │ │ ├── scratch_vit_rgb.yaml │ │ ├── pretrained_r3m_rgb.yaml │ │ ├── scratch_resnet50_rgb.yaml │ │ ├── scratch_vit_rgbd.yaml │ │ ├── pretrained_multimae_rgbd.yaml │ │ ├── scratch_multivit_rgbd.yaml │ │ ├── scratch_resnet50_rgbd.yaml │ │ ├── pretrained_ponderv2_pcd.yaml │ │ ├── scratch_spunet_pcd.yaml │ │ └── scratch_pointnet_pcd.yaml │ └── base.yaml ├── paths │ └── default.yaml ├── test_rlbench_dp.yaml ├── train.yaml └── test_rlbench_act.yaml ├── src ├── __init__.py ├── data │ ├── __init__.py │ └── components │ │ ├── rlbench │ │ └── __init__.py │ │ ├── __init__.py │ │ ├── maniskill2 │ │ └── __init__.py │ │ └── misc.py ├── models │ ├── __init__.py │ └── components │ │ ├── diffusion_policy │ │ ├── vision │ │ │ └── __init__.py │ │ ├── diffusion │ │ │ ├── __init__.py │ │ │ ├── positional_embedding.py │ │ │ └── conv1d_components.py │ │ ├── __init__.py │ │ └── base_image_policy.py │ │ ├── loss │ │ ├── __init__.py │ │ └── misc.py │ │ ├── pcd_encoder │ │ └── __init__.py │ │ ├── img_encoder │ │ └── __init__.py │ │ ├── act │ │ ├── __init__.py │ │ └── utils.py │ │ └── __init__.py ├── envs │ ├── __init__.py │ └── custom_maniskill2 │ │ ├── mpm │ │ ├── .gitignore │ │ ├── RopeInit.pkl │ │ ├── __init__.py │ │ └── shader │ │ │ ├── common │ │ │ ├── view.glsl │ │ │ └── shadow.glsl │ │ │ └── point │ │ │ ├── deferred.vert │ │ │ ├── composite.vert │ │ │ ├── line.frag │ │ │ ├── shadow.vert │ │ │ ├── gbuffer2.frag │ │ │ ├── gbuffer2.vert │ │ │ ├── shadow_point.vert │ │ │ ├── line.vert │ │ │ ├── shadow_point.frag │ │ │ ├── gbuffer1.vert │ │ │ ├── point.frag │ │ │ └── gbuffer.vert │ │ ├── misc │ │ └── __init__.py │ │ ├── assembly │ │ └── __init__.py │ │ ├── pick_and_place │ │ └── __init__.py │ │ └── __init__.py └── utils │ ├── diffusion_policy │ ├── __init__.py │ ├── module_attr_mixin.py │ ├── shape_util.py │ └── dict_of_tensor_mixin.py │ ├── dist.py │ ├── __init__.py │ ├── spconv_utils.py │ ├── metrics.py │ ├── io.py │ ├── logging_utils.py │ └── instantiators.py ├── libs └── pointops │ ├── src │ ├── __init__.py │ ├── sampling │ │ ├── sampling_cuda_kernel.h │ │ └── sampling_cuda.cpp │ ├── knn_query │ │ ├── knn_query_cuda_kernel.h │ │ └── knn_query_cuda.cpp │ ├── cuda_utils.h │ ├── grouping │ │ ├── grouping_cuda_kernel.h │ │ ├── grouping_cuda.cpp │ │ └── grouping_cuda_kernel.cu │ ├── interpolation │ │ ├── interpolation_cuda_kernel.h │ │ └── interpolation_cuda.cpp │ ├── subtraction │ │ ├── subtraction_cuda_kernel.h │ │ └── subtraction_cuda.cpp │ ├── ball_query │ │ ├── ball_query_cuda.cpp │ │ └── ball_query_cuda_kernel.h │ ├── random_ball_query │ │ ├── random_ball_query_cuda_kernel.h │ │ └── random_ball_query_cuda.cpp │ └── aggregation │ │ ├── aggregation_cuda_kernel.h │ │ └── aggregation_cuda.cpp │ ├── __init__.py │ ├── functions │ ├── __init__.py │ ├── sampling.py │ ├── subtraction.py │ └── aggregation.py │ └── setup.py ├── .project-root ├── assets └── overview.png ├── .github ├── PULL_REQUEST_TEMPLATE.md └── workflows │ └── formatter.yaml ├── requirements.txt ├── setup.py ├── LICENSE └── scripts ├── run_rlbench_camera_view.sh └── run_maniskill2_camera_view.sh /configs/local/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /src/data/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /configs/callbacks/none.yaml: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/envs/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /libs/pointops/src/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/.gitignore: -------------------------------------------------------------------------------- 1 | *.sdf -------------------------------------------------------------------------------- /libs/pointops/__init__.py: -------------------------------------------------------------------------------- 1 | from .functions import * 2 | -------------------------------------------------------------------------------- /src/models/components/diffusion_policy/vision/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/models/components/diffusion_policy/diffusion/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/misc/__init__.py: -------------------------------------------------------------------------------- 1 | from . import turn_faucet 2 | -------------------------------------------------------------------------------- /src/models/components/loss/__init__.py: -------------------------------------------------------------------------------- 1 | from .misc import KLDivergence 2 | -------------------------------------------------------------------------------- /configs/trainer/cpu.yaml: -------------------------------------------------------------------------------- 1 | defaults: 2 | - default 3 | 4 | accelerator: cpu 5 | devices: 1 6 | -------------------------------------------------------------------------------- /configs/trainer/gpu.yaml: -------------------------------------------------------------------------------- 1 | defaults: 2 | - default 3 | 4 | accelerator: gpu 5 | devices: 1 6 | -------------------------------------------------------------------------------- /configs/trainer/mps.yaml: -------------------------------------------------------------------------------- 1 | defaults: 2 | - default 3 | 4 | accelerator: mps 5 | devices: 1 6 | -------------------------------------------------------------------------------- /.project-root: -------------------------------------------------------------------------------- 1 | # this file is required for inferring the project root directory 2 | # do not delete 3 | -------------------------------------------------------------------------------- /assets/overview.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HaoyiZhu/PointCloudMatters/HEAD/assets/overview.png -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/assembly/__init__.py: -------------------------------------------------------------------------------- 1 | # isort: off 2 | from . import peg_insertion_side 3 | -------------------------------------------------------------------------------- /configs/__init__.py: -------------------------------------------------------------------------------- 1 | # this file is needed here to include configs when building project as a package 2 | -------------------------------------------------------------------------------- /src/models/components/pcd_encoder/__init__.py: -------------------------------------------------------------------------------- 1 | from .pointnet import PointNet 2 | from .spunet import SpUNet 3 | -------------------------------------------------------------------------------- /src/models/components/diffusion_policy/__init__.py: -------------------------------------------------------------------------------- 1 | from .diffusion_unet_image_policy import DiffusionUnetImagePolicy 2 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/pick_and_place/__init__.py: -------------------------------------------------------------------------------- 1 | # isort: off 2 | from . import pick_cube 3 | from . import stack_cube 4 | -------------------------------------------------------------------------------- /configs/callbacks/device_stats_monitor.yaml: -------------------------------------------------------------------------------- 1 | device_stat_monitor: 2 | _target_: lightning.pytorch.callbacks.DeviceStatsMonitor 3 | cpu_stats: null -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/__init__.py: -------------------------------------------------------------------------------- 1 | from .assembly import * 2 | from .misc import * 3 | from .mpm import * 4 | from .pick_and_place import * 5 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/RopeInit.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HaoyiZhu/PointCloudMatters/HEAD/src/envs/custom_maniskill2/mpm/RopeInit.pkl -------------------------------------------------------------------------------- /configs/trainer/ddp_sim.yaml: -------------------------------------------------------------------------------- 1 | defaults: 2 | - default 3 | 4 | # simulate DDP on CPU, useful for debugging 5 | accelerator: cpu 6 | devices: 2 7 | strategy: ddp_spawn 8 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/__init__.py: -------------------------------------------------------------------------------- 1 | # isort: off 2 | from . import excavate_env 3 | from . import fill_env 4 | from . import pour_env 5 | from . import hang_env 6 | -------------------------------------------------------------------------------- /configs/debug/fdr.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # runs 1 train, 1 validation and 1 test step 4 | 5 | defaults: 6 | - default 7 | 8 | trainer: 9 | fast_dev_run: true 10 | -------------------------------------------------------------------------------- /src/models/components/img_encoder/__init__.py: -------------------------------------------------------------------------------- 1 | from .multivit import MultiMAEModel, MultiViTModel 2 | from .resnet import R3MResNet, ResNetTorchVision 3 | from .vit import MAEViT, VC1ViT, ViT 4 | -------------------------------------------------------------------------------- /configs/logger/csv.yaml: -------------------------------------------------------------------------------- 1 | # csv logger built in lightning 2 | 3 | csv: 4 | _target_: lightning.pytorch.loggers.csv_logs.CSVLogger 5 | save_dir: "${paths.output_dir}" 6 | name: "csv/" 7 | prefix: "" 8 | -------------------------------------------------------------------------------- /configs/logger/many_loggers.yaml: -------------------------------------------------------------------------------- 1 | # train with many loggers at once 2 | 3 | defaults: 4 | # - comet.yaml 5 | - csv 6 | # - mlflow.yaml 7 | # - neptune.yaml 8 | - tensorboard 9 | - wandb 10 | -------------------------------------------------------------------------------- /configs/callbacks/rich_progress_bar.yaml: -------------------------------------------------------------------------------- 1 | # https://lightning.ai/docs/pytorch/latest/api/lightning.pytorch.callbacks.RichProgressBar.html 2 | 3 | rich_progress_bar: 4 | _target_: lightning.pytorch.callbacks.RichProgressBar 5 | -------------------------------------------------------------------------------- /configs/debug/profiler.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # runs with execution time profiling 4 | 5 | defaults: 6 | - default 7 | 8 | trainer: 9 | max_epochs: 1 10 | profiler: "simple" 11 | # profiler: "advanced" 12 | # profiler: "pytorch" 13 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/common/view.glsl: -------------------------------------------------------------------------------- 1 | float ndc_depth_to_z(float z, float proj22, float proj32) { 2 | return -proj32 / (z + proj22); 3 | } 4 | 5 | float z_to_ndc_depth(float z, float proj22, float proj32) { 6 | return -proj22 - proj32 / z; 7 | } 8 | -------------------------------------------------------------------------------- /src/models/components/act/__init__.py: -------------------------------------------------------------------------------- 1 | from .act import ACT, ACTPCD 2 | from .backbone import Backbone, Joiner 3 | from .positional_encoding import PositionEmbeddingLearned, PositionEmbeddingSine 4 | from .transformer import Transformer, TransformerDecoder, TransformerEncoder 5 | -------------------------------------------------------------------------------- /configs/debug/overfit.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # overfits to 3 batches 4 | 5 | defaults: 6 | - default 7 | 8 | trainer: 9 | max_epochs: 20 10 | overfit_batches: 3 11 | 12 | # model ckpt and early stopping need to be disabled during overfitting 13 | callbacks: 14 | -------------------------------------------------------------------------------- /configs/extras/default.yaml: -------------------------------------------------------------------------------- 1 | # disable python warnings if they annoy you 2 | ignore_warnings: false 3 | 4 | # ask user for tags if none are provided in the config 5 | enforce_tags: true 6 | 7 | # pretty print config tree at the start of the run using Rich library 8 | print_config: true 9 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/deferred.vert: -------------------------------------------------------------------------------- 1 | #version 450 2 | 3 | layout (location = 0) out vec2 outUV; 4 | 5 | void main() 6 | { 7 | outUV = vec2((gl_VertexIndex << 1) & 2, gl_VertexIndex & 2); 8 | gl_Position = vec4(outUV * 2.0f + -1.0f, 0.0f, 1.0f); 9 | } 10 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/composite.vert: -------------------------------------------------------------------------------- 1 | #version 450 2 | 3 | layout (location = 0) out vec2 outUV; 4 | 5 | void main() 6 | { 7 | outUV = vec2((gl_VertexIndex << 1) & 2, gl_VertexIndex & 2); 8 | gl_Position = vec4(outUV * 2.0f + -1.0f, 0.0f, 1.0f); 9 | } 10 | -------------------------------------------------------------------------------- /configs/callbacks/lr_monitor.yaml: -------------------------------------------------------------------------------- 1 | # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.LearningRateMonitor.html#lightning.pytorch.callbacks.LearningRateMonitor 2 | 3 | lr_monitor: 4 | _target_: lightning.pytorch.callbacks.LearningRateMonitor 5 | logging_interval: 'step' 6 | -------------------------------------------------------------------------------- /configs/callbacks/model_summary.yaml: -------------------------------------------------------------------------------- 1 | # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.RichModelSummary.html 2 | 3 | model_summary: 4 | _target_: lightning.pytorch.callbacks.RichModelSummary 5 | max_depth: 1 # the maximum depth of layer nesting that the summary will include 6 | -------------------------------------------------------------------------------- /src/utils/diffusion_policy/__init__.py: -------------------------------------------------------------------------------- 1 | from . import tensor_util 2 | from .dict_of_tensor_mixin import DictOfTensorMixin 3 | from .module_attr_mixin import ModuleAttrMixin 4 | from .normalizer import LinearNormalizer, SingleFieldLinearNormalizer 5 | from .shape_util import get_module_device, get_output_shape 6 | -------------------------------------------------------------------------------- /configs/callbacks/stochastic_weight_averaging.yaml: -------------------------------------------------------------------------------- 1 | # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.StochasticWeightAveraging.html#stochasticweightaveraging 2 | 3 | stochastic_weight_averaging: 4 | _target_: lightning.pytorch.callbacks.StochasticWeightAveraging 5 | swa_lrs: 0.05 6 | -------------------------------------------------------------------------------- /src/data/components/rlbench/__init__.py: -------------------------------------------------------------------------------- 1 | from .rlbench_single_task_act import ( 2 | RLBenchSingleTaskACTPCDDataset, 3 | RLBenchSingleTaskACTRGBDDataset, 4 | ) 5 | from .rlbench_single_task_dp import ( 6 | RLBenchSingleTaskDiffusionPolicyPCDDataset, 7 | RLBenchSingleTaskDiffusionPolicyRGBDDataset, 8 | ) 9 | -------------------------------------------------------------------------------- /configs/debug/limit.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # uses only 1% of the training data and 5% of validation/test data 4 | 5 | defaults: 6 | - default 7 | 8 | trainer: 9 | accelerator: gpu 10 | max_epochs: 3 11 | limit_train_batches: 0.01 12 | limit_val_batches: 0.05 13 | limit_test_batches: 0.05 14 | -------------------------------------------------------------------------------- /configs/logger/tensorboard.yaml: -------------------------------------------------------------------------------- 1 | # https://www.tensorflow.org/tensorboard/ 2 | 3 | tensorboard: 4 | _target_: lightning.pytorch.loggers.tensorboard.TensorBoardLogger 5 | save_dir: "${paths.output_dir}/tensorboard/" 6 | name: 7 | log_graph: false 8 | default_hp_metric: true 9 | prefix: "" 10 | version: ${task_name} 11 | -------------------------------------------------------------------------------- /configs/logger/neptune.yaml: -------------------------------------------------------------------------------- 1 | # https://neptune.ai 2 | 3 | neptune: 4 | _target_: lightning.pytorch.loggers.neptune.NeptuneLogger 5 | api_key: ${oc.env:NEPTUNE_API_TOKEN} # api key is loaded from environment variable 6 | project: username/lightning-hydra-template 7 | # name: "" 8 | log_model_checkpoints: true 9 | prefix: "" 10 | -------------------------------------------------------------------------------- /configs/trainer/ddp.yaml: -------------------------------------------------------------------------------- 1 | defaults: 2 | - default 3 | 4 | strategy: ddp # _find_unused_parameters_true # ddp 5 | 6 | accelerator: gpu 7 | devices: 1 8 | num_nodes: 1 9 | sync_batchnorm: true 10 | max_epochs: 100 11 | check_val_every_n_epoch: 2 12 | gradient_clip_val: 0.5 13 | accumulate_grad_batches: 1 14 | precision: 32-true 15 | log_every_n_steps: 50 -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/rlbench_model/scratch_vit_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_act_rgbd_dataset 4 | - override /model: rlbench_act_model 5 | 6 | task_name: scratch_vit_rgb 7 | 8 | model: 9 | policy: 10 | backbone: 11 | _target_: src.models.components.img_encoder.vit.ViT 12 | model_name: vit_base_patch16 13 | channels: 3 -------------------------------------------------------------------------------- /configs/logger/mlflow.yaml: -------------------------------------------------------------------------------- 1 | # https://mlflow.org 2 | 3 | mlflow: 4 | _target_: lightning.pytorch.loggers.mlflow.MLFlowLogger 5 | # experiment_name: "" 6 | # run_name: "" 7 | tracking_uri: ${paths.log_dir}/mlflow/mlruns # run `mlflow ui` command inside the `logs/mlflow/` dir to open the UI 8 | tags: 9 | # save_dir: "./mlruns" 10 | prefix: "" 11 | artifact_location: 12 | # run_id: "" 13 | -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/rlbench_model/pretrained_vc1_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_act_rgbd_dataset 4 | - override /model: rlbench_act_model 5 | 6 | task_name: pretrained_vc1_rgb 7 | 8 | model: 9 | policy: 10 | backbone: 11 | _target_: src.models.components.img_encoder.vit.VC1ViT 12 | model_name: vit_base_patch16 13 | channels: 3 -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/line.frag: -------------------------------------------------------------------------------- 1 | #version 450 2 | #extension GL_ARB_separate_shader_objects : enable 3 | #extension GL_ARB_shading_language_420pack : enable 4 | 5 | layout(location = 0) in vec4 inPosition; 6 | layout(location = 1) in vec4 inPrevPosition; 7 | layout(location = 2) in vec4 inColor; 8 | 9 | layout(location = 0) out vec4 outLine; 10 | 11 | void main() { 12 | outLine = inColor; 13 | } 14 | -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/rlbench_model/pretrained_r3m_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_act_rgbd_dataset 4 | - override /model: rlbench_act_model 5 | 6 | task_name: pretrained_r3m_rgb 7 | 8 | model: 9 | policy: 10 | backbone: 11 | _target_: src.models.components.img_encoder.resnet.R3MResNet 12 | resnet_model: resnet50 13 | channels: 3 14 | avg_pool: true -------------------------------------------------------------------------------- /configs/logger/comet.yaml: -------------------------------------------------------------------------------- 1 | # https://www.comet.ml 2 | 3 | comet: 4 | _target_: lightning.pytorch.loggers.comet.CometLogger 5 | api_key: ${oc.env:COMET_API_TOKEN} # api key is loaded from environment variable 6 | save_dir: "${paths.output_dir}" 7 | project_name: "lightning-hydra-template" 8 | rest_api_key: 9 | # experiment_name: "" 10 | experiment_key: # set to resume experiment 11 | offline: false 12 | prefix: "" 13 | -------------------------------------------------------------------------------- /src/utils/diffusion_policy/module_attr_mixin.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class ModuleAttrMixin(nn.Module): 5 | def __init__(self): 6 | super().__init__() 7 | self._dummy_variable = nn.Parameter() 8 | 9 | @property 10 | def device(self): 11 | return next(iter(self.parameters())).device 12 | 13 | @property 14 | def dtype(self): 15 | return next(iter(self.parameters())).dtype 16 | -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/rlbench_model/scratch_vit_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_act_rgbd_dataset 4 | - override /model: rlbench_act_model 5 | 6 | task_name: scratch_vit_rgbd 7 | 8 | data: 9 | train: 10 | include_depth: true 11 | 12 | model: 13 | policy: 14 | backbone: 15 | _target_: src.models.components.img_encoder.vit.ViT 16 | model_name: vit_base_patch16 17 | channels: 4 -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/rlbench_model/scratch_resnet50_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_act_rgbd_dataset 4 | - override /model: rlbench_act_model 5 | 6 | task_name: scratch_resnet50_rgb 7 | 8 | model: 9 | policy: 10 | backbone: 11 | _target_: src.models.components.img_encoder.resnet.ResNetTorchVision 12 | resnet_model: resnet50 13 | pretrained: false 14 | channels: 3 15 | avg_pool: true -------------------------------------------------------------------------------- /src/models/components/__init__.py: -------------------------------------------------------------------------------- 1 | from .act import ( 2 | ACT, 3 | ACTPCD, 4 | PositionEmbeddingLearned, 5 | PositionEmbeddingSine, 6 | Transformer, 7 | TransformerDecoder, 8 | TransformerEncoder, 9 | ) 10 | from .diffusion_policy import DiffusionUnetImagePolicy 11 | from .img_encoder import ( 12 | MultiMAEModel, 13 | MultiViTModel, 14 | R3MResNet, 15 | ResNetTorchVision, 16 | VC1ViT, 17 | ViT, 18 | ) 19 | from .pcd_encoder import PointNet, SpUNet 20 | -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/rlbench_model/scratch_pointnet_pcd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_act_pcd_dataset 4 | - override /model: rlbench_act_pcd_model 5 | 6 | task_name: scratch_pointnet_pcd 7 | 8 | data: 9 | batch_size_train: 8 10 | 11 | trainer: 12 | accumulate_grad_batches: 4 13 | 14 | model: 15 | policy: 16 | backbone: 17 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 18 | in_channels: 6 19 | num_classes: 0 -------------------------------------------------------------------------------- /configs/data/maniskill2_act_rgbd_dataset.yaml: -------------------------------------------------------------------------------- 1 | _target_: src.data.base_datamodule.BaseDataModule 2 | 3 | batch_size_train: 16 4 | batch_size_val: 1 5 | batch_size_test: 1 6 | num_workers: 16 7 | pin_memory: true 8 | train: 9 | dataset_file: 10 | load_count: -1 11 | camera_names: 12 | - base_camera 13 | include_depth: false 14 | scale_rgb_only: false 15 | goal_cond_keys: 16 | chunk_size: 100 17 | 18 | val: 19 | _target_: src.data.components.misc.DummyDataset 20 | size: 400 # 400 eval trials -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/rlbench_model/pretrained_multimae_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_act_rgbd_dataset 4 | - override /model: rlbench_act_model 5 | 6 | task_name: pretrained_multimae_rgbd 7 | 8 | data: 9 | batch_size_train: 16 10 | train: 11 | include_depth: true 12 | 13 | trainer: 14 | accumulate_grad_batches: 2 15 | 16 | model: 17 | policy: 18 | backbone: 19 | _target_: src.models.components.img_encoder.multivit.MultiViTModel 20 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_vit_pointmap.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_act_model 4 | - override /paths: default 5 | 6 | task_name: scratch_vit_pointmap 7 | 8 | data: 9 | batch_size_train: 16 10 | num_workers: 16 11 | train: 12 | pointmap: true 13 | 14 | model: 15 | policy: 16 | backbone: 17 | _target_: src.models.components.img_encoder.vit.ViT 18 | model_name: vit_base_patch16 19 | pretrained: false 20 | channels: 6 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_vit_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_act_model 4 | - override /paths: default 5 | 6 | task_name: scratch_vit_rgbd 7 | 8 | data: 9 | batch_size_train: 16 10 | num_workers: 16 11 | train: 12 | include_depth: true 13 | 14 | model: 15 | policy: 16 | backbone: 17 | _target_: src.models.components.img_encoder.vit.ViT 18 | model_name: vit_base_patch16 19 | pretrained: false 20 | channels: 4 -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/rlbench_model/scratch_resnet50_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_act_rgbd_dataset 4 | - override /model: rlbench_act_model 5 | 6 | task_name: scratch_resnet50_rgbd 7 | 8 | data: 9 | train: 10 | include_depth: true 11 | 12 | model: 13 | policy: 14 | backbone: 15 | _target_: src.models.components.img_encoder.resnet.ResNetTorchVision 16 | resnet_model: resnet50 17 | pretrained: false 18 | channels: 4 19 | avg_pool: true -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/rlbench_model/scratch_spunet_pcd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_act_pcd_dataset 4 | - override /model: rlbench_act_pcd_model 5 | 6 | task_name: scratch_spunet_pcd 7 | 8 | data: 9 | batch_size_train: 8 10 | 11 | trainer: 12 | accumulate_grad_batches: 4 13 | 14 | model: 15 | policy: 16 | backbone: 17 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 18 | in_channels: 6 19 | num_classes: 0 20 | pretrained_path: null -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/rlbench_model/scratch_multivit_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_act_rgbd_dataset 4 | - override /model: rlbench_act_model 5 | 6 | task_name: scratch_multivit_rgbd 7 | 8 | data: 9 | batch_size_train: 16 10 | train: 11 | include_depth: true 12 | 13 | trainer: 14 | accumulate_grad_batches: 2 15 | 16 | 17 | model: 18 | policy: 19 | backbone: 20 | _target_: src.models.components.img_encoder.multivit.MultiViTModel 21 | ckpt_path: null -------------------------------------------------------------------------------- /configs/data/maniskill2_diffusion_policy_rgbd_dataset.yaml: -------------------------------------------------------------------------------- 1 | _target_: src.data.base_datamodule.BaseDataModule 2 | 3 | batch_size_train: 32 4 | batch_size_val: 1 5 | batch_size_test: 1 6 | num_workers: 16 7 | pin_memory: true 8 | train: 9 | dataset_file: 10 | load_count: -1 11 | camera_names: 12 | - base_camera 13 | include_depth: false 14 | scale_rgb_only: false 15 | goal_cond_keys: 16 | chunk_size: 16 # equal to horizon 17 | 18 | val: 19 | _target_: src.data.components.misc.DummyDataset 20 | size: 400 # 400 eval trials -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_resnet50_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_act_model 4 | - override /paths: default 5 | 6 | task_name: scratch_resnet50_rgbd 7 | 8 | data: 9 | batch_size_train: 16 10 | num_workers: 16 11 | train: 12 | include_depth: true 13 | 14 | model: 15 | policy: 16 | backbone: 17 | _target_: src.models.components.img_encoder.resnet.ResNetTorchVision 18 | resnet_model: resnet50 19 | pretrained: false 20 | channels: 4 21 | -------------------------------------------------------------------------------- /configs/trainer/default.yaml: -------------------------------------------------------------------------------- 1 | _target_: lightning.pytorch.trainer.Trainer 2 | 3 | default_root_dir: ${paths.output_dir} 4 | 5 | min_epochs: 1 # prevents early stopping 6 | max_epochs: 10 7 | 8 | accelerator: ddp 9 | devices: 1 10 | 11 | # mixed precision for extra speed-up 12 | # precision: 16 13 | 14 | # perform a validation loop every N training epochs 15 | check_val_every_n_epoch: 1 16 | 17 | # set True to to ensure deterministic results 18 | # makes training slower but gives more reproducibility than just setting seeds 19 | deterministic: false 20 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_resnet50_pointmap.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_act_model 4 | - override /paths: default 5 | 6 | task_name: scratch_resnet50_pointmap 7 | 8 | data: 9 | batch_size_train: 16 10 | num_workers: 16 11 | train: 12 | pointmap: true 13 | 14 | model: 15 | policy: 16 | backbone: 17 | _target_: src.models.components.img_encoder.resnet.ResNetTorchVision 18 | resnet_model: resnet50 19 | pretrained: false 20 | channels: 6 21 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_vit_depth_only.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_act_model 4 | - override /paths: default 5 | 6 | task_name: scratch_vit_depth_only 7 | 8 | data: 9 | batch_size_train: 16 10 | num_workers: 16 11 | train: 12 | include_depth: true 13 | only_depth: true 14 | 15 | model: 16 | policy: 17 | backbone: 18 | _target_: src.models.components.img_encoder.vit.ViT 19 | model_name: vit_base_patch16 20 | pretrained: false 21 | channels: 1 -------------------------------------------------------------------------------- /libs/pointops/functions/__init__.py: -------------------------------------------------------------------------------- 1 | from .aggregation import aggregation 2 | from .attention import attention_fusion_step, attention_relation_step 3 | from .grouping import grouping, grouping2 4 | from .interpolation import interpolation, interpolation2 5 | from .query import ball_query, knn_query, random_ball_query 6 | from .sampling import farthest_point_sampling 7 | from .subtraction import subtraction 8 | from .utils import ( 9 | ball_query_and_group, 10 | batch2offset, 11 | knn_query_and_group, 12 | offset2batch, 13 | query_and_group, 14 | ) 15 | -------------------------------------------------------------------------------- /configs/callbacks/default.yaml: -------------------------------------------------------------------------------- 1 | defaults: 2 | - model_checkpoint 3 | - early_stopping 4 | - model_summary 5 | - rich_progress_bar 6 | - lr_monitor 7 | - device_stats_monitor 8 | - _self_ 9 | 10 | model_checkpoint: 11 | dirpath: ${paths.output_dir}/checkpoints 12 | filename: "epoch_{epoch:03d}" 13 | monitor: "val/loss" 14 | mode: "min" 15 | save_last: true 16 | auto_insert_metric_name: false 17 | save_top_k: 3 18 | 19 | early_stopping: 20 | monitor: "val/loss" 21 | patience: 100 22 | mode: "min" 23 | 24 | model_summary: 25 | max_depth: -1 26 | -------------------------------------------------------------------------------- /src/data/components/__init__.py: -------------------------------------------------------------------------------- 1 | from .maniskill2 import ( 2 | ManiSkill2GoalPosSingleTaskACTPCDDataset, 3 | ManiSkill2GoalPosSingleTaskACTRGBDDataset, 4 | ManiSkill2GoalPosSingleTaskDiffusionPolicyPCDDataset, 5 | ManiSkill2GoalPosSingleTaskDiffusionPolicyRGBDDataset, 6 | ManiSkill2NullGoalSingleTaskACTPCDDataset, 7 | ManiSkill2NullGoalSingleTaskACTRGBDDataset, 8 | ManiSkill2NullGoalSingleTaskDiffusionPolicyPCDDataset, 9 | ManiSkill2NullGoalSingleTaskDiffusionPolicyRGBDDataset, 10 | ) 11 | from .misc import DummyDataset, ExperienceSourceDataset 12 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_resnet50_depth_only.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_act_model 4 | - override /paths: default 5 | 6 | task_name: scratch_resnet50_depth_only 7 | 8 | data: 9 | batch_size_train: 16 10 | num_workers: 16 11 | train: 12 | include_depth: true 13 | only_depth: true 14 | 15 | model: 16 | policy: 17 | backbone: 18 | _target_: src.models.components.img_encoder.resnet.ResNetTorchVision 19 | resnet_model: resnet50 20 | pretrained: false 21 | channels: 1 22 | -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/rlbench_model/pretrained_ponderv2_pcd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_act_pcd_dataset 4 | - override /model: rlbench_act_pcd_model 5 | 6 | task_name: pretrained_ponderv2_pcd 7 | 8 | data: 9 | batch_size_train: 16 10 | 11 | trainer: 12 | accumulate_grad_batches: 2 13 | 14 | model: 15 | policy: 16 | backbone: 17 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 18 | in_channels: 6 19 | num_classes: 0 20 | pretrained_path: checkpoints/ponderv2-ppt-pretrain-scannet-s3dis-structured3d.pth -------------------------------------------------------------------------------- /src/utils/dist.py: -------------------------------------------------------------------------------- 1 | import torch.distributed as dist 2 | 3 | 4 | def is_dist_avail_and_initialized(): 5 | if not dist.is_available(): 6 | return False 7 | if not dist.is_initialized(): 8 | return False 9 | return True 10 | 11 | 12 | def get_world_size(): 13 | if not is_dist_avail_and_initialized(): 14 | return 1 15 | return dist.get_world_size() 16 | 17 | 18 | def get_rank(): 19 | if not is_dist_avail_and_initialized(): 20 | return 0 21 | return dist.get_rank() 22 | 23 | 24 | def is_main_process(): 25 | return get_rank() == 0 26 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_pointnet_pcd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_pointnet_pcd 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | 13 | model: 14 | obs_mode: pointcloud 15 | policy: 16 | backbone: 17 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 18 | in_channels: 6 19 | num_classes: 0 20 | 21 | trainer: 22 | accumulate_grad_batches: 2 23 | -------------------------------------------------------------------------------- /configs/logger/wandb.yaml: -------------------------------------------------------------------------------- 1 | # https://wandb.ai 2 | 3 | wandb: 4 | _target_: lightning.pytorch.loggers.wandb.WandbLogger 5 | # name: "" # name of the run (normally generated by wandb) 6 | save_dir: "${paths.output_dir}" 7 | offline: false 8 | id: # pass correct id to resume experiment! 9 | anonymous: # enable anonymous logging 10 | project: "lightning-hydra-template" 11 | log_model: false # upload lightning ckpts 12 | prefix: "" # a string to put at the beginning of metric keys 13 | # entity: "" # set to name of your wandb team 14 | group: "" 15 | tags: [] 16 | job_type: "" 17 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_spunet_pcd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_spunet_pcd 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | 13 | model: 14 | obs_mode: pointcloud 15 | policy: 16 | backbone: 17 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 18 | in_channels: 6 19 | num_classes: 0 20 | pretrained_path: null 21 | 22 | trainer: 23 | accumulate_grad_batches: 2 -------------------------------------------------------------------------------- /.github/PULL_REQUEST_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | ## What does this PR do? 2 | 3 | Fixes #\ 4 | 5 | ## Before submitting 6 | 7 | - [ ] Did you make sure **title is self-explanatory** and **the description concisely explains the PR**? 8 | - [ ] Did you make sure your **PR does only one thing**, instead of bundling different changes together? 9 | - [ ] Did you list all the **breaking changes** introduced by this pull request? 10 | - [ ] Did you **test your PR locally**? 11 | - [ ] Did you **format your code** with `isort . --profile=black;black .` command? 12 | 13 | ## Did you have fun? 14 | 15 | Make sure you had fun coding 🙃 16 | -------------------------------------------------------------------------------- /src/models/components/diffusion_policy/diffusion/positional_embedding.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import torch 4 | import torch.nn as nn 5 | 6 | 7 | class SinusoidalPosEmb(nn.Module): 8 | def __init__(self, dim): 9 | super().__init__() 10 | self.dim = dim 11 | 12 | def forward(self, x): 13 | device = x.device 14 | half_dim = self.dim // 2 15 | emb = math.log(10000) / (half_dim - 1) 16 | emb = torch.exp(torch.arange(half_dim, device=device) * -emb) 17 | emb = x[:, None] * emb[None, :] 18 | emb = torch.cat((emb.sin(), emb.cos()), dim=-1) 19 | return emb 20 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_vit_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_vit_rgb 7 | 8 | model: 9 | policy: 10 | shape_meta: 11 | obs: 12 | base_camera_rgb: 13 | shape: [3, 128, 128] 14 | type: rgb 15 | qpos: 16 | shape: [7] 17 | type: low_dim 18 | obs_encoder: 19 | rgb_model: 20 | _target_: src.models.components.img_encoder.vit.ViT 21 | model_name: vit_base_patch16 22 | channels: 3 23 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_pointnet_pcd_presample.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_pointnet_pcd_presample 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | 13 | model: 14 | obs_mode: pointcloud 15 | policy: 16 | backbone: 17 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 18 | in_channels: 6 19 | num_classes: 0 20 | pre_sample: true 21 | 22 | trainer: 23 | accumulate_grad_batches: 2 -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/pretrained_vc1_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: pretrained_vc1_rgb 7 | 8 | model: 9 | policy: 10 | shape_meta: 11 | obs: 12 | base_camera_rgb: 13 | shape: [3, 128, 128] 14 | type: rgb 15 | qpos: 16 | shape: [7] 17 | type: low_dim 18 | obs_encoder: 19 | rgb_model: 20 | _target_: src.models.components.img_encoder.vit.VC1ViT 21 | model_name: vit_base_patch16 22 | channels: 3 23 | -------------------------------------------------------------------------------- /libs/pointops/src/sampling/sampling_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _SAMPLING_CUDA_KERNEL 2 | #define _SAMPLING_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void farthest_point_sampling_cuda(int b, int n, at::Tensor xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor tmp_tensor, at::Tensor idx_tensor); 8 | 9 | #ifdef __cplusplus 10 | extern "C" { 11 | #endif 12 | 13 | void farthest_point_sampling_cuda_launcher(int b, int n, const float *xyz, const int *offset, const int *new_offset, float *tmp, int *idx); 14 | 15 | #ifdef __cplusplus 16 | } 17 | #endif 18 | #endif 19 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_spunet_pcd_presample.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_spunet_pcd_presample 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | 13 | model: 14 | obs_mode: pointcloud 15 | policy: 16 | backbone: 17 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 18 | in_channels: 6 19 | num_classes: 512 20 | pretrained_path: null 21 | pre_sample: true 22 | 23 | trainer: 24 | accumulate_grad_batches: 2 -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/pretrained_r3m_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: pretrained_r3m_rgb 7 | 8 | model: 9 | policy: 10 | shape_meta: 11 | obs: 12 | base_camera_rgb: 13 | shape: [3, 128, 128] 14 | type: rgb 15 | qpos: 16 | shape: [7] 17 | type: low_dim 18 | obs_encoder: 19 | rgb_model: 20 | _target_: src.models.components.img_encoder.resnet.R3MResNet 21 | resnet_model: resnet50 22 | channels: 3 23 | avg_pool: true 24 | -------------------------------------------------------------------------------- /configs/hydra/default.yaml: -------------------------------------------------------------------------------- 1 | # https://hydra.cc/docs/configure_hydra/intro/ 2 | 3 | # enable color logging 4 | defaults: 5 | - override hydra_logging: colorlog 6 | - override job_logging: colorlog 7 | 8 | # output directory, generated dynamically on each run 9 | run: 10 | dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S} 11 | sweep: 12 | dir: ${paths.log_dir}/${task_name}/multiruns/${now:%Y-%m-%d}_${now:%H-%M-%S} 13 | subdir: ${hydra.job.num} 14 | 15 | job_logging: 16 | handlers: 17 | file: 18 | # Incorporates fix from https://github.com/facebookresearch/hydra/pull/2242 19 | filename: ${hydra.runtime.output_dir}/${task_name}.log 20 | -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/rlbench_model/pretrained_vc1_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_diffusion_policy_rgbd_dataset 4 | - override /model: rlbench_diffusion_policy_model 5 | 6 | task_name: pretrained_vc1_rgb 7 | 8 | model: 9 | policy: 10 | shape_meta: 11 | obs: 12 | front_rgb: 13 | shape: [3, 128, 128] 14 | type: rgb 15 | qpos: 16 | shape: 17 | - ${data.train.action_dim} 18 | type: low_dim 19 | obs_encoder: 20 | rgb_model: 21 | _target_: src.models.components.img_encoder.vit.VC1ViT 22 | model_name: vit_base_patch16 23 | channels: 3 -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/rlbench_model/scratch_vit_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_diffusion_policy_rgbd_dataset 4 | - override /model: rlbench_diffusion_policy_model 5 | 6 | task_name: scratch_vit_rgb 7 | 8 | model: 9 | policy: 10 | shape_meta: 11 | obs: 12 | front_rgb: 13 | shape: [3, 128, 128] 14 | type: rgb 15 | qpos: 16 | shape: 17 | - ${data.train.action_dim} 18 | type: low_dim 19 | obs_encoder: 20 | rgb_model: 21 | _target_: src.models.components.img_encoder.vit.ViT 22 | model_name: vit_base_patch16 23 | channels: 3 24 | -------------------------------------------------------------------------------- /libs/pointops/src/knn_query/knn_query_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _KNN_QUERY_CUDA_KERNEL 2 | #define _KNN_QUERY_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void knn_query_cuda(int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor); 8 | 9 | #ifdef __cplusplus 10 | extern "C" { 11 | #endif 12 | 13 | void knn_query_cuda_launcher(int m, int nsample, const float *xyz, const float *new_xyz, const int *offset, const int *new_offset, int *idx, float *dist2); 14 | 15 | #ifdef __cplusplus 16 | } 17 | #endif 18 | #endif 19 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_resnet50_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_resnet50_rgb 7 | 8 | model: 9 | policy: 10 | shape_meta: 11 | obs: 12 | base_camera_rgb: 13 | shape: [3, 128, 128] 14 | type: rgb 15 | qpos: 16 | shape: [7] 17 | type: low_dim 18 | obs_encoder: 19 | rgb_model: 20 | _target_: src.models.components.img_encoder.resnet.ResNetTorchVision 21 | resnet_model: resnet50 22 | pretrained: false 23 | channels: 3 24 | avg_pool: true -------------------------------------------------------------------------------- /src/data/components/maniskill2/__init__.py: -------------------------------------------------------------------------------- 1 | from .maniskill2_single_task_pcd_act import ( 2 | ManiSkill2GoalPosSingleTaskACTPCDDataset, 3 | ManiSkill2NullGoalSingleTaskACTPCDDataset, 4 | ) 5 | from .maniskill2_single_task_pcd_dp import ( 6 | ManiSkill2GoalPosSingleTaskDiffusionPolicyPCDDataset, 7 | ManiSkill2NullGoalSingleTaskDiffusionPolicyPCDDataset, 8 | ) 9 | from .maniskill2_single_task_rgbd_act import ( 10 | ManiSkill2GoalPosSingleTaskACTRGBDDataset, 11 | ManiSkill2NullGoalSingleTaskACTRGBDDataset, 12 | ) 13 | from .maniskill2_single_task_rgbd_dp import ( 14 | ManiSkill2GoalPosSingleTaskDiffusionPolicyRGBDDataset, 15 | ManiSkill2NullGoalSingleTaskDiffusionPolicyRGBDDataset, 16 | ) 17 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_vit_pointmap.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_vit_pointmap 7 | 8 | data: 9 | train: 10 | pointmap: true 11 | 12 | model: 13 | obs_mode: pointcloud 14 | policy: 15 | shape_meta: 16 | obs: 17 | base_camera_rgb: 18 | shape: [6, 128, 128] 19 | type: rgb 20 | qpos: 21 | shape: [7] 22 | type: low_dim 23 | obs_encoder: 24 | rgb_model: 25 | _target_: src.models.components.img_encoder.vit.ViT 26 | model_name: vit_base_patch16 27 | channels: 6 -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/rlbench_model/pretrained_r3m_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_diffusion_policy_rgbd_dataset 4 | - override /model: rlbench_diffusion_policy_model 5 | 6 | task_name: pretrained_r3m_rgb 7 | 8 | model: 9 | policy: 10 | shape_meta: 11 | obs: 12 | front_rgb: 13 | shape: [3, 128, 128] 14 | type: rgb 15 | qpos: 16 | shape: 17 | - ${data.train.action_dim} 18 | type: low_dim 19 | obs_encoder: 20 | rgb_model: 21 | _target_: src.models.components.img_encoder.resnet.R3MResNet 22 | resnet_model: resnet50 23 | channels: 3 24 | avg_pool: true -------------------------------------------------------------------------------- /.github/workflows/formatter.yaml: -------------------------------------------------------------------------------- 1 | name: Formatter 2 | 3 | on: 4 | workflow_dispatch: 5 | push: 6 | branches: 7 | - main 8 | pull_request: 9 | types: [opened, reopened, synchronize] 10 | 11 | concurrency: 12 | group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} 13 | cancel-in-progress: true 14 | 15 | jobs: 16 | formatter: 17 | runs-on: ubuntu-latest 18 | steps: 19 | - uses: actions/checkout@v3 20 | - uses: actions/setup-python@v2 21 | with: 22 | python-version: '3.x' 23 | - run: python -m pip install --upgrade pip 24 | - run: pip install isort black 25 | - run: isort . --profile black 26 | - uses: psf/black@stable -------------------------------------------------------------------------------- /configs/paths/default.yaml: -------------------------------------------------------------------------------- 1 | # path to root directory 2 | # this requires PROJECT_ROOT environment variable to exist 3 | # you can replace it with "." if you want the root to be the current working directory 4 | root_dir: ${oc.env:PROJECT_ROOT} 5 | 6 | # path to data directory 7 | data_dir: ${paths.root_dir}/data/ 8 | 9 | # path to logging directory 10 | log_dir: ${paths.root_dir}/logs/ 11 | 12 | # path to output directory, created dynamically by hydra 13 | # path generation pattern is specified in `configs/hydra/default.yaml` 14 | # use it to store all files generated during the run, like ckpts and metrics 15 | output_dir: ${hydra:runtime.output_dir} 16 | 17 | # path to working directory 18 | work_dir: ${hydra:runtime.cwd} 19 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_pcd_task/Excavate-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_act/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_pcd_dataset 8 | 9 | data: 10 | pin_memory: true 11 | train: 12 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_act.ManiSkill2GoalPosSingleTaskACTPCDDataset 13 | dataset_file: data/maniskill2/demos/v0/soft_body/Excavate-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 14 | goal_cond_keys: 15 | - target 16 | 17 | model: 18 | env_id: Excavate-v0 19 | obs_mode: pointcloud 20 | policy: 21 | action_dim: 6 22 | qpos_dim: 7 23 | goal_cond_dim: 1 24 | -------------------------------------------------------------------------------- /libs/pointops/src/sampling/sampling_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "sampling_cuda_kernel.h" 5 | 6 | 7 | void farthest_point_sampling_cuda(int b, int n, at::Tensor xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor tmp_tensor, at::Tensor idx_tensor) 8 | { 9 | const float *xyz = xyz_tensor.data_ptr(); 10 | const int *offset = offset_tensor.data_ptr(); 11 | const int *new_offset = new_offset_tensor.data_ptr(); 12 | float *tmp = tmp_tensor.data_ptr(); 13 | int *idx = idx_tensor.data_ptr(); 14 | farthest_point_sampling_cuda_launcher(b, n, xyz, offset, new_offset, tmp, idx); 15 | } 16 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_pcd_task/Fill-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_pcd_dataset 8 | 9 | data: 10 | pin_memory: true 11 | train: 12 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_act.ManiSkill2GoalPosSingleTaskACTPCDDataset 13 | dataset_file: data/maniskill2/demos/v0/soft_body/Fill-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 14 | goal_cond_keys: 15 | - target 16 | 17 | model: 18 | env_id: Fill-v0 19 | obs_mode: pointcloud 20 | policy: 21 | action_dim: 6 22 | qpos_dim: 7 23 | goal_cond_dim: 2 24 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_pcd_task/Hang-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_pcd_dataset 8 | 9 | data: 10 | pin_memory: true 11 | train: 12 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_act.ManiSkill2GoalPosSingleTaskACTPCDDataset 13 | dataset_file: data/maniskill2/demos/v0/soft_body/Hang-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 14 | goal_cond_keys: 15 | - target 16 | 17 | model: 18 | env_id: Hang-v0 19 | obs_mode: pointcloud 20 | policy: 21 | action_dim: 7 22 | qpos_dim: 9 23 | goal_cond_dim: 7 24 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_pcd_task/Pour-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_pcd_dataset 8 | 9 | data: 10 | pin_memory: true 11 | train: 12 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_act.ManiSkill2GoalPosSingleTaskACTPCDDataset 13 | dataset_file: data/maniskill2/demos/v0/soft_body/Pour-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 14 | goal_cond_keys: 15 | - target 16 | 17 | model: 18 | env_id: Pour-v0 19 | obs_mode: pointcloud 20 | policy: 21 | action_dim: 7 22 | qpos_dim: 9 23 | goal_cond_dim: 1 24 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_pcd_task/StackCube-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_pcd_dataset 8 | 9 | data: 10 | pin_memory: true 11 | train: 12 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_act.ManiSkill2NullGoalSingleTaskACTPCDDataset 13 | dataset_file: data/maniskill2/demos/v0/rigid_body/StackCube-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 14 | goal_cond_keys: null 15 | 16 | model: 17 | env_id: StackCube-v0 18 | obs_mode: pointcloud 19 | policy: 20 | action_dim: 7 21 | qpos_dim: 9 22 | goal_cond_dim: 0 23 | 24 | -------------------------------------------------------------------------------- /libs/pointops/src/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | 4 | #include 5 | #include 6 | 7 | #define TOTAL_THREADS 1024 8 | #define THREADS_PER_BLOCK 512 9 | #define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0)) 10 | 11 | inline int opt_n_threads(int work_size) { 12 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 13 | return std::max(std::min(1 << pow_2, TOTAL_THREADS), 1); 14 | } 15 | 16 | inline dim3 opt_block_config(int x, int y) { 17 | const int x_threads = opt_n_threads(x); 18 | const int y_threads = std::max(std::min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 19 | dim3 block_config(x_threads, y_threads, 1); 20 | return block_config; 21 | } 22 | 23 | #endif 24 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | h5py==3.11.0 2 | pyyaml==6.0.1 3 | sharedarray==3.2.3 4 | tensorboard==2.16.2 5 | tensorboardx==2.6.2.2 6 | yapf==0.40.2 7 | addict==2.4.0 8 | einops==0.8.0 9 | scipy==1.13.1 10 | open3d==0.18.0 11 | plyfile==1.0.3 12 | termcolor==2.4.0 13 | timm==0.9.10 14 | hydra-colorlog==1.2.0 15 | hydra-core==1.3.2 16 | hydra-optuna-sweeper==1.2.0 17 | lightning>=2.0.0 18 | opencv-python==4.10.0.82 19 | pre-commit==3.7.1 20 | pytest==8.2.2 21 | rich==13.7.1 22 | transforms3d==0.4.1 23 | isort==5.13.2 24 | black==24.4.2 25 | torchmetrics>=0.11.4 26 | torch-geometric==2.5.3 27 | rootutils==1.0.7 28 | imageio==2.34.1 29 | python-dotenv==1.0.1 30 | zarr==2.18.2 31 | diffusers==0.29.0 32 | git+https://github.com/openai/CLIP.git 33 | gym==0.25.2 34 | gymnasium==0.29.1 35 | -------------------------------------------------------------------------------- /src/utils/diffusion_policy/shape_util.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, Dict, List, Tuple 2 | 3 | import torch 4 | import torch.nn as nn 5 | 6 | 7 | def get_module_device(m: nn.Module): 8 | device = torch.device("cpu") 9 | try: 10 | param = next(iter(m.parameters())) 11 | device = param.device 12 | except StopIteration: 13 | pass 14 | return device 15 | 16 | 17 | @torch.no_grad() 18 | def get_output_shape( 19 | input_shape: Tuple[int], net: Callable[[torch.Tensor], torch.Tensor] 20 | ): 21 | device = get_module_device(net) 22 | test_input = torch.zeros((1,) + tuple(input_shape), device=device) 23 | test_output = net(test_input) 24 | output_shape = tuple(test_output.shape[1:]) 25 | return output_shape 26 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from setuptools import find_packages, setup 4 | 5 | setup( 6 | name="src", 7 | version="0.0.1", 8 | description="Point Cloud Matters: Rethinking the Impact of Different Observation Spaces on Robot Learning", 9 | author="Haoyi Zhu", 10 | author_email="hyizhu1108@gmail.com", 11 | url="https://github.com/HaoyiZhu/PointCloudMatters", 12 | install_requires=["lightning", "hydra-core"], 13 | packages=find_packages(), 14 | # use this to customize global commands available in the terminal after installing the package 15 | entry_points={ 16 | "console_scripts": [ 17 | "train_command = src.train:main", 18 | "eval_command = src.eval:main", 19 | ] 20 | }, 21 | ) 22 | -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/rlbench_model/scratch_resnet50_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_diffusion_policy_rgbd_dataset 4 | - override /model: rlbench_diffusion_policy_model 5 | 6 | task_name: scratch_resnet50_rgb 7 | 8 | model: 9 | policy: 10 | shape_meta: 11 | obs: 12 | front_rgb: 13 | shape: [3, 128, 128] 14 | type: rgb 15 | qpos: 16 | shape: 17 | - ${data.train.action_dim} 18 | type: low_dim 19 | obs_encoder: 20 | rgb_model: 21 | rgb_model: 22 | _target_: src.models.components.img_encoder.resnet.ResNetTorchVision 23 | resnet_model: resnet50 24 | pretrained: false 25 | channels: 3 26 | avg_pool: true -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_pcd_task/PickCube-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_pcd_dataset 8 | - override /trainer: ddp 9 | 10 | data: 11 | pin_memory: true 12 | train: 13 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_act.ManiSkill2GoalPosSingleTaskACTPCDDataset 14 | dataset_file: data/maniskill2/demos/v0/rigid_body/PickCube-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 15 | goal_cond_keys: 16 | - goal_pos 17 | 18 | model: 19 | env_id: PickCube-v0 20 | obs_mode: pointcloud 21 | policy: 22 | action_dim: 7 23 | qpos_dim: 9 24 | goal_cond_dim: 3 25 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_resnet50_pointmap.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_resnet50_pointmap 7 | 8 | data: 9 | train: 10 | pointmap: true 11 | 12 | model: 13 | obs_mode: pointcloud 14 | policy: 15 | shape_meta: 16 | obs: 17 | base_camera_rgb: 18 | shape: [6, 128, 128] 19 | type: rgb 20 | qpos: 21 | shape: [7] 22 | type: low_dim 23 | obs_encoder: 24 | rgb_model: 25 | _target_: src.models.components.img_encoder.resnet.ResNetTorchVision 26 | resnet_model: resnet50 27 | pretrained: false 28 | channels: 6 29 | avg_pool: true -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_task/Fill-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_rgbd_dataset 8 | - override /model: maniskill2_act_model 9 | 10 | data: 11 | pin_memory: true 12 | train: 13 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_act.ManiSkill2GoalPosSingleTaskACTRGBDDataset 14 | dataset_file: data/maniskill2/demos/v0/soft_body/Fill-v0/trajectory.rgbd.pd_ee_delta_pose.h5 15 | camera_names: 16 | - base_camera 17 | goal_cond_keys: 18 | - target 19 | chunk_size: 100 20 | 21 | model: 22 | env_id: Fill-v0 23 | policy: 24 | action_dim: 6 25 | qpos_dim: 7 26 | goal_cond_dim: 2 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_task/Hang-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_rgbd_dataset 8 | - override /model: maniskill2_act_model 9 | 10 | data: 11 | pin_memory: true 12 | train: 13 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_act.ManiSkill2GoalPosSingleTaskACTRGBDDataset 14 | dataset_file: data/maniskill2/demos/v0/soft_body/Hang-v0/trajectory.rgbd.pd_ee_delta_pose.h5 15 | camera_names: 16 | - base_camera 17 | goal_cond_keys: 18 | - target 19 | chunk_size: 100 20 | 21 | model: 22 | env_id: Hang-v0 23 | policy: 24 | action_dim: 7 25 | qpos_dim: 9 26 | goal_cond_dim: 7 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_task/Pour-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_rgbd_dataset 8 | - override /model: maniskill2_act_model 9 | 10 | data: 11 | pin_memory: true 12 | train: 13 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_act.ManiSkill2GoalPosSingleTaskACTRGBDDataset 14 | dataset_file: data/maniskill2/demos/v0/soft_body/Pour-v0/trajectory.rgbd.pd_ee_delta_pose.h5 15 | camera_names: 16 | - base_camera 17 | goal_cond_keys: 18 | - target 19 | chunk_size: 100 20 | 21 | model: 22 | env_id: Pour-v0 23 | policy: 24 | action_dim: 7 25 | qpos_dim: 9 26 | goal_cond_dim: 1 -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_vit_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_vit_rgbd 7 | 8 | data: 9 | train: 10 | include_depth: true 11 | 12 | model: 13 | policy: 14 | shape_meta: 15 | obs: 16 | base_camera_rgb: 17 | shape: [3, 128, 128] 18 | type: rgb 19 | base_camera_depth: 20 | shape: [1, 128, 128] 21 | type: depth 22 | qpos: 23 | shape: [7] 24 | type: low_dim 25 | obs_encoder: 26 | use_depth: true 27 | rgb_model: 28 | _target_: src.models.components.img_encoder.vit.ViT 29 | model_name: vit_base_patch16 30 | channels: 4 31 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_task/StackCube-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_rgbd_dataset 8 | - override /model: maniskill2_act_model 9 | 10 | data: 11 | pin_memory: true 12 | train: 13 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_act.ManiSkill2NullGoalSingleTaskACTRGBDDataset 14 | dataset_file: data/maniskill2/demos/v0/rigid_body/StackCube-v0/trajectory.rgbd.pd_ee_delta_pose.h5 15 | camera_names: 16 | - base_camera 17 | goal_cond_keys: null 18 | chunk_size: 100 19 | 20 | model: 21 | env_id: StackCube-v0 22 | policy: 23 | action_dim: 7 24 | qpos_dim: 9 25 | goal_cond_dim: 0 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_pcd_task/TurnFaucet-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_pcd_dataset 8 | 9 | data: 10 | pin_memory: true 11 | train: 12 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_act.ManiSkill2GoalPosSingleTaskACTPCDDataset 13 | dataset_file: data/maniskill2/demos/v0/rigid_body/TurnFaucet-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 14 | goal_cond_keys: 15 | - target_angle_diff 16 | - target_joint_axis 17 | - target_link_pos 18 | 19 | model: 20 | env_id: TurnFaucet-v0 21 | obs_mode: pointcloud 22 | policy: 23 | action_dim: 7 24 | qpos_dim: 9 25 | goal_cond_dim: 7 26 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_task/PickCube-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_rgbd_dataset 8 | - override /model: maniskill2_act_model 9 | 10 | data: 11 | pin_memory: true 12 | train: 13 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_act.ManiSkill2GoalPosSingleTaskACTRGBDDataset 14 | dataset_file: data/maniskill2/demos/v0/rigid_body/PickCube-v0/trajectory.rgbd.pd_ee_delta_pose.h5 15 | camera_names: 16 | - base_camera 17 | goal_cond_keys: 18 | - goal_pos 19 | chunk_size: 100 20 | 21 | model: 22 | env_id: PickCube-v0 23 | policy: 24 | action_dim: 7 25 | qpos_dim: 9 26 | goal_cond_dim: 3 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_task/Excavate-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_rgbd_dataset 8 | - override /model: maniskill2_act_model 9 | 10 | data: 11 | pin_memory: true 12 | train: 13 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_act.ManiSkill2GoalPosSingleTaskACTRGBDDataset 14 | dataset_file: data/maniskill2/demos/v0/soft_body/Excavate-v0/trajectory.rgbd.pd_ee_delta_pose.h5 15 | camera_names: 16 | - base_camera 17 | goal_cond_keys: 18 | - target 19 | chunk_size: 100 20 | 21 | model: 22 | env_id: Excavate-v0 23 | policy: 24 | action_dim: 6 25 | qpos_dim: 7 26 | goal_cond_dim: 1 27 | -------------------------------------------------------------------------------- /libs/pointops/src/knn_query/knn_query_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "knn_query_cuda_kernel.h" 5 | 6 | 7 | void knn_query_cuda(int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor) 8 | { 9 | const float *xyz = xyz_tensor.data_ptr(); 10 | const float *new_xyz = new_xyz_tensor.data_ptr(); 11 | const int *offset = offset_tensor.data_ptr(); 12 | const int *new_offset = new_offset_tensor.data_ptr(); 13 | int *idx = idx_tensor.data_ptr(); 14 | float *dist2 = dist2_tensor.data_ptr(); 15 | knn_query_cuda_launcher(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2); 16 | } 17 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/pretrained_multimae_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_diffusion_policy_rgbd_dataset 4 | - override /model: maniskill2_diffusion_policy_model 5 | - override /paths: default 6 | 7 | task_name: pretrained_multimae_rgbd 8 | 9 | data: 10 | train: 11 | include_depth: true 12 | 13 | model: 14 | policy: 15 | shape_meta: 16 | obs: 17 | base_camera_rgb: 18 | shape: [3, 128, 128] 19 | type: rgb 20 | base_camera_depth: 21 | shape: [1, 128, 128] 22 | type: depth 23 | qpos: 24 | shape: [7] 25 | type: low_dim 26 | obs_encoder: 27 | rgb_model: 28 | _target_: src.models.components.img_encoder.multivit.MultiViTModel 29 | use_depth: true -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_vit_depth_only.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_vit_depth_only 7 | 8 | data: 9 | train: 10 | include_depth: true 11 | 12 | model: 13 | policy: 14 | shape_meta: 15 | obs: 16 | base_camera_rgb: 17 | shape: [3, 128, 128] 18 | type: rgb 19 | base_camera_depth: 20 | shape: [1, 128, 128] 21 | type: depth 22 | qpos: 23 | shape: [7] 24 | type: low_dim 25 | obs_encoder: 26 | use_depth: true 27 | only_depth: true 28 | rgb_model: 29 | _target_: src.models.components.img_encoder.vit.ViT 30 | model_name: vit_base_patch16 31 | channels: 1 32 | -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/rlbench_model/scratch_vit_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_diffusion_policy_rgbd_dataset 4 | - override /model: rlbench_diffusion_policy_model 5 | 6 | task_name: scratch_vit_rgbd 7 | 8 | data: 9 | train: 10 | include_depth: true 11 | 12 | model: 13 | policy: 14 | shape_meta: 15 | obs: 16 | front_rgb: 17 | shape: [3, 128, 128] 18 | type: rgb 19 | front_depth: 20 | shape: [1, 128, 128] 21 | type: depth 22 | qpos: 23 | shape: 24 | - ${data.train.action_dim} 25 | type: low_dim 26 | obs_encoder: 27 | rgb_model: 28 | _target_: src.models.components.img_encoder.vit.ViT 29 | model_name: vit_base_patch16 30 | channels: 4 31 | use_depth: true -------------------------------------------------------------------------------- /libs/pointops/src/grouping/grouping_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _GROUPING_CUDA_KERNEL 2 | #define _GROUPING_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void grouping_forward_cuda(int m, int nsample, int c, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void grouping_backward_cuda(int m, int nsample, int c, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor grad_input_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void grouping_forward_cuda_launcher(int m, int nsample, int c, const float *input, const int *idx, float *output); 15 | void grouping_backward_cuda_launcher(int m, int nsample, int c, const float *grad_output, const int *idx, float *grad_input); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_multivit_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_diffusion_policy_rgbd_dataset 4 | - override /model: maniskill2_diffusion_policy_model 5 | - override /paths: default 6 | 7 | task_name: scratch_multivit_rgbd 8 | 9 | data: 10 | train: 11 | include_depth: true 12 | 13 | model: 14 | policy: 15 | shape_meta: 16 | obs: 17 | base_camera_rgb: 18 | shape: [3, 128, 128] 19 | type: rgb 20 | base_camera_depth: 21 | shape: [1, 128, 128] 22 | type: depth 23 | qpos: 24 | shape: [7] 25 | type: low_dim 26 | obs_encoder: 27 | rgb_model: 28 | _target_: src.models.components.img_encoder.multivit.MultiViTModel 29 | ckpt_path: null 30 | use_depth: true 31 | -------------------------------------------------------------------------------- /src/models/components/loss/misc.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class KLDivergence(nn.Module): 7 | def __init__(self): 8 | super().__init__() 9 | 10 | def forward(self, mu, logvar): 11 | if mu is None: 12 | return 0 13 | 14 | batch_size = mu.size(0) 15 | assert batch_size != 0 16 | if mu.data.ndimension() == 4: 17 | mu = mu.view(mu.size(0), mu.size(1)) 18 | if logvar.data.ndimension() == 4: 19 | logvar = logvar.view(logvar.size(0), logvar.size(1)) 20 | 21 | klds = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()) 22 | total_kld = klds.sum(1).mean(0, True) 23 | dimension_wise_kld = klds.mean(0) 24 | mean_kld = klds.mean(1).mean(0, True) 25 | 26 | return total_kld[0] # , dimension_wise_kld, mean_kld 27 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_resnet50_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_resnet50_rgbd 7 | 8 | data: 9 | train: 10 | include_depth: true 11 | 12 | model: 13 | policy: 14 | shape_meta: 15 | obs: 16 | base_camera_rgb: 17 | shape: [3, 128, 128] 18 | type: rgb 19 | base_camera_depth: 20 | shape: [1, 128, 128] 21 | type: depth 22 | qpos: 23 | shape: [7] 24 | type: low_dim 25 | obs_encoder: 26 | use_depth: true 27 | rgb_model: 28 | _target_: src.models.components.img_encoder.resnet.ResNetTorchVision 29 | resnet_model: resnet50 30 | pretrained: false 31 | channels: 4 32 | avg_pool: true -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_task/TurnFaucet-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_rgbd_dataset 8 | - override /model: maniskill2_act_model 9 | 10 | data: 11 | pin_memory: true 12 | train: 13 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_act.ManiSkill2GoalPosSingleTaskACTRGBDDataset 14 | dataset_file: data/maniskill2/demos/v0/rigid_body/TurnFaucet-v0/trajectory.rgbd.pd_ee_delta_pose.h5 15 | camera_names: 16 | - base_camera 17 | goal_cond_keys: 18 | - target_angle_diff 19 | - target_joint_axis 20 | - target_link_pos 21 | chunk_size: 100 22 | 23 | model: 24 | env_id: TurnFaucet-v0 25 | policy: 26 | action_dim: 7 27 | qpos_dim: 9 28 | goal_cond_dim: 7 -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/rlbench_model/pretrained_multimae_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_diffusion_policy_rgbd_dataset 4 | - override /model: rlbench_diffusion_policy_model 5 | 6 | task_name: scratch_multivit_rgbd 7 | 8 | data: 9 | batch_size_train: 16 10 | train: 11 | include_depth: true 12 | 13 | trainer: 14 | accumulate_grad_batches: 2 15 | 16 | model: 17 | policy: 18 | shape_meta: 19 | obs: 20 | front_rgb: 21 | shape: [3, 128, 128] 22 | type: rgb 23 | front_depth: 24 | shape: [1, 128, 128] 25 | type: depth 26 | qpos: 27 | shape: 28 | - ${data.train.action_dim} 29 | type: low_dim 30 | obs_encoder: 31 | rgb_model: 32 | _target_: src.models.components.img_encoder.multivit.MultiViTModel 33 | type: MultiViT 34 | use_depth: true -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/rlbench_model/scratch_multivit_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_diffusion_policy_rgbd_dataset 4 | - override /model: rlbench_diffusion_policy_model 5 | 6 | task_name: scratch_multivit_rgbd 7 | 8 | data: 9 | batch_size_train: 16 10 | train: 11 | include_depth: true 12 | 13 | trainer: 14 | accumulate_grad_batches: 2 15 | 16 | model: 17 | policy: 18 | shape_meta: 19 | obs: 20 | front_rgb: 21 | shape: [3, 128, 128] 22 | type: rgb 23 | front_depth: 24 | shape: [1, 128, 128] 25 | type: depth 26 | qpos: 27 | shape: 28 | - ${data.train.action_dim} 29 | type: low_dim 30 | obs_encoder: 31 | rgb_model: 32 | _target_: src.models.components.img_encoder.multivit.MultiViTModel 33 | ckpt_path: null 34 | use_depth: true -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_resnet50_depth_only.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_resnet50_depth_only 7 | 8 | data: 9 | train: 10 | include_depth: true 11 | 12 | model: 13 | policy: 14 | shape_meta: 15 | obs: 16 | base_camera_rgb: 17 | shape: [3, 128, 128] 18 | type: rgb 19 | base_camera_depth: 20 | shape: [1, 128, 128] 21 | type: depth 22 | qpos: 23 | shape: [7] 24 | type: low_dim 25 | obs_encoder: 26 | use_depth: true 27 | only_depth: true 28 | rgb_model: 29 | _target_: src.models.components.img_encoder.resnet.ResNetTorchVision 30 | resnet_model: resnet50 31 | pretrained: false 32 | channels: 1 33 | avg_pool: true -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/rlbench_model/scratch_resnet50_rgbd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_diffusion_policy_rgbd_dataset 4 | - override /model: rlbench_diffusion_policy_model 5 | 6 | task_name: scratch_resnet50_rgbd 7 | 8 | data: 9 | train: 10 | include_depth: true 11 | 12 | model: 13 | policy: 14 | shape_meta: 15 | obs: 16 | front_rgb: 17 | shape: [3, 128, 128] 18 | type: rgb 19 | front_depth: 20 | shape: [1, 128, 128] 21 | type: depth 22 | qpos: 23 | shape: 24 | - ${data.train.action_dim} 25 | type: low_dim 26 | obs_encoder: 27 | rgb_model: 28 | _target_: src.models.components.img_encoder.resnet.ResNetTorchVision 29 | resnet_model: resnet50 30 | pretrained: false 31 | channels: 4 32 | avg_pool: true 33 | use_depth: true -------------------------------------------------------------------------------- /libs/pointops/functions/sampling.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pointops._C import farthest_point_sampling_cuda 3 | from torch.autograd import Function 4 | 5 | 6 | class FarthestPointSampling(Function): 7 | @staticmethod 8 | def forward(ctx, xyz, offset, new_offset): 9 | """ 10 | input: coords: (n, 3), offset: (b), new_offset: (b) 11 | output: idx: (m) 12 | """ 13 | assert xyz.is_contiguous() 14 | n, b, n_max = xyz.shape[0], offset.shape[0], offset[0] 15 | for i in range(1, b): 16 | n_max = max(offset[i] - offset[i - 1], n_max) 17 | idx = torch.cuda.IntTensor(new_offset[b - 1].item()).zero_() 18 | tmp = torch.cuda.FloatTensor(n).fill_(1e10) 19 | farthest_point_sampling_cuda( 20 | b, n_max, xyz, offset.int(), new_offset.int(), tmp, idx 21 | ) 22 | del tmp 23 | return idx 24 | 25 | 26 | farthest_point_sampling = FarthestPointSampling.apply 27 | -------------------------------------------------------------------------------- /libs/pointops/src/interpolation/interpolation_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _INTERPOLATION_CUDA_KERNEL 2 | #define _INTERPOLATION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void interpolation_forward_cuda(int n, int c, int k, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor output_tensor); 8 | void interpolation_backward_cuda(int n, int c, int k, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_input_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void interpolation_forward_cuda_launcher(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output); 15 | void interpolation_backward_cuda_launcher(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /libs/pointops/src/subtraction/subtraction_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _SUBTRACTION_CUDA_KERNEL 2 | #define _SUBTRACTION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void subtraction_forward_cuda(int n, int nsample, int c, at::Tensor input1_tensor, at::Tensor input2_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void subtraction_backward_cuda(int n, int nsample, int c, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input1_tensor, at::Tensor grad_input2_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void subtraction_forward_cuda_launcher(int n, int nsample, int c, const float *input1, const float *input2, const int *idx, float *output); 15 | void subtraction_backward_cuda_launcher(int n, int nsample, int c, const int *idx, const float *grad_output, float *grad_input1, float *grad_input2); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/shadow.vert: -------------------------------------------------------------------------------- 1 | #version 450 2 | #extension GL_ARB_separate_shader_objects : enable 3 | #extension GL_ARB_shading_language_420pack : enable 4 | 5 | layout(set = 0, binding = 0) uniform LightBuffer { 6 | mat4 viewMatrix; 7 | mat4 viewMatrixInverse; 8 | mat4 projectionMatrix; 9 | mat4 projectionMatrixInverse; 10 | int width; 11 | int height; 12 | } lightBuffer; 13 | 14 | layout(set = 1, binding = 0) uniform ObjectBuffer { 15 | mat4 modelMatrix; 16 | mat4 prevModelMatrix; 17 | uvec4 segmentation; 18 | float transparency; 19 | int shadeFlat; 20 | } objectBuffer; 21 | 22 | layout(location = 0) in vec3 position; 23 | layout(location = 1) in vec3 normal; 24 | layout(location = 2) in vec2 uv; 25 | layout(location = 3) in vec3 tangent; 26 | layout(location = 4) in vec3 bitangent; 27 | 28 | void main() { 29 | gl_Position = lightBuffer.projectionMatrix * lightBuffer.viewMatrix * objectBuffer.modelMatrix * vec4(position, 1.f); 30 | } 31 | -------------------------------------------------------------------------------- /libs/pointops/src/ball_query/ball_query_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "ball_query_cuda_kernel.h" 5 | 6 | 7 | void ball_query_cuda(int m, int nsample, 8 | float min_radius, float max_radius, 9 | at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, 10 | at::Tensor offset_tensor, at::Tensor new_offset_tensor, 11 | at::Tensor idx_tensor, at::Tensor dist2_tensor) 12 | { 13 | const float *xyz = xyz_tensor.data_ptr(); 14 | const float *new_xyz = new_xyz_tensor.data_ptr(); 15 | const int *offset = offset_tensor.data_ptr(); 16 | const int *new_offset = new_offset_tensor.data_ptr(); 17 | int *idx = idx_tensor.data_ptr(); 18 | float *dist2 = dist2_tensor.data_ptr(); 19 | ball_query_cuda_launcher(m, nsample, min_radius, max_radius, xyz, new_xyz, offset, new_offset, idx, dist2); 20 | } 21 | -------------------------------------------------------------------------------- /libs/pointops/setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | from distutils.sysconfig import get_config_vars 3 | 4 | from setuptools import setup 5 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 6 | 7 | (opt,) = get_config_vars("OPT") 8 | os.environ["OPT"] = " ".join( 9 | flag for flag in opt.split() if flag != "-Wstrict-prototypes" 10 | ) 11 | 12 | src = "src" 13 | sources = [ 14 | os.path.join(root, file) 15 | for root, dirs, files in os.walk(src) 16 | for file in files 17 | if file.endswith(".cpp") or file.endswith(".cu") 18 | ] 19 | 20 | setup( 21 | name="pointops", 22 | version="1.0", 23 | install_requires=["torch", "numpy"], 24 | packages=["pointops"], 25 | package_dir={"pointops": "functions"}, 26 | ext_modules=[ 27 | CUDAExtension( 28 | name="pointops._C", 29 | sources=sources, 30 | extra_compile_args={"cxx": ["-g"], "nvcc": ["-O2"]}, 31 | ) 32 | ], 33 | cmdclass={"build_ext": BuildExtension}, 34 | ) 35 | -------------------------------------------------------------------------------- /libs/pointops/src/grouping/grouping_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "grouping_cuda_kernel.h" 5 | 6 | 7 | void grouping_forward_cuda(int m, int nsample, int c, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 8 | { 9 | const float *input = input_tensor.data_ptr(); 10 | const int *idx = idx_tensor.data_ptr(); 11 | float *output = output_tensor.data_ptr(); 12 | grouping_forward_cuda_launcher(m, nsample, c, input, idx, output); 13 | } 14 | 15 | void grouping_backward_cuda(int m, int nsample, int c, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor grad_input_tensor) 16 | { 17 | const float *grad_output = grad_output_tensor.data_ptr(); 18 | const int *idx = idx_tensor.data_ptr(); 19 | float *grad_input = grad_input_tensor.data_ptr(); 20 | grouping_backward_cuda_launcher(m, nsample, c, grad_output, idx, grad_input); 21 | } 22 | -------------------------------------------------------------------------------- /libs/pointops/src/ball_query/ball_query_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _BALL_QUERY_CUDA_KERNEL 2 | #define _BALL_QUERY_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void ball_query_cuda(int m, int nsample, 8 | float min_radius, float max_radius, 9 | at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, 10 | at::Tensor offset_tensor, at::Tensor new_offset_tensor, 11 | at::Tensor idx_tensor, at::Tensor dist2_tensor); 12 | 13 | #ifdef __cplusplus 14 | extern "C" { 15 | #endif 16 | 17 | void ball_query_cuda_launcher(int m, int nsample, 18 | float min_radius, float max_radius, 19 | const float *xyz, const float *new_xyz, 20 | const int *offset, const int *new_offset, 21 | int *idx, float *dist2); 22 | 23 | #ifdef __cplusplus 24 | } 25 | #endif 26 | #endif 27 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_pcd_task/Fill-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_pcd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 8 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyPCDDataset 16 | dataset_file: data/maniskill2/demos/v0/soft_body/Fill-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 17 | goal_cond_keys: 18 | - target 19 | 20 | model: 21 | env_id: Fill-v0 22 | policy: 23 | shape_meta: 24 | action: 25 | shape: 26 | - 6 27 | obs: 28 | qpos: 29 | shape: 30 | - 7 31 | type: low_dim 32 | goal: 33 | task_emb: 34 | shape: [2] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_pcd_task/Hang-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_pcd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 8 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyPCDDataset 16 | dataset_file: data/maniskill2/demos/v0/soft_body/Hang-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 17 | goal_cond_keys: 18 | - target 19 | 20 | model: 21 | env_id: Hang-v0 22 | policy: 23 | shape_meta: 24 | action: 25 | shape: 26 | - 7 27 | obs: 28 | qpos: 29 | shape: 30 | - 9 31 | type: low_dim 32 | goal: 33 | task_emb: 34 | shape: [7] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_pcd_task/Pour-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_pcd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 8 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyPCDDataset 16 | dataset_file: data/maniskill2/demos/v0/soft_body/Pour-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 17 | goal_cond_keys: 18 | - target 19 | 20 | model: 21 | env_id: Pour-v0 22 | policy: 23 | shape_meta: 24 | action: 25 | shape: 26 | - 7 27 | obs: 28 | qpos: 29 | shape: 30 | - 9 31 | type: low_dim 32 | goal: 33 | task_emb: 34 | shape: [1] -------------------------------------------------------------------------------- /src/models/components/diffusion_policy/base_image_policy.py: -------------------------------------------------------------------------------- 1 | """ 2 | Reference: 3 | - https://github.com/real-stanford/diffusion_policy 4 | """ 5 | 6 | from typing import Dict 7 | 8 | import torch 9 | import torch.nn as nn 10 | 11 | from src.utils.diffusion_policy import LinearNormalizer, ModuleAttrMixin 12 | 13 | 14 | class BaseImagePolicy(ModuleAttrMixin): 15 | # init accepts keyword argument shape_meta, see config/task/*_image.yaml 16 | 17 | def predict_action( 18 | self, obs_dict: Dict[str, torch.Tensor] 19 | ) -> Dict[str, torch.Tensor]: 20 | """ 21 | obs_dict: 22 | str: B,To,* 23 | return: B,Ta,Da 24 | """ 25 | raise NotImplementedError() 26 | 27 | # reset state for stateful policies 28 | def reset(self): 29 | pass 30 | 31 | # ========== training =========== 32 | # no standard training interface except setting normalizer 33 | def set_normalizer(self, normalizer: LinearNormalizer): 34 | raise NotImplementedError() 35 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_pcd_task/Excavate-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_pcd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 8 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyPCDDataset 16 | dataset_file: data/maniskill2/demos/v0/soft_body/Excavate-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 17 | goal_cond_keys: 18 | - target 19 | 20 | model: 21 | env_id: Excavate-v0 22 | policy: 23 | shape_meta: 24 | action: 25 | shape: 26 | - 6 27 | obs: 28 | qpos: 29 | shape: 30 | - 7 31 | type: low_dim 32 | goal: 33 | task_emb: 34 | shape: [1] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_pcd_task/PickCube-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_pcd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 8 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyPCDDataset 16 | dataset_file: data/maniskill2/demos/v0/rigid_body/PickCube-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 17 | goal_cond_keys: 18 | - goal_pos 19 | 20 | model: 21 | env_id: PickCube-v0 22 | policy: 23 | shape_meta: 24 | action: 25 | shape: 26 | - 7 27 | obs: 28 | qpos: 29 | shape: 30 | - 9 31 | type: low_dim 32 | goal: 33 | task_emb: 34 | shape: [3] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_pcd_task/StackCube-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_pcd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 8 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_dp.ManiSkill2NullGoalSingleTaskDiffusionPolicyPCDDataset 16 | dataset_file: data/maniskill2/demos/v0/rigid_body/StackCube-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 17 | goal_cond_keys: null 18 | 19 | model: 20 | env_id: StackCube-v0 21 | policy: 22 | shape_meta: 23 | action: 24 | shape: 25 | - 7 26 | obs: 27 | qpos: 28 | shape: 29 | - 9 30 | type: low_dim 31 | goal: null 32 | # task_emb: 33 | # shape: [0] -------------------------------------------------------------------------------- /configs/debug/default.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # default debugging setup, runs 1 full epoch 4 | # other debugging configs can inherit from this one 5 | 6 | # overwrite task name so debugging logs are stored in separate folder 7 | task_name: "debug" 8 | 9 | extras: 10 | ignore_warnings: false 11 | enforce_tags: false 12 | 13 | # sets level of all command line loggers to 'DEBUG' 14 | # https://hydra.cc/docs/tutorials/basic/running_your_app/logging/ 15 | hydra: 16 | job_logging: 17 | root: 18 | level: DEBUG 19 | 20 | # use this to also set hydra loggers to 'DEBUG' 21 | # verbose: True 22 | 23 | trainer: 24 | max_epochs: 1 25 | accelerator: cpu # debuggers don't like gpus 26 | devices: 1 # debuggers don't like multiprocessing 27 | detect_anomaly: true # raise exception if NaN or +/-inf is detected in any tensor 28 | limit_train_batches: 2 29 | limit_val_batches: 1 30 | limit_test_batches: 1 31 | 32 | data: 33 | num_workers: 0 # debuggers don't like multiprocessing 34 | pin_memory: false # disable gpu memory pin 35 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_task/Fill-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_rgbd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 16 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyRGBDDataset 16 | dataset_file: data/maniskill2/demos/v0/soft_body/Fill-v0/trajectory.rgbd.pd_ee_delta_pose.h5 17 | camera_names: 18 | - base_camera 19 | goal_cond_keys: 20 | - target 21 | 22 | model: 23 | env_id: Fill-v0 24 | policy: 25 | shape_meta: 26 | action: 27 | shape: 28 | - 6 29 | obs: 30 | qpos: 31 | shape: 32 | - 7 33 | type: low_dim 34 | goal: 35 | task_emb: 36 | shape: [2] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_task/Hang-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_rgbd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 16 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyRGBDDataset 16 | dataset_file: data/maniskill2/demos/v0/soft_body/Hang-v0/trajectory.rgbd.pd_ee_delta_pose.h5 17 | camera_names: 18 | - base_camera 19 | goal_cond_keys: 20 | - target 21 | 22 | model: 23 | env_id: Hang-v0 24 | policy: 25 | shape_meta: 26 | action: 27 | shape: 28 | - 7 29 | obs: 30 | qpos: 31 | shape: 32 | - 9 33 | type: low_dim 34 | goal: 35 | task_emb: 36 | shape: [7] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_task/Pour-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_rgbd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 16 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyRGBDDataset 16 | dataset_file: data/maniskill2/demos/v0/soft_body/Pour-v0/trajectory.rgbd.pd_ee_delta_pose.h5 17 | camera_names: 18 | - base_camera 19 | goal_cond_keys: 20 | - target 21 | 22 | model: 23 | env_id: Pour-v0 24 | policy: 25 | shape_meta: 26 | action: 27 | shape: 28 | - 7 29 | obs: 30 | qpos: 31 | shape: 32 | - 9 33 | type: low_dim 34 | goal: 35 | task_emb: 36 | shape: [1] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_task/Excavate-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_rgbd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 16 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyRGBDDataset 16 | dataset_file: data/maniskill2/demos/v0/soft_body/Excavate-v0/trajectory.rgbd.pd_ee_delta_pose.h5 17 | camera_names: 18 | - base_camera 19 | goal_cond_keys: 20 | - target 21 | 22 | model: 23 | env_id: Excavate-v0 24 | policy: 25 | shape_meta: 26 | action: 27 | shape: 28 | - 6 29 | obs: 30 | qpos: 31 | shape: 32 | - 7 33 | type: low_dim 34 | goal: 35 | task_emb: 36 | shape: [1] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_task/StackCube-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_rgbd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 16 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_dp.ManiSkill2NullGoalSingleTaskDiffusionPolicyRGBDDataset 16 | dataset_file: data/maniskill2/demos/v0/rigid_body/StackCube-v0/trajectory.rgbd.pd_ee_delta_pose.h5 17 | camera_names: 18 | - base_camera 19 | goal_cond_keys: null 20 | 21 | model: 22 | env_id: StackCube-v0 23 | policy: 24 | shape_meta: 25 | action: 26 | shape: 27 | - 7 28 | obs: 29 | qpos: 30 | shape: 31 | - 9 32 | type: low_dim 33 | goal: null 34 | # task_emb: 35 | # shape: [0] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_task/PickCube-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_rgbd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 16 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyRGBDDataset 16 | dataset_file: data/maniskill2/demos/v0/rigid_body/PickCube-v0/trajectory.rgbd.pd_ee_delta_pose.h5 17 | camera_names: 18 | - base_camera 19 | goal_cond_keys: 20 | - goal_pos 21 | 22 | model: 23 | env_id: PickCube-v0 24 | policy: 25 | shape_meta: 26 | action: 27 | shape: 28 | - 7 29 | obs: 30 | qpos: 31 | shape: 32 | - 9 33 | type: low_dim 34 | goal: 35 | task_emb: 36 | shape: [3] -------------------------------------------------------------------------------- /src/utils/__init__.py: -------------------------------------------------------------------------------- 1 | import src.utils.io as io_utils 2 | from src.utils.dist import ( 3 | get_rank, 4 | get_world_size, 5 | is_dist_avail_and_initialized, 6 | is_main_process, 7 | ) 8 | from src.utils.instantiators import instantiate_callbacks, instantiate_loggers 9 | from src.utils.logging_utils import log_hyperparameters 10 | from src.utils.misc import ( 11 | TemporalAgg, 12 | import_modules_from_strings, 13 | interpolate_linear, 14 | is_seq_of, 15 | make_dirs, 16 | ) 17 | from src.utils.pylogger import RankedLogger 18 | from src.utils.registry import Registry, build_from_cfg 19 | from src.utils.rich_utils import enforce_tags, print_config_tree 20 | from src.utils.sparse_tensor_utils import ( 21 | batch2offset, 22 | off_diagonal, 23 | offset2batch, 24 | pcd_collate_fn, 25 | point_collate_fn, 26 | ) 27 | from src.utils.utils import extras, get_metric_value, task_wrapper 28 | 29 | from . import diffusion_policy as dp_utils 30 | from .optimizer import build_optimizer, build_optimizer_v2 31 | from .scheduler import build_scheduler 32 | -------------------------------------------------------------------------------- /configs/data/rlbench_act_rgbd_dataset.yaml: -------------------------------------------------------------------------------- 1 | _target_: src.data.base_datamodule.BaseDataModule 2 | 3 | batch_size_train: 32 4 | batch_size_val: 1 5 | batch_size_test: 1 6 | num_workers: 16 7 | pin_memory: true 8 | train: 9 | _target_: src.data.components.rlbench.rlbench_single_task_act.RLBenchSingleTaskACTRGBDDataset 10 | root: data/rlbench/processed/train/ 11 | task_names: 12 | - ${rlbench_task} 13 | chunk_size: 100 14 | max_episodes_per_task: 100 15 | cameras: 16 | - front 17 | action_dim: 11 18 | collision: true 19 | rot_type: 6d 20 | include_depth: false 21 | loop: 1 22 | 23 | val: 24 | _target_: src.data.components.rlbench.rlbench_single_task_act.RLBenchSingleTaskACTRGBDDataset 25 | root: data/rlbench/processed/val/ 26 | task_names: ${data.train.task_names} 27 | chunk_size: ${data.train.chunk_size} 28 | max_episodes_per_task: 25 29 | cameras: ${data.train.cameras} 30 | action_dim: ${data.train.action_dim} 31 | collision: ${data.train.collision} 32 | rot_type: ${data.train.rot_type} 33 | include_depth: ${data.train.include_depth} 34 | loop: 1 35 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_spunet_pcd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_spunet_pcd 7 | 8 | model: 9 | obs_mode: pointcloud 10 | policy: 11 | shape_meta: 12 | obs: 13 | pcds: 14 | shape: [6] 15 | type: pcd 16 | qpos: 17 | shape: [7] 18 | type: low_dim 19 | obs_encoder: 20 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 21 | shape_meta: ${model.policy.shape_meta} 22 | pcd_model: 23 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 24 | in_channels: 6 25 | num_classes: 0 26 | pretrained_path: null 27 | share_pcd_model: true 28 | n_obs_step: ${model.policy.n_obs_steps} 29 | pcd_nsample: 16 30 | pcd_npoints: 2048 31 | use_mask: false 32 | bg_ratio: 0.0 33 | pcd_hidden_dim: 96 34 | projector_layers: 2 35 | projector_channels: [96, 128, 128] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_pcd_task/TurnFaucet-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_pcd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 8 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyPCDDataset 16 | dataset_file: data/maniskill2/demos/v0/rigid_body/TurnFaucet-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 17 | goal_cond_keys: 18 | - target_angle_diff 19 | - target_joint_axis 20 | - target_link_pos 21 | 22 | model: 23 | env_id: TurnFaucet-v0 24 | policy: 25 | shape_meta: 26 | action: 27 | shape: 28 | - 7 29 | obs: 30 | qpos: 31 | shape: 32 | - 9 33 | type: low_dim 34 | goal: 35 | task_emb: 36 | shape: [7] -------------------------------------------------------------------------------- /src/data/components/misc.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, Iterable 2 | 3 | import torch 4 | from torch.utils.data import Dataset, IterableDataset 5 | 6 | 7 | class DummyDataset(Dataset): 8 | def __init__(self, size: int): 9 | self.size = size 10 | 11 | def __getitem__(self, index: int): 12 | return torch.tensor(index) 13 | 14 | def __len__(self): 15 | return self.size 16 | 17 | 18 | class ExperienceSourceDataset(IterableDataset): 19 | """ 20 | Implementation from PyTorch Lightning Bolts: 21 | https://github.com/PyTorchLightning/pytorch-lightning-bolts/blob/master/pl_bolts/datamodules/experience_source.py 22 | 23 | Basic experience source dataset. Takes a generate_batch function that returns an iterator. 24 | The logic for the experience source and how the batch is generated is defined the Lightning model itself 25 | """ 26 | 27 | def __init__(self, generate_batch: Callable): 28 | self.generate_batch = generate_batch 29 | 30 | def __iter__(self) -> Iterable: 31 | iterator = self.generate_batch() 32 | return iterator 33 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_pointnet_pcd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_pointnet_pcd 7 | 8 | model: 9 | obs_mode: pointcloud 10 | policy: 11 | shape_meta: 12 | obs: 13 | pcds: 14 | shape: [6] 15 | type: pcd 16 | qpos: 17 | shape: [7] 18 | type: low_dim 19 | obs_encoder: 20 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 21 | shape_meta: ${model.policy.shape_meta} 22 | pcd_model: 23 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 24 | in_channels: 6 25 | num_classes: 96 26 | pretrained_path: null 27 | share_pcd_model: true 28 | n_obs_step: ${model.policy.n_obs_steps} 29 | pcd_nsample: 16 30 | pcd_npoints: 2048 31 | use_mask: false 32 | bg_ratio: 0.0 33 | pcd_hidden_dim: 96 34 | projector_layers: 1 35 | projector_channels: [96, 128, 128] -------------------------------------------------------------------------------- /configs/data/rlbench_diffusion_policy_rgbd_dataset.yaml: -------------------------------------------------------------------------------- 1 | _target_: src.data.base_datamodule.BaseDataModule 2 | 3 | batch_size_train: 32 4 | batch_size_val: 1 5 | batch_size_test: 1 6 | num_workers: 16 7 | pin_memory: true 8 | train: 9 | _target_: src.data.components.rlbench.rlbench_single_task_dp.RLBenchSingleTaskDiffusionPolicyRGBDDataset 10 | root: data/rlbench/processed/train/ 11 | task_names: 12 | - ${rlbench_task} 13 | chunk_size: 16 14 | max_episodes_per_task: 100 15 | cameras: 16 | - front 17 | action_dim: 11 18 | collision: true 19 | rot_type: 6d 20 | include_depth: false 21 | loop: 1 22 | 23 | val: 24 | _target_: src.data.components.rlbench.rlbench_single_task_dp.RLBenchSingleTaskDiffusionPolicyRGBDDataset 25 | root: data/rlbench/processed/val/ 26 | task_names: ${data.train.task_names} 27 | chunk_size: ${data.train.chunk_size} 28 | max_episodes_per_task: 25 29 | cameras: ${data.train.cameras} 30 | action_dim: ${data.train.action_dim} 31 | collision: ${data.train.collision} 32 | rot_type: ${data.train.rot_type} 33 | include_depth: ${data.train.include_depth} 34 | loop: 1 35 | 36 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Haoyi Zhu 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /libs/pointops/src/random_ball_query/random_ball_query_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _RANDOM_BALL_QUERY_CUDA_KERNEL 2 | #define _RANDOM_BALL_QUERY_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void random_ball_query_cuda(int m, int nsample, 8 | float min_radius, float max_radius, at::Tensor order_tensor, 9 | at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, 10 | at::Tensor offset_tensor, at::Tensor new_offset_tensor, 11 | at::Tensor idx_tensor, at::Tensor dist2_tensor); 12 | 13 | #ifdef __cplusplus 14 | extern "C" { 15 | #endif 16 | 17 | void random_ball_query_cuda_launcher(int m, int nsample, 18 | float min_radius, float max_radius, const int *order, 19 | const float *xyz, const float *new_xyz, 20 | const int *offset, const int *new_offset, 21 | int *idx, float *dist2); 22 | 23 | #ifdef __cplusplus 24 | } 25 | #endif 26 | #endif 27 | -------------------------------------------------------------------------------- /libs/pointops/src/random_ball_query/random_ball_query_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "random_ball_query_cuda_kernel.h" 5 | 6 | 7 | void random_ball_query_cuda(int m, int nsample, 8 | float min_radius, float max_radius, at::Tensor order_tensor, 9 | at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, 10 | at::Tensor offset_tensor, at::Tensor new_offset_tensor, 11 | at::Tensor idx_tensor, at::Tensor dist2_tensor) 12 | { 13 | const int *order = order_tensor.data_ptr(); 14 | const float *xyz = xyz_tensor.data_ptr(); 15 | const float *new_xyz = new_xyz_tensor.data_ptr(); 16 | const int *offset = offset_tensor.data_ptr(); 17 | const int *new_offset = new_offset_tensor.data_ptr(); 18 | int *idx = idx_tensor.data_ptr(); 19 | float *dist2 = dist2_tensor.data_ptr(); 20 | random_ball_query_cuda_launcher(m, nsample, min_radius, max_radius, order, xyz, new_xyz, offset, new_offset, idx, dist2); 21 | } 22 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_task/TurnFaucet-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_rgbd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 16 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_dp.ManiSkill2GoalPosSingleTaskDiffusionPolicyRGBDDataset 16 | dataset_file: data/maniskill2/demos/v0/rigid_body/TurnFaucet-v0/trajectory.rgbd.pd_ee_delta_pose.h5 17 | camera_names: 18 | - base_camera 19 | goal_cond_keys: 20 | - target_angle_diff 21 | - target_joint_axis 22 | - target_link_pos 23 | 24 | model: 25 | env_id: TurnFaucet-v0 26 | policy: 27 | shape_meta: 28 | action: 29 | shape: 30 | - 7 31 | obs: 32 | qpos: 33 | shape: 34 | - 9 35 | type: low_dim 36 | goal: 37 | task_emb: 38 | shape: [7] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/base.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | defaults: 4 | - maniskill2_model: 5 | - maniskill2_task: 6 | - maniskill2_pcd_task: 7 | - override /trainer: ddp 8 | - override /callbacks: default 9 | - override /paths: default 10 | - override /hydra: default 11 | 12 | # all parameters below will be merged with parameters from default configurations set above 13 | # this allows you to overwrite only specified parameters 14 | 15 | hydra: 16 | run: 17 | dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S} 18 | sweep: 19 | dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S} 20 | subdir: "" 21 | 22 | paths: 23 | log_dir: ${paths.root_dir}/logs/maniskill2_dp/${model.env_id} 24 | 25 | trainer: 26 | devices: 1 27 | max_epochs: 1800 28 | check_val_every_n_epoch: 100 29 | accelerator: gpu 30 | strategy: auto 31 | # limit_val_batches: 0.0 32 | 33 | callbacks: 34 | model_checkpoint: 35 | monitor: val/mean_success 36 | mode: max 37 | filename: "epoch={epoch:03d}-val_mean_success={val/mean_success:.4f}" 38 | save_top_k: 5 39 | early_stopping: null 40 | -------------------------------------------------------------------------------- /configs/data/maniskill2_act_pcd_dataset.yaml: -------------------------------------------------------------------------------- 1 | _target_: src.data.base_datamodule.BaseDataModule 2 | 3 | batch_size_train: 16 4 | batch_size_val: 1 5 | batch_size_test: 1 6 | num_workers: 16 7 | pin_memory: true 8 | train: 9 | dataset_file: 10 | load_count: -1 11 | goal_cond_keys: 12 | chunk_size: 100 # equal to horizon 13 | camera_ids: [0] 14 | point_num_per_cam: 16384 # 128 * 128 15 | transform_pcd: 16 | - _target_: src.data.components.transformpcd.GridSamplePCD 17 | grid_size: 0.005 18 | hash_type: fnv 19 | mode: train 20 | return_grid_coord: true 21 | return_displacement: false 22 | keys: 23 | - coord 24 | - color 25 | - _target_: src.data.components.transformpcd.NormalizeColorPCD 26 | - _target_: src.data.components.transformpcd.ShufflePointPCD 27 | - _target_: src.data.components.transformpcd.ToTensorPCD 28 | - _target_: src.data.components.transformpcd.CollectPCD 29 | keys: 30 | - coord 31 | - grid_coord 32 | feat_keys: 33 | - color 34 | - coord 35 | 36 | val: 37 | _target_: src.data.components.misc.DummyDataset 38 | size: 400 # 400 eval trials -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/pretrained_ponderv2_pcd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: pretrained_ponderv2_pcd 7 | 8 | model: 9 | obs_mode: pointcloud 10 | policy: 11 | shape_meta: 12 | obs: 13 | pcds: 14 | shape: [6] 15 | type: pcd 16 | qpos: 17 | shape: [7] 18 | type: low_dim 19 | obs_encoder: 20 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 21 | shape_meta: ${model.policy.shape_meta} 22 | pcd_model: 23 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 24 | in_channels: 6 25 | num_classes: 0 26 | pretrained_path: checkpoints/ponderv2-ppt-pretrain-scannet-s3dis-structured3d.pth 27 | share_pcd_model: true 28 | n_obs_step: ${model.policy.n_obs_steps} 29 | pcd_nsample: 16 30 | pcd_npoints: 2048 31 | use_mask: false 32 | bg_ratio: 0.0 33 | pcd_hidden_dim: 96 34 | projector_layers: 2 35 | projector_channels: [96, 128, 128] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_spunet_pcd_presample.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_spunet_pcd_presample 7 | 8 | model: 9 | obs_mode: pointcloud 10 | policy: 11 | shape_meta: 12 | obs: 13 | pcds: 14 | shape: [6] 15 | type: pcd 16 | qpos: 17 | shape: [7] 18 | type: low_dim 19 | obs_encoder: 20 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 21 | shape_meta: ${model.policy.shape_meta} 22 | pcd_model: 23 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 24 | in_channels: 6 25 | num_classes: 0 26 | pretrained_path: null 27 | share_pcd_model: true 28 | n_obs_step: ${model.policy.n_obs_steps} 29 | pcd_nsample: 16 30 | pcd_npoints: 2048 31 | use_mask: false 32 | bg_ratio: 0.0 33 | pcd_hidden_dim: 96 34 | projector_layers: 2 35 | projector_channels: [96, 128, 128] 36 | pre_sample: true 37 | in_channel: 6 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/base.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | defaults: 4 | - maniskill2_model: 5 | - maniskill2_task: 6 | - maniskill2_pcd_task: 7 | - override /trainer: ddp 8 | - override /callbacks: default 9 | - override /paths: default 10 | - override /hydra: default 11 | 12 | # all parameters below will be merged with parameters from default configurations set above 13 | # this allows you to overwrite only specified parameters 14 | 15 | hydra: 16 | run: 17 | dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S} 18 | sweep: 19 | dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S} 20 | subdir: '' 21 | 22 | paths: 23 | log_dir: ${paths.root_dir}/logs/maniskill2_act/${model.env_id} 24 | 25 | 26 | trainer: 27 | devices: 1 28 | max_epochs: 500 29 | check_val_every_n_epoch: 100 30 | accelerator: gpu 31 | strategy: auto 32 | # limit_val_batches: 0.0 33 | 34 | callbacks: 35 | model_checkpoint: 36 | monitor: val/mean_success 37 | mode: max 38 | filename: "epoch={epoch:03d}-val_mean_success={val/mean_success:.4f}" 39 | save_top_k: 5 40 | early_stopping: null 41 | 42 | 43 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_pointnet_pcd_presample.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_pointnet_pcd_presample 7 | 8 | model: 9 | obs_mode: pointcloud 10 | policy: 11 | shape_meta: 12 | obs: 13 | pcds: 14 | shape: [6] 15 | type: pcd 16 | qpos: 17 | shape: [7] 18 | type: low_dim 19 | obs_encoder: 20 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 21 | shape_meta: ${model.policy.shape_meta} 22 | pcd_model: 23 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 24 | in_channels: 6 25 | num_classes: 96 26 | pretrained_path: null 27 | share_pcd_model: true 28 | n_obs_step: ${model.policy.n_obs_steps} 29 | pcd_nsample: 16 30 | pcd_npoints: 2048 31 | use_mask: false 32 | bg_ratio: 0.0 33 | pcd_hidden_dim: 96 34 | projector_layers: 1 35 | projector_channels: [96, 128, 128] 36 | pre_sample: true 37 | in_channel: 6 -------------------------------------------------------------------------------- /src/utils/spconv_utils.py: -------------------------------------------------------------------------------- 1 | from typing import Set 2 | 3 | import spconv 4 | 5 | if float(spconv.__version__[2:]) >= 2.2: 6 | spconv.constants.SPCONV_USE_DIRECT_TABLE = False 7 | 8 | try: 9 | import spconv.pytorch as spconv 10 | except: 11 | import spconv as spconv 12 | 13 | import torch.nn as nn 14 | 15 | 16 | def find_all_spconv_keys(model: nn.Module, prefix="") -> Set[str]: 17 | """ 18 | Finds all spconv keys that need to have weight's transposed 19 | """ 20 | found_keys: Set[str] = set() 21 | for name, child in model.named_children(): 22 | new_prefix = f"{prefix}.{name}" if prefix != "" else name 23 | 24 | if isinstance(child, spconv.conv.SparseConvolution): 25 | new_prefix = f"{new_prefix}.weight" 26 | found_keys.add(new_prefix) 27 | 28 | found_keys.update(find_all_spconv_keys(child, prefix=new_prefix)) 29 | 30 | return found_keys 31 | 32 | 33 | def replace_feature(out, new_features): 34 | if "replace_feature" in out.__dir__(): 35 | # spconv 2.x behaviour 36 | return out.replace_feature(new_features) 37 | else: 38 | out.features = new_features 39 | return out 40 | -------------------------------------------------------------------------------- /configs/callbacks/early_stopping.yaml: -------------------------------------------------------------------------------- 1 | # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.EarlyStopping.html 2 | 3 | early_stopping: 4 | _target_: lightning.pytorch.callbacks.EarlyStopping 5 | monitor: ??? # quantity to be monitored, must be specified !!! 6 | min_delta: 0. # minimum change in the monitored quantity to qualify as an improvement 7 | patience: 3 # number of checks with no improvement after which training will be stopped 8 | verbose: false # verbosity mode 9 | mode: "min" # "max" means higher metric value is better, can be also "min" 10 | strict: true # whether to crash the training if monitor is not found in the validation metrics 11 | check_finite: true # when set True, stops training when the monitor becomes NaN or infinite 12 | stopping_threshold: # stop training immediately once the monitored quantity reaches this threshold 13 | divergence_threshold: # stop training as soon as the monitored quantity becomes worse than this threshold 14 | check_on_train_epoch_end: # whether to run early stopping at the end of the training epoch 15 | # log_rank_zero_only: False # this keyword argument isn't available in stable version 16 | -------------------------------------------------------------------------------- /libs/pointops/src/interpolation/interpolation_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "interpolation_cuda_kernel.h" 5 | 6 | 7 | void interpolation_forward_cuda(int n, int c, int k, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor output_tensor) 8 | { 9 | const float *input = input_tensor.data_ptr(); 10 | const int *idx = idx_tensor.data_ptr(); 11 | const float *weight = weight_tensor.data_ptr(); 12 | float *output = output_tensor.data_ptr(); 13 | interpolation_forward_cuda_launcher(n, c, k, input, idx, weight, output); 14 | } 15 | 16 | void interpolation_backward_cuda(int n, int c, int k, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_input_tensor) 17 | { 18 | const float *grad_output = grad_output_tensor.data_ptr(); 19 | const int *idx = idx_tensor.data_ptr(); 20 | const float *weight = weight_tensor.data_ptr(); 21 | float *grad_input = grad_input_tensor.data_ptr(); 22 | interpolation_backward_cuda_launcher(n, c, k, grad_output, idx, weight, grad_input); 23 | } 24 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/common/shadow.glsl: -------------------------------------------------------------------------------- 1 | vec3 project(mat4 proj, vec3 point) { 2 | vec4 v = proj * vec4(point, 1); 3 | return v.xyz / v.w; 4 | } 5 | 6 | const int PCF_SampleCount = 25; 7 | vec2 PCF_Samples[PCF_SampleCount] = { 8 | {-2, -2}, {-1, -2}, {0, -2}, {1, -2}, {2, -2}, 9 | {-2, -1}, {-1, -1}, {0, -1}, {1, -1}, {2, -1}, 10 | {-2, 0}, {-1, 0}, {0, 0}, {1, 0}, {2, 0}, 11 | {-2, 1}, {-1, 1}, {0, 1}, {1, 1}, {2, 1}, 12 | {-2, 2}, {-1, 2}, {0, 2}, {1, 2}, {2, 2} 13 | }; 14 | 15 | float ShadowMapPCF( 16 | sampler2D shadowTex, vec3 projCoord, float resolution, float searchUV, float filterSize) 17 | { 18 | float shadow = 0.0f; 19 | vec2 grad = fract(projCoord.xy * resolution + 0.5f); 20 | 21 | for (int i = 0; i < PCF_SampleCount; i++) 22 | { 23 | vec4 tmp = textureGather(shadowTex, projCoord.xy + filterSize * PCF_Samples[i] * searchUV); 24 | tmp.x = tmp.x < projCoord.z ? 0.0f : 1.0f; 25 | tmp.y = tmp.y < projCoord.z ? 0.0f : 1.0f; 26 | tmp.z = tmp.z < projCoord.z ? 0.0f : 1.0f; 27 | tmp.w = tmp.w < projCoord.z ? 0.0f : 1.0f; 28 | shadow += mix(mix(tmp.w, tmp.z, grad.x), mix(tmp.x, tmp.y, grad.x), grad.y); 29 | } 30 | return shadow / PCF_SampleCount; 31 | } 32 | -------------------------------------------------------------------------------- /configs/data/maniskill2_diffusion_policy_pcd_dataset.yaml: -------------------------------------------------------------------------------- 1 | _target_: src.data.base_datamodule.BaseDataModule 2 | 3 | batch_size_train: 32 4 | batch_size_val: 1 5 | batch_size_test: 1 6 | num_workers: 16 7 | pin_memory: true 8 | train: 9 | dataset_file: 10 | load_count: -1 11 | goal_cond_keys: 12 | n_obs_steps: ${model.policy.n_obs_steps} 13 | chunk_size: 16 # equal to horizon 14 | camera_ids: [0] 15 | point_num_per_cam: 16384 16 | transform_pcd: 17 | - _target_: src.data.components.transformpcd.GridSamplePCD 18 | grid_size: 0.005 19 | hash_type: fnv 20 | mode: train 21 | return_grid_coord: true 22 | return_displacement: false 23 | keys: 24 | - coord 25 | - color 26 | - _target_: src.data.components.transformpcd.NormalizeColorPCD 27 | - _target_: src.data.components.transformpcd.ShufflePointPCD 28 | - _target_: src.data.components.transformpcd.ToTensorPCD 29 | - _target_: src.data.components.transformpcd.CollectPCD 30 | keys: 31 | - coord 32 | - grid_coord 33 | feat_keys: 34 | - color 35 | - coord 36 | 37 | val: 38 | _target_: src.data.components.misc.DummyDataset 39 | size: 400 # 400 eval trials -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/rlbench_model/pretrained_ponderv2_pcd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_diffusion_policy_pcd_dataset 4 | - override /model: rlbench_diffusion_policy_model 5 | 6 | task_name: pretrained_ponderv2_pcd 7 | 8 | data: 9 | batch_size_train: 16 10 | 11 | trainer: 12 | accumulate_grad_batches: 2 13 | 14 | model: 15 | policy: 16 | shape_meta: 17 | obs: 18 | pcds: 19 | shape: [6] 20 | type: pcd 21 | qpos: 22 | shape: 23 | - ${data.train.action_dim} 24 | type: low_dim 25 | obs_encoder: 26 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 27 | shape_meta: ${model.policy.shape_meta} 28 | pcd_model: 29 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 30 | in_channels: 6 31 | num_classes: 0 32 | pretrained_path: checkpoints/ponderv2-ppt-pretrain-scannet-s3dis-structured3d.pth 33 | pcd_nsample: 16 34 | pcd_npoints: 2048 35 | use_mask: false 36 | bg_ratio: 0.0 37 | pcd_hidden_dim: 96 38 | projector_layers: 2 39 | projector_channels: [96, 128, 128] -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/gbuffer2.frag: -------------------------------------------------------------------------------- 1 | #version 450 2 | #extension GL_ARB_separate_shader_objects : enable 3 | #extension GL_ARB_shading_language_420pack : enable 4 | 5 | layout(set = 2, binding = 0) uniform MaterialBuffer { 6 | vec4 emission; 7 | vec4 baseColor; 8 | float fresnel; 9 | float roughness; 10 | float metallic; 11 | float transmission; 12 | float ior; 13 | float transmissionRoughness; 14 | int textureMask; 15 | int padding1; 16 | } materialBuffer; 17 | 18 | layout(set = 2, binding = 1) uniform sampler2D colorTexture; 19 | layout(set = 2, binding = 2) uniform sampler2D roughnessTexture; 20 | layout(set = 2, binding = 3) uniform sampler2D normalTexture; 21 | layout(set = 2, binding = 4) uniform sampler2D metallicTexture; 22 | 23 | layout(location = 0) in vec4 inPosition; 24 | layout(location = 1) in vec2 inUV; 25 | 26 | layout(location = 0) out vec4 outAlbedo2; 27 | 28 | void main() { 29 | if ((materialBuffer.textureMask & 1) != 0) { 30 | outAlbedo2 = texture(colorTexture, inUV); 31 | outAlbedo2.rgb = outAlbedo2.rgb; // sRGB to linear 32 | } else { 33 | outAlbedo2 = materialBuffer.baseColor; 34 | } 35 | if (outAlbedo2.a == 0) { 36 | discard; 37 | } 38 | } 39 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/gbuffer2.vert: -------------------------------------------------------------------------------- 1 | #version 450 2 | #extension GL_ARB_separate_shader_objects : enable 3 | #extension GL_ARB_shading_language_420pack : enable 4 | 5 | layout(set = 0, binding = 0) uniform CameraBuffer { 6 | mat4 viewMatrix; 7 | mat4 projectionMatrix; 8 | mat4 viewMatrixInverse; 9 | mat4 projectionMatrixInverse; 10 | mat4 prevViewMatrix; 11 | mat4 prevViewMatrixInverse; 12 | float width; 13 | float height; 14 | } cameraBuffer; 15 | 16 | layout(set = 1, binding = 0) uniform ObjectBuffer { 17 | mat4 modelMatrix; 18 | mat4 prevModelMatrix; 19 | uvec4 segmentation; 20 | float transparency; 21 | int shadeFlat; 22 | } objectBuffer; 23 | 24 | layout(location = 0) in vec3 position; 25 | layout(location = 1) in vec3 normal; 26 | layout(location = 2) in vec2 uv; 27 | layout(location = 3) in vec3 tangent; 28 | layout(location = 4) in vec3 bitangent; 29 | 30 | layout(location = 0) out vec4 outPosition; 31 | layout(location = 1) out vec2 outUV; 32 | 33 | void main() { 34 | mat4 modelView = cameraBuffer.viewMatrix * objectBuffer.modelMatrix; 35 | outPosition = modelView * vec4(position, 1); 36 | outUV = uv; 37 | gl_Position = cameraBuffer.projectionMatrix * outPosition; 38 | } 39 | -------------------------------------------------------------------------------- /configs/callbacks/model_checkpoint.yaml: -------------------------------------------------------------------------------- 1 | # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html 2 | 3 | model_checkpoint: 4 | _target_: lightning.pytorch.callbacks.ModelCheckpoint 5 | dirpath: # directory to save the model file 6 | filename: # checkpoint filename 7 | monitor: # name of the logged metric which determines when model is improving 8 | verbose: false # verbosity mode 9 | save_last: # additionally always save an exact copy of the last checkpoint to a file last.ckpt 10 | save_top_k: 10 # save k best models (determined by above metric) 11 | mode: "min" # "max" means higher metric value is better, can be also "min" 12 | auto_insert_metric_name: true # when True, the checkpoints filenames will contain the metric name 13 | save_weights_only: false # if True, then only the model’s weights will be saved 14 | every_n_train_steps: # number of training steps between checkpoints 15 | train_time_interval: # checkpoints are monitored at the specified time interval 16 | every_n_epochs: # number of epochs between checkpoints 17 | save_on_train_epoch_end: # whether to run checkpointing at the end of the training epoch or the end of validation 18 | -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/rlbench_model/scratch_spunet_pcd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_diffusion_policy_pcd_dataset 4 | - override /model: rlbench_diffusion_policy_model 5 | 6 | task_name: scratch_spunet_pcd 7 | 8 | data: 9 | batch_size_train: 16 10 | 11 | trainer: 12 | accumulate_grad_batches: 2 13 | 14 | model: 15 | policy: 16 | shape_meta: 17 | obs: 18 | pcds: 19 | shape: [6] 20 | type: pcd 21 | qpos: 22 | shape: 23 | - ${data.train.action_dim} 24 | type: low_dim 25 | obs_encoder: 26 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 27 | shape_meta: ${model.policy.shape_meta} 28 | pcd_model: 29 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 30 | in_channels: 6 31 | num_classes: 0 32 | pretrained_path: null 33 | share_pcd_model: true 34 | n_obs_step: ${model.policy.n_obs_steps} 35 | pcd_nsample: 16 36 | pcd_npoints: 2048 37 | use_mask: false 38 | bg_ratio: 0.0 39 | pcd_hidden_dim: 96 40 | projector_layers: 2 41 | projector_channels: [96, 128, 128] -------------------------------------------------------------------------------- /libs/pointops/src/subtraction/subtraction_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "subtraction_cuda_kernel.h" 5 | 6 | 7 | void subtraction_forward_cuda(int n, int nsample, int c, at::Tensor input1_tensor, at::Tensor input2_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 8 | { 9 | const float *input1 = input1_tensor.data_ptr(); 10 | const float *input2 = input2_tensor.data_ptr(); 11 | const int *idx = idx_tensor.data_ptr(); 12 | float *output = output_tensor.data_ptr(); 13 | subtraction_forward_cuda_launcher(n, nsample, c, input1, input2, idx, output); 14 | } 15 | 16 | void subtraction_backward_cuda(int n, int nsample, int c, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input1_tensor, at::Tensor grad_input2_tensor) 17 | { 18 | const int *idx = idx_tensor.data_ptr(); 19 | const float *grad_output = grad_output_tensor.data_ptr(); 20 | float *grad_input1 = grad_input1_tensor.data_ptr(); 21 | float *grad_input2 = grad_input2_tensor.data_ptr(); 22 | subtraction_backward_cuda_launcher(n, nsample, c, idx, grad_output, grad_input1, grad_input2); 23 | } 24 | -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/rlbench_model/scratch_pointnet_pcd.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: rlbench_diffusion_policy_pcd_dataset 4 | - override /model: rlbench_diffusion_policy_model 5 | 6 | task_name: scratch_pointnet_pcd 7 | 8 | data: 9 | batch_size_train: 16 10 | 11 | trainer: 12 | accumulate_grad_batches: 2 13 | 14 | model: 15 | policy: 16 | shape_meta: 17 | obs: 18 | pcds: 19 | shape: [6] 20 | type: pcd 21 | qpos: 22 | shape: 23 | - ${data.train.action_dim} 24 | type: low_dim 25 | obs_encoder: 26 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 27 | shape_meta: ${model.policy.shape_meta} 28 | pcd_model: 29 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 30 | in_channels: 6 31 | num_classes: 96 32 | pretrained_path: null 33 | share_pcd_model: true 34 | n_obs_step: ${model.policy.n_obs_steps} 35 | pcd_nsample: 16 36 | pcd_npoints: 2048 37 | use_mask: false 38 | bg_ratio: 0.0 39 | pcd_hidden_dim: 96 40 | projector_layers: 1 41 | projector_channels: [96, 128, 128] -------------------------------------------------------------------------------- /libs/pointops/src/aggregation/aggregation_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _AGGREGATION_CUDA_KERNEL 2 | #define _AGGREGATION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void aggregation_forward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void aggregation_backward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input_tensor, at::Tensor grad_position_tensor, at::Tensor grad_weight_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void aggregation_forward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output); 15 | void aggregation_backward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/shadow_point.vert: -------------------------------------------------------------------------------- 1 | #version 450 2 | #extension GL_ARB_separate_shader_objects : enable 3 | #extension GL_ARB_shading_language_420pack : enable 4 | 5 | layout(set = 0, binding = 0) uniform LightBuffer { 6 | mat4 viewMatrix; 7 | mat4 viewMatrixInverse; 8 | mat4 projectionMatrix; 9 | mat4 projectionMatrixInverse; 10 | int width; 11 | int height; 12 | } lightBuffer; 13 | 14 | layout(set = 1, binding = 0) uniform ObjectBuffer { 15 | mat4 modelMatrix; 16 | mat4 prevModelMatrix; 17 | uvec4 segmentation; 18 | float transparency; 19 | int shadeFlat; 20 | } objectBuffer; 21 | 22 | layout(location = 0) in vec3 position; 23 | layout(location = 1) in float scale; 24 | layout(location = 2) in vec4 color; 25 | 26 | layout(location = 0) out vec4 outPosition; 27 | layout(location = 1) out flat vec4 outNdcRadius; 28 | 29 | void main() { 30 | mat4 modelView = lightBuffer.viewMatrix * objectBuffer.modelMatrix; 31 | outPosition = modelView * vec4(position, 1); 32 | 33 | float radius = scale; 34 | 35 | gl_PointSize = lightBuffer.projectionMatrix[0][0] * lightBuffer.width * radius; 36 | 37 | gl_Position = lightBuffer.projectionMatrix * outPosition; 38 | outNdcRadius = vec4(gl_Position.xyz / gl_Position.w, radius); 39 | } 40 | -------------------------------------------------------------------------------- /configs/exp_rlbench_act_policy/base.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | defaults: 4 | - rlbench_model: 5 | - override /trainer: ddp 6 | - override /callbacks: default 7 | - override /paths: default 8 | - override /hydra: default 9 | 10 | # all parameters below will be merged with parameters from default configurations set above 11 | # this allows you to overwrite only specified parameters 12 | 13 | rlbench_task: null 14 | 15 | hydra: 16 | run: 17 | dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S} 18 | sweep: 19 | dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S} 20 | subdir: "" 21 | 22 | paths: 23 | log_dir: ${paths.root_dir}/logs/rlbench_act/${rlbench_task} 24 | 25 | data: 26 | num_workers: 16 27 | pin_memory: false 28 | train: 29 | task_names: 30 | - ${rlbench_task} 31 | loop: 1 32 | cache_episode: true 33 | 34 | trainer: 35 | devices: 1 36 | max_epochs: 2000 37 | check_val_every_n_epoch: 100 38 | accelerator: gpu 39 | strategy: auto 40 | 41 | callbacks: 42 | model_checkpoint: 43 | monitor: val/loss 44 | mode: min 45 | filename: "epoch={epoch:03d}-val_loss={val/loss:.4f}" 46 | save_top_k: 3 47 | early_stopping: 48 | monitor: val/loss 49 | mode: "min" 50 | -------------------------------------------------------------------------------- /configs/exp_rlbench_diffusion_policy/base.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | defaults: 4 | - rlbench_model: 5 | - override /trainer: ddp 6 | - override /callbacks: default 7 | - override /paths: default 8 | - override /hydra: default 9 | 10 | # all parameters below will be merged with parameters from default configurations set above 11 | # this allows you to overwrite only specified parameters 12 | 13 | rlbench_task: null 14 | 15 | hydra: 16 | run: 17 | dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S} 18 | sweep: 19 | dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S} 20 | subdir: "" 21 | 22 | paths: 23 | log_dir: ${paths.root_dir}/logs/rlbench_dp/${rlbench_task} 24 | 25 | data: 26 | num_workers: 16 27 | pin_memory: false 28 | train: 29 | task_names: 30 | - ${rlbench_task} 31 | loop: 1 32 | cache_episode: true 33 | 34 | trainer: 35 | devices: 1 36 | max_epochs: 2000 37 | check_val_every_n_epoch: 100 38 | accelerator: gpu 39 | strategy: auto 40 | 41 | callbacks: 42 | model_checkpoint: 43 | monitor: val/loss 44 | mode: min 45 | filename: "epoch={epoch:03d}-val_loss={val/loss:.4f}" 46 | save_top_k: 3 47 | early_stopping: 48 | monitor: val/loss 49 | mode: "min" 50 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/line.vert: -------------------------------------------------------------------------------- 1 | #version 450 2 | #extension GL_ARB_separate_shader_objects : enable 3 | #extension GL_ARB_shading_language_420pack : enable 4 | 5 | layout(set = 0, binding = 0) uniform CameraBuffer { 6 | mat4 viewMatrix; 7 | mat4 projectionMatrix; 8 | mat4 viewMatrixInverse; 9 | mat4 projectionMatrixInverse; 10 | mat4 prevViewMatrix; 11 | mat4 prevViewMatrixInverse; 12 | float width; 13 | float height; 14 | } cameraBuffer; 15 | 16 | layout(set = 1, binding = 0) uniform ObjectBuffer { 17 | mat4 modelMatrix; 18 | mat4 prevModelMatrix; 19 | uvec4 segmentation; 20 | float transparency; 21 | int shadeFlat; 22 | } objectBuffer; 23 | 24 | layout(location = 0) in vec3 position; 25 | layout(location = 1) in float scale; 26 | layout(location = 2) in vec4 color; 27 | 28 | layout(location = 0) out vec4 outPosition; 29 | layout(location = 1) out vec4 outPrevPosition; 30 | layout(location = 2) out vec4 outColor; 31 | 32 | void main() { 33 | mat4 modelView = cameraBuffer.viewMatrix * objectBuffer.modelMatrix; 34 | outPosition = modelView * vec4(position, 1); 35 | outPrevPosition = cameraBuffer.prevViewMatrix * objectBuffer.prevModelMatrix * vec4(position, 1); 36 | gl_Position = cameraBuffer.projectionMatrix * outPosition; 37 | outColor = color; 38 | } 39 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_pointnet_pcd_wo_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_pointnet_pcd_wo_rgb 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | train: 13 | transform_pcd: 14 | - _target_: src.data.components.transformpcd.GridSamplePCD 15 | grid_size: 0.005 16 | hash_type: fnv 17 | mode: train 18 | return_grid_coord: true 19 | return_displacement: false 20 | keys: 21 | - coord 22 | - color 23 | - _target_: src.data.components.transformpcd.NormalizeColorPCD 24 | - _target_: src.data.components.transformpcd.ShufflePointPCD 25 | - _target_: src.data.components.transformpcd.ToTensorPCD 26 | - _target_: src.data.components.transformpcd.CollectPCD 27 | keys: 28 | - coord 29 | - grid_coord 30 | feat_keys: 31 | - coord 32 | 33 | model: 34 | obs_mode: pointcloud 35 | policy: 36 | backbone: 37 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 38 | in_channels: 3 39 | num_classes: 0 40 | 41 | trainer: 42 | accumulate_grad_batches: 2 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_pointnet_pcd_wo_xyz.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_pointnet_pcd_wo_xyz 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | train: 13 | transform_pcd: 14 | - _target_: src.data.components.transformpcd.GridSamplePCD 15 | grid_size: 0.005 16 | hash_type: fnv 17 | mode: train 18 | return_grid_coord: true 19 | return_displacement: false 20 | keys: 21 | - coord 22 | - color 23 | - _target_: src.data.components.transformpcd.NormalizeColorPCD 24 | - _target_: src.data.components.transformpcd.ShufflePointPCD 25 | - _target_: src.data.components.transformpcd.ToTensorPCD 26 | - _target_: src.data.components.transformpcd.CollectPCD 27 | keys: 28 | - coord 29 | - grid_coord 30 | feat_keys: 31 | - color 32 | 33 | model: 34 | obs_mode: pointcloud 35 | policy: 36 | backbone: 37 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 38 | in_channels: 3 39 | num_classes: 0 40 | 41 | trainer: 42 | accumulate_grad_batches: 2 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_spunet_pcd_wo_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_spunet_pcd_wo_rgb 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | train: 13 | transform_pcd: 14 | - _targt_: src.data.components.transformpcd.GridSamplePCD 15 | grid_size: 0.005 16 | hash_type: fnv 17 | mode: train 18 | return_grid_coord: true 19 | return_displacement: false 20 | keys: 21 | - coord 22 | - color 23 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 24 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 25 | - _targt_: src.data.components.transformpcd.ToTensorPCD 26 | - _targt_: src.data.components.transformpcd.CollectPCD 27 | keys: 28 | - coord 29 | - grid_coord 30 | feat_keys: 31 | - coord 32 | 33 | model: 34 | obs_mode: pointcloud 35 | policy: 36 | backbone: 37 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 38 | in_channels: 3 39 | num_classes: 0 40 | pretrained_path: null 41 | 42 | trainer: 43 | accumulate_grad_batches: 2 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_spunet_pcd_wo_xyz.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_spunet_pcd_wo_xyz 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | train: 13 | transform_pcd: 14 | - _targt_: src.data.components.transformpcd.GridSamplePCD 15 | grid_size: 0.005 16 | hash_type: fnv 17 | mode: train 18 | return_grid_coord: true 19 | return_displacement: false 20 | keys: 21 | - coord 22 | - color 23 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 24 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 25 | - _targt_: src.data.components.transformpcd.ToTensorPCD 26 | - _targt_: src.data.components.transformpcd.CollectPCD 27 | keys: 28 | - coord 29 | - grid_coord 30 | feat_keys: 31 | - color 32 | 33 | model: 34 | obs_mode: pointcloud 35 | policy: 36 | backbone: 37 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 38 | in_channels: 3 39 | num_classes: 0 40 | pretrained_path: null 41 | 42 | trainer: 43 | accumulate_grad_batches: 2 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_pointnet_pcd_presample_wo_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_pointnet_pcd_presample_wo_rgb 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | train: 13 | transform_pcd: 14 | - _target_: src.data.components.transformpcd.GridSamplePCD 15 | grid_size: 0.005 16 | hash_type: fnv 17 | mode: train 18 | return_grid_coord: true 19 | return_displacement: false 20 | keys: 21 | - coord 22 | - color 23 | - _target_: src.data.components.transformpcd.NormalizeColorPCD 24 | - _target_: src.data.components.transformpcd.ShufflePointPCD 25 | - _target_: src.data.components.transformpcd.ToTensorPCD 26 | - _target_: src.data.components.transformpcd.CollectPCD 27 | keys: 28 | - coord 29 | - grid_coord 30 | feat_keys: 31 | - coord 32 | 33 | model: 34 | obs_mode: pointcloud 35 | policy: 36 | backbone: 37 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 38 | in_channels: 3 39 | num_classes: 0 40 | pre_sample: true 41 | 42 | trainer: 43 | accumulate_grad_batches: 2 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_pointnet_pcd_presample_wo_xyz.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_pointnet_pcd_presample_wo_xyz 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | train: 13 | transform_pcd: 14 | - _target_: src.data.components.transformpcd.GridSamplePCD 15 | grid_size: 0.005 16 | hash_type: fnv 17 | mode: train 18 | return_grid_coord: true 19 | return_displacement: false 20 | keys: 21 | - coord 22 | - color 23 | - _target_: src.data.components.transformpcd.NormalizeColorPCD 24 | - _target_: src.data.components.transformpcd.ShufflePointPCD 25 | - _target_: src.data.components.transformpcd.ToTensorPCD 26 | - _target_: src.data.components.transformpcd.CollectPCD 27 | keys: 28 | - coord 29 | - grid_coord 30 | feat_keys: 31 | - color 32 | 33 | model: 34 | obs_mode: pointcloud 35 | policy: 36 | policy: 37 | backbone: 38 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 39 | in_channels: 3 40 | num_classes: 0 41 | pre_sample: true 42 | 43 | trainer: 44 | accumulate_grad_batches: 2 -------------------------------------------------------------------------------- /configs/logger/aim.yaml: -------------------------------------------------------------------------------- 1 | # https://aimstack.io/ 2 | 3 | # example usage in lightning module: 4 | # https://github.com/aimhubio/aim/blob/main/examples/pytorch_lightning_track.py 5 | 6 | # open the Aim UI with the following command (run in the folder containing the `.aim` folder): 7 | # `aim up` 8 | 9 | aim: 10 | _target_: aim.pytorch_lightning.AimLogger 11 | repo: ${paths.root_dir} # .aim folder will be created here 12 | # repo: "aim://ip_address:port" # can instead provide IP address pointing to Aim remote tracking server which manages the repo, see https://aimstack.readthedocs.io/en/latest/using/remote_tracking.html# 13 | 14 | # aim allows to group runs under experiment name 15 | experiment: # any string, set to "default" if not specified 16 | 17 | train_metric_prefix: "train/" 18 | val_metric_prefix: "val/" 19 | test_metric_prefix: "test/" 20 | 21 | # sets the tracking interval in seconds for system usage metrics (CPU, GPU, memory, etc.) 22 | system_tracking_interval: 10 # set to null to disable system metrics tracking 23 | 24 | # enable/disable logging of system params such as installed packages, git info, env vars, etc. 25 | log_system_params: true 26 | 27 | # enable/disable tracking console logs (default value is true) 28 | capture_terminal_logs: false # set to false to avoid infinite console log loop issue https://github.com/aimhubio/aim/issues/2550 29 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/shadow_point.frag: -------------------------------------------------------------------------------- 1 | #version 450 2 | #extension GL_ARB_separate_shader_objects : enable 3 | #extension GL_ARB_shading_language_420pack : enable 4 | 5 | layout(location = 0) in vec4 inPosition; 6 | layout(location = 1) in flat vec4 inNdcRadius; 7 | 8 | layout(set = 0, binding = 0) uniform LightBuffer { 9 | mat4 viewMatrix; 10 | mat4 viewMatrixInverse; 11 | mat4 projectionMatrix; 12 | mat4 projectionMatrixInverse; 13 | int width; 14 | int height; 15 | } lightBuffer; 16 | 17 | layout(set = 1, binding = 0) uniform ObjectBuffer { 18 | mat4 modelMatrix; 19 | mat4 prevModelMatrix; 20 | uvec4 segmentation; 21 | float transparency; 22 | int shadeFlat; 23 | } objectBuffer; 24 | 25 | void main() { 26 | if (gl_PointCoord.s * gl_PointCoord.s + gl_PointCoord.t * gl_PointCoord.t > 1) { 27 | discard; 28 | } 29 | // vec2 centerNdc = inNdcRadius.xy; 30 | // vec2 res = vec2(lightBuffer.width, lightBuffer.height) * RESOLUTION_SCALE; 31 | // vec2 pixelNdc = gl_FragCoord.xy / res * 2.0 - 1.0; 32 | // vec2 offsetNdc = pixelNdc - centerNdc; 33 | // vec2 offset = offsetNdc * (-inPosition.z) / vec2(lightBuffer.projectionMatrix[0][0], lightBuffer.projectionMatrix[1][1]); 34 | // float radius = inNdcRadius.w; 35 | // offset /= radius; 36 | // if (offset.x * offset.x + offset.y * offset.y > 1) { 37 | // discard; 38 | // } 39 | } 40 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_spunet_pcd_presample_wo_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_spunet_pcd_presample_wo_rgb 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | train: 13 | transform_pcd: 14 | - _target_: src.data.components.transformpcd.GridSamplePCD 15 | grid_size: 0.005 16 | hash_type: fnv 17 | mode: train 18 | return_grid_coord: true 19 | return_displacement: false 20 | keys: 21 | - coord 22 | - color 23 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 24 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 25 | - _targt_: src.data.components.transformpcd.ToTensorPCD 26 | - _targt_: src.data.components.transformpcd.CollectPCD 27 | keys: 28 | - coord 29 | - grid_coord 30 | feat_keys: 31 | - coord 32 | 33 | model: 34 | obs_mode: pointcloud 35 | policy: 36 | backbone: 37 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 38 | in_channels: 3 39 | num_classes: 512 40 | pretrained_path: null 41 | pre_sample: true 42 | 43 | trainer: 44 | accumulate_grad_batches: 2 -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_model/scratch_spunet_pcd_presample_wo_xyz.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_act_pcd_dataset 4 | - override /model: maniskill2_act_pcd_model 5 | - override /paths: default 6 | 7 | task_name: scratch_spunet_pcd_presample_wo_xyz 8 | 9 | data: 10 | batch_size_train: 8 11 | num_workers: 16 12 | train: 13 | transform_pcd: 14 | - _targt_: src.data.components.transformpcd.GridSamplePCD 15 | grid_size: 0.005 16 | hash_type: fnv 17 | mode: train 18 | return_grid_coord: true 19 | return_displacement: false 20 | keys: 21 | - coord 22 | - color 23 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 24 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 25 | - _targt_: src.data.components.transformpcd.ToTensorPCD 26 | - _targt_: src.data.components.transformpcd.CollectPCD 27 | keys: 28 | - coord 29 | - grid_coord 30 | feat_keys: 31 | - color 32 | 33 | model: 34 | obs_mode: pointcloud 35 | policy: 36 | backbone: 37 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 38 | in_channels: 3 39 | num_classes: 512 40 | pretrained_path: null 41 | pre_sample: true 42 | 43 | trainer: 44 | accumulate_grad_batches: 2 -------------------------------------------------------------------------------- /libs/pointops/functions/subtraction.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pointops._C import subtraction_backward_cuda, subtraction_forward_cuda 3 | from torch.autograd import Function 4 | 5 | 6 | class Subtraction(Function): 7 | @staticmethod 8 | def forward(ctx, input1, input2, idx): 9 | """ 10 | input: input1: (n, c), input2: (n, c), idx: (n, nsample) 11 | output: (n, nsample, c) 12 | """ 13 | assert input1.is_contiguous() and input2.is_contiguous() 14 | n, c = input1.shape 15 | nsample = idx.shape[-1] 16 | output = torch.cuda.FloatTensor(n, nsample, c).zero_() 17 | subtraction_forward_cuda(n, nsample, c, input1, input2, idx, output) 18 | ctx.save_for_backward(idx) 19 | return output 20 | 21 | @staticmethod 22 | def backward(ctx, grad_output): 23 | """ 24 | input: grad_out: (n, nsample, c) 25 | output: grad_input1: (n, c), grad_input2: (n, c) 26 | """ 27 | (idx,) = ctx.saved_tensors 28 | n, nsample, c = grad_output.shape 29 | grad_input1 = torch.cuda.FloatTensor(n, c).zero_() 30 | grad_input2 = torch.cuda.FloatTensor(n, c).zero_() 31 | subtraction_backward_cuda( 32 | n, nsample, c, idx, grad_output, grad_input1, grad_input2 33 | ) 34 | return grad_input1, grad_input2, None 35 | 36 | 37 | subtraction = Subtraction.apply 38 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_pcd_task/PegInsertionSide-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_pcd_dataset 8 | 9 | data: 10 | pin_memory: true 11 | train: 12 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_act.ManiSkill2NullGoalSingleTaskACTPCDDataset 13 | dataset_file: data/maniskill2/demos/v0/rigid_body/PegInsertionSide-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 14 | goal_cond_keys: null 15 | 16 | model: 17 | env_id: PegInsertionSide-v0 18 | obs_mode: pointcloud 19 | policy: 20 | action_dim: 7 21 | qpos_dim: 9 22 | goal_cond_dim: 0 23 | val_metrics: 24 | metrics: 25 | - _target_: torchmetrics.MeanMetric 26 | - _target_: torchmetrics.MeanMetric 27 | - _target_: torchmetrics.MeanMetric 28 | input_keys: 29 | - mean_success 30 | - mean_grasp 31 | - mean_align 32 | output_keys: 33 | - val/mean_success 34 | - val/mean_grasp 35 | - val/mean_align 36 | best_val_metrics: 37 | metrics: 38 | - _target_: torchmetrics.MaxMetric 39 | - _target_: torchmetrics.MaxMetric 40 | - _target_: torchmetrics.MaxMetric 41 | input_keys: 42 | - val/mean_success 43 | - val/mean_grasp 44 | - val/mean_align 45 | output_keys: 46 | - val/mean_success 47 | - val/mean_grasp 48 | - val/mean_align 49 | -------------------------------------------------------------------------------- /src/models/components/diffusion_policy/diffusion/conv1d_components.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | # from einops.layers.torch import Rearrange 6 | 7 | 8 | class Downsample1d(nn.Module): 9 | def __init__(self, dim): 10 | super().__init__() 11 | self.conv = nn.Conv1d(dim, dim, 3, 2, 1) 12 | 13 | def forward(self, x): 14 | return self.conv(x) 15 | 16 | 17 | class Upsample1d(nn.Module): 18 | def __init__(self, dim): 19 | super().__init__() 20 | self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1) 21 | 22 | def forward(self, x): 23 | return self.conv(x) 24 | 25 | 26 | class Conv1dBlock(nn.Module): 27 | """ 28 | Conv1d --> GroupNorm --> Mish 29 | """ 30 | 31 | def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8): 32 | super().__init__() 33 | 34 | self.block = nn.Sequential( 35 | nn.Conv1d( 36 | inp_channels, out_channels, kernel_size, padding=kernel_size // 2 37 | ), 38 | # Rearrange('batch channels horizon -> batch channels 1 horizon'), 39 | nn.GroupNorm(n_groups, out_channels), 40 | # Rearrange('batch channels 1 horizon -> batch channels horizon'), 41 | nn.Mish(), 42 | ) 43 | 44 | def forward(self, x): 45 | return self.block(x) 46 | 47 | 48 | def test(): 49 | cb = Conv1dBlock(256, 128, kernel_size=3) 50 | x = torch.zeros((1, 256, 16)) 51 | o = cb(x) 52 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_act_policy/maniskill2_task/PegInsertionSide-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_act_rgbd_dataset 8 | - override /model: maniskill2_act_model 9 | 10 | data: 11 | pin_memory: true 12 | train: 13 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_act.ManiSkill2NullGoalSingleTaskACTRGBDDataset 14 | dataset_file: data/maniskill2/demos/v0/rigid_body/PegInsertionSide-v0/trajectory.rgbd.pd_ee_delta_pose.h5 15 | camera_names: 16 | - base_camera 17 | goal_cond_keys: null 18 | chunk_size: 100 19 | 20 | model: 21 | env_id: PegInsertionSide-v0 22 | policy: 23 | action_dim: 7 24 | qpos_dim: 9 25 | goal_cond_dim: 0 26 | val_metrics: 27 | metrics: 28 | - _target_: torchmetrics.MeanMetric 29 | - _target_: torchmetrics.MeanMetric 30 | - _target_: torchmetrics.MeanMetric 31 | input_keys: 32 | - mean_success 33 | - mean_grasp 34 | - mean_align 35 | output_keys: 36 | - val/mean_success 37 | - val/mean_grasp 38 | - val/mean_align 39 | best_val_metrics: 40 | metrics: 41 | - _target_: torchmetrics.MaxMetric 42 | - _target_: torchmetrics.MaxMetric 43 | - _target_: torchmetrics.MaxMetric 44 | input_keys: 45 | - val/mean_success 46 | - val/mean_grasp 47 | - val/mean_align 48 | output_keys: 49 | - val/mean_success 50 | - val/mean_grasp 51 | - val/mean_align 52 | -------------------------------------------------------------------------------- /configs/test_rlbench_dp.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # specify here default configuration 4 | # order of defaults determines the order in which configs override each other 5 | defaults: 6 | - _self_ 7 | - data: 8 | - model: 9 | - callbacks: default 10 | - logger: tensorboard # set logger here or use command line (e.g. `python train.py logger=tensorboard`) 11 | - trainer: ddp 12 | - paths: default 13 | - extras: default 14 | - hydra: default 15 | - exp_rlbench_diffusion_policy: 16 | 17 | # optional local config for machine/user specific settings 18 | # it's optional since it doesn't need to exist and is excluded from version control 19 | - optional local: default 20 | 21 | # debugging config (enable through command line, e.g. `python train.py debug=default) 22 | - debug: 23 | 24 | # task name, determines output directory path 25 | task_name: "test" 26 | 27 | # tags to help you identify your experiments 28 | # you can overwrite this in experiment configs 29 | # overwrite from command line with `python train.py tags="[first_tag, second_tag]"` 30 | tags: ["test"] 31 | 32 | # simply provide checkpoint path to resume training 33 | ckpt_path: 34 | 35 | # seed for random number generators in pytorch, numpy and python.random 36 | seed: 1000 37 | 38 | # rlbench test options 39 | live_demos: false 40 | data_root: data/rlbench/raw/test 41 | headless: true 42 | episodes_num: 25 43 | max_steps: 200 44 | 45 | result_path: logs/rlbench_dp/${rlbench_task}/${task_name} 46 | result_file: test_result.txt 47 | result_name: ${ckpt_path} 48 | 49 | offline: false 50 | 51 | camera_view_test: 52 | apply: false 53 | rot_axis: y 54 | rot_angle: 0 55 | transl: [0, 0, 0] -------------------------------------------------------------------------------- /src/utils/diffusion_policy/dict_of_tensor_mixin.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class DictOfTensorMixin(nn.Module): 6 | def __init__(self, params_dict=None): 7 | super().__init__() 8 | if params_dict is None: 9 | params_dict = nn.ParameterDict() 10 | self.params_dict = params_dict 11 | 12 | @property 13 | def device(self): 14 | return next(iter(self.parameters())).device 15 | 16 | def _load_from_state_dict( 17 | self, 18 | state_dict, 19 | prefix, 20 | local_metadata, 21 | strict, 22 | missing_keys, 23 | unexpected_keys, 24 | error_msgs, 25 | ): 26 | def dfs_add(dest, keys, value: torch.Tensor): 27 | if len(keys) == 1: 28 | dest[keys[0]] = value 29 | return 30 | 31 | if keys[0] not in dest: 32 | dest[keys[0]] = nn.ParameterDict() 33 | dfs_add(dest[keys[0]], keys[1:], value) 34 | 35 | def load_dict(state_dict, prefix): 36 | out_dict = nn.ParameterDict() 37 | for key, value in state_dict.items(): 38 | value: torch.Tensor 39 | if key.startswith(prefix): 40 | param_keys = key[len(prefix) :].split(".")[1:] 41 | # if len(param_keys) == 0: 42 | # import pdb; pdb.set_trace() 43 | dfs_add(out_dict, param_keys, value.clone()) 44 | return out_dict 45 | 46 | self.params_dict = load_dict(state_dict, prefix + "params_dict") 47 | self.params_dict.requires_grad_(False) 48 | return 49 | -------------------------------------------------------------------------------- /configs/train.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # specify here default configuration 4 | # order of defaults determines the order in which configs override each other 5 | defaults: 6 | - _self_ 7 | - data: 8 | - model: 9 | - callbacks: default 10 | - logger: tensorboard # set logger here or use command line (e.g. `python train.py logger=tensorboard`) 11 | - trainer: ddp 12 | - paths: default 13 | - extras: default 14 | - hydra: default 15 | - exp_maniskill2_act_policy: 16 | - exp_maniskill2_diffusion_policy: 17 | - exp_rlbench_act_policy: 18 | - exp_rlbench_diffusion_policy: 19 | 20 | # optional local config for machine/user specific settings 21 | # it's optional since it doesn't need to exist and is excluded from version control 22 | - optional local: default 23 | 24 | # debugging config (enable through command line, e.g. `python train.py debug=default) 25 | - debug: 26 | 27 | # task name, determines output directory path 28 | task_name: "train" 29 | 30 | # tags to help you identify your experiments 31 | # you can overwrite this in experiment configs 32 | # overwrite from command line with `python train.py tags="[first_tag, second_tag]"` 33 | tags: ["dev"] 34 | 35 | # set False to skip model training 36 | train: true 37 | 38 | # evaluate on test set, using best model weights achieved during training 39 | # lightning chooses best weights based on the metric specified in checkpoint callback 40 | test: false 41 | 42 | # compile model for faster training with pytorch 2.0 43 | compile: false 44 | 45 | # simply provide checkpoint path to resume training 46 | ckpt_path: 47 | 48 | # seed for random number generators in pytorch, numpy and python.random 49 | seed: 1000 50 | -------------------------------------------------------------------------------- /src/models/components/act/utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | """ 4 | Reference: 5 | - https://github.com/tonyzhaozh/act 6 | """ 7 | from typing import Optional 8 | 9 | import numpy as np 10 | import torch 11 | 12 | 13 | class NestedTensor: 14 | def __init__(self, tensors, mask: Optional[torch.Tensor]): 15 | self.tensors = tensors 16 | self.mask = mask 17 | 18 | def to(self, device): 19 | # type: (Device) -> NestedTensor # noqa 20 | cast_tensor = self.tensors.to(device) 21 | mask = self.mask 22 | if mask is not None: 23 | assert mask is not None 24 | cast_mask = mask.to(device) 25 | else: 26 | cast_mask = None 27 | return NestedTensor(cast_tensor, cast_mask) 28 | 29 | def decompose(self): 30 | return self.tensors, self.mask 31 | 32 | def __repr__(self): 33 | return str(self.tensors) 34 | 35 | 36 | def reparametrize(mu, logvar): 37 | std = logvar.div(2).exp() 38 | eps = torch.autograd.Variable(std.data.new(std.size()).normal_()) 39 | return mu + std * eps 40 | 41 | 42 | def get_sinusoid_encoding_table(n_position, d_hid): 43 | def get_position_angle_vec(position): 44 | return [ 45 | position / np.power(10000, 2 * (hid_j // 2) / d_hid) 46 | for hid_j in range(d_hid) 47 | ] 48 | 49 | sinusoid_table = np.array( 50 | [get_position_angle_vec(pos_i) for pos_i in range(n_position)] 51 | ) 52 | sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i 53 | sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 54 | 55 | return torch.FloatTensor(sinusoid_table).unsqueeze(0) 56 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/gbuffer1.vert: -------------------------------------------------------------------------------- 1 | #version 450 2 | #extension GL_ARB_separate_shader_objects : enable 3 | #extension GL_ARB_shading_language_420pack : enable 4 | 5 | layout(set = 0, binding = 0) uniform CameraBuffer { 6 | mat4 viewMatrix; 7 | mat4 projectionMatrix; 8 | mat4 viewMatrixInverse; 9 | mat4 projectionMatrixInverse; 10 | mat4 prevViewMatrix; 11 | mat4 prevViewMatrixInverse; 12 | float width; 13 | float height; 14 | } cameraBuffer; 15 | 16 | layout(set = 1, binding = 0) uniform ObjectBuffer { 17 | mat4 modelMatrix; 18 | mat4 prevModelMatrix; 19 | uvec4 segmentation; 20 | float transparency; 21 | int shadeFlat; 22 | } objectBuffer; 23 | 24 | layout(location = 0) in vec3 position; 25 | layout(location = 1) in vec3 normal; 26 | layout(location = 2) in vec2 uv; 27 | layout(location = 3) in vec3 tangent; 28 | layout(location = 4) in vec3 bitangent; 29 | 30 | layout(location = 0) out vec4 outPosition; 31 | layout(location = 1) out vec2 outUV; 32 | layout(location = 2) out flat uvec4 outSegmentation; 33 | layout(location = 3) out vec3 objectCoord; 34 | layout(location = 4) out mat3 outTbn; 35 | 36 | void main() { 37 | outSegmentation = objectBuffer.segmentation; 38 | mat4 modelView = cameraBuffer.viewMatrix * objectBuffer.modelMatrix; 39 | mat3 normalMatrix = mat3(transpose(inverse(modelView))); 40 | objectCoord = position; 41 | outPosition = modelView * vec4(position, 1); 42 | outUV = uv; 43 | gl_Position = cameraBuffer.projectionMatrix * outPosition; 44 | vec3 outTangent = normalize(normalMatrix * tangent); 45 | vec3 outBitangent = normalize(normalMatrix * bitangent); 46 | vec3 outNormal = normalize(normalMatrix * normal); 47 | outTbn = mat3(outTangent, outBitangent, outNormal); 48 | } 49 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_pcd_task/PegInsertionSide-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_pcd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 8 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_pcd_dp.ManiSkill2NullGoalSingleTaskDiffusionPolicyPCDDataset 16 | dataset_file: data/maniskill2/demos/v0/rigid_body/PegInsertionSide-v0/trajectory.pointcloud.pd_ee_delta_pose.h5 17 | goal_cond_keys: null 18 | 19 | model: 20 | env_id: PegInsertionSide-v0 21 | policy: 22 | shape_meta: 23 | action: 24 | shape: 25 | - 7 26 | obs: 27 | qpos: 28 | shape: 29 | - 9 30 | type: low_dim 31 | goal: null 32 | val_metrics: 33 | metrics: 34 | - _target_: torchmetrics.MeanMetric 35 | - _target_: torchmetrics.MeanMetric 36 | - _target_: torchmetrics.MeanMetric 37 | input_keys: 38 | - mean_success 39 | - mean_grasp 40 | - mean_align 41 | output_keys: 42 | - val/mean_success 43 | - val/mean_grasp 44 | - val/mean_align 45 | best_val_metrics: 46 | metrics: 47 | - _target_: torchmetrics.MaxMetric 48 | - _target_: torchmetrics.MaxMetric 49 | - _target_: torchmetrics.MaxMetric 50 | input_keys: 51 | - val/mean_success 52 | - val/mean_grasp 53 | - val/mean_align 54 | output_keys: 55 | - val/mean_success 56 | - val/mean_grasp 57 | - val/mean_align -------------------------------------------------------------------------------- /libs/pointops/src/aggregation/aggregation_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "aggregation_cuda_kernel.h" 5 | 6 | 7 | void aggregation_forward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 8 | { 9 | const float *input = input_tensor.data_ptr(); 10 | const float *position = position_tensor.data_ptr(); 11 | const float *weight = weight_tensor.data_ptr(); 12 | const int *idx = idx_tensor.data_ptr(); 13 | float *output = output_tensor.data_ptr(); 14 | aggregation_forward_cuda_launcher(n, nsample, c, w_c, input, position, weight, idx, output); 15 | } 16 | 17 | void aggregation_backward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input_tensor, at::Tensor grad_position_tensor, at::Tensor grad_weight_tensor) 18 | { 19 | const float *input = input_tensor.data_ptr(); 20 | const float *position = position_tensor.data_ptr(); 21 | const float *weight = weight_tensor.data_ptr(); 22 | const int *idx = idx_tensor.data_ptr(); 23 | const float *grad_output = grad_output_tensor.data_ptr(); 24 | float *grad_input = grad_input_tensor.data_ptr(); 25 | float *grad_position = grad_position_tensor.data_ptr(); 26 | float *grad_weight = grad_weight_tensor.data_ptr(); 27 | aggregation_backward_cuda_launcher(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight); 28 | } 29 | -------------------------------------------------------------------------------- /scripts/run_rlbench_camera_view.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | policy=$1 # policy can be either "diffusion" or act 4 | ckpt_path=$2 5 | model=$3 6 | task=$4 7 | seed=$5 8 | 9 | if [ "$policy" != "diffusion" ] && [ "$policy" != "act" ]; then 10 | echo "Policy should be either diffusion or act, but got $policy" 11 | exit 1 12 | fi 13 | 14 | echo "Starting testing ${policy} policy for task: ${task} with model: ${model} and checkpoint path: ${ckpt_path}" 15 | echo "Seed: ${seed}" 16 | 17 | 18 | rot_axes=("y" "z") 19 | rot_angles=(-5 -10) 20 | 21 | for rot_axis in "${rot_axes[@]}"; do 22 | for rot_angle in "${rot_angles[@]}"; do 23 | transl=(0 0 0) 24 | 25 | if [ "$rot_axis" == "y" ]; then 26 | transl[0]=$(echo "$rot_angle * -0.02" | bc) 27 | transl[2]=$(echo "$rot_angle * -0.02" | bc) 28 | fi 29 | 30 | echo "Rotation Axis: $rot_axis, Rotation Angle: $rot_angle, Translation: ${transl[@]}" 31 | 32 | if [ "$policy" == "diffusion" ]; then 33 | test_script=test_rlbench_dp.py 34 | elif [ "$policy" == "act" ]; then 35 | test_script=test_rlbench_act.py 36 | fi 37 | 38 | python src/${test_script} exp_rlbench_${policy}_policy=base \ 39 | rlbench_task=$task exp_rlbench_${policy}_policy/rlbench_model@rlbench_model=$model \ 40 | ckpt_path=$ckpt_path seed=$seed \ 41 | camera_view_test.apply=true camera_view_test.rot_axis=${rot_axis} camera_view_test.rot_angle=${rot_angle} \ 42 | camera_view_test.transl="[${transl[0]},${transl[1]},${transl[2]}]" \ 43 | result_file=test_camera_view_${rot_axis}_${rot_angle}.txt 44 | 45 | done 46 | done 47 | 48 | echo "Testing completed." 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_task/PegInsertionSide-v0.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # to execute this task use: 4 | # exp_maniskill2_diffusion_policy/maniskill2_task@maniskill2_task=StackCube 5 | 6 | defaults: 7 | - override /data: maniskill2_diffusion_policy_rgbd_dataset 8 | - override /model: maniskill2_diffusion_policy_model 9 | 10 | data: 11 | batch_size_train: 64 12 | num_workers: 16 13 | pin_memory: true 14 | train: 15 | _target_: src.data.components.maniskill2.maniskill2_single_task_rgbd_dp.ManiSkill2NullGoalSingleTaskDiffusionPolicyRGBDDataset 16 | dataset_file: data/maniskill2/demos/v0/rigid_body/PegInsertionSide-v0/trajectory.rgbd.pd_ee_delta_pose.h5 17 | camera_names: 18 | - base_camera 19 | goal_cond_keys: null 20 | 21 | model: 22 | env_id: PegInsertionSide-v0 23 | policy: 24 | shape_meta: 25 | action: 26 | shape: 27 | - 7 28 | obs: 29 | qpos: 30 | shape: 31 | - 9 32 | type: low_dim 33 | goal: null 34 | val_metrics: 35 | metrics: 36 | - _target_: torchmetrics.MeanMetric 37 | - _target_: torchmetrics.MeanMetric 38 | - _target_: torchmetrics.MeanMetric 39 | input_keys: 40 | - mean_success 41 | - mean_grasp 42 | - mean_align 43 | output_keys: 44 | - val/mean_success 45 | - val/mean_grasp 46 | - val/mean_align 47 | best_val_metrics: 48 | metrics: 49 | - _target_: torchmetrics.MaxMetric 50 | - _target_: torchmetrics.MaxMetric 51 | - _target_: torchmetrics.MaxMetric 52 | input_keys: 53 | - val/mean_success 54 | - val/mean_grasp 55 | - val/mean_align 56 | output_keys: 57 | - val/mean_success 58 | - val/mean_grasp 59 | - val/mean_align -------------------------------------------------------------------------------- /configs/data/rlbench_act_pcd_dataset.yaml: -------------------------------------------------------------------------------- 1 | _target_: src.data.base_datamodule.BaseDataModule 2 | 3 | batch_size_train: 32 4 | batch_size_val: 1 5 | batch_size_test: 1 6 | num_workers: 16 7 | pin_memory: true 8 | train: 9 | _target_: src.data.components.rlbench.rlbench_single_task_act.RLBenchSingleTaskACTPCDDataset 10 | root: data/rlbench/processed/train/ 11 | task_names: 12 | - ${rlbench_task} 13 | chunk_size: 100 14 | max_episodes_per_task: 100 15 | cameras: 16 | - front 17 | action_dim: 11 18 | collision: true 19 | rot_type: 6d 20 | include_depth: null 21 | transform_pcd: 22 | - _target_: src.data.components.transformpcd.GridSamplePCD 23 | grid_size: 0.005 24 | hash_type: fnv 25 | mode: train 26 | return_grid_coord: true 27 | return_displacement: false 28 | keys: 29 | - coord 30 | - color 31 | - _target_: src.data.components.transformpcd.NormalizeColorPCD 32 | - _target_: src.data.components.transformpcd.ShufflePointPCD 33 | - _target_: src.data.components.transformpcd.ToTensorPCD 34 | - _target_: src.data.components.transformpcd.CollectPCD 35 | keys: 36 | - coord 37 | - grid_coord 38 | feat_keys: 39 | - color 40 | - coord 41 | loop: 1 42 | 43 | val: 44 | _target_: src.data.components.rlbench.rlbench_single_task_act.RLBenchSingleTaskACTPCDDataset 45 | root: data/rlbench/processed/val/ 46 | task_names: ${data.train.task_names} 47 | chunk_size: ${data.train.chunk_size} 48 | max_episodes_per_task: 25 49 | cameras: ${data.train.cameras} 50 | action_dim: ${data.train.action_dim} 51 | collision: ${data.train.collision} 52 | rot_type: ${data.train.rot_type} 53 | include_depth: ${data.train.include_depth} 54 | transform_pcd: ${data.train.transform_pcd} 55 | loop: 1 56 | -------------------------------------------------------------------------------- /configs/test_rlbench_act.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | 3 | # specify here default configuration 4 | # order of defaults determines the order in which configs override each other 5 | defaults: 6 | - _self_ 7 | - data: 8 | - model: 9 | - callbacks: default 10 | - logger: tensorboard # set logger here or use command line (e.g. `python train.py logger=tensorboard`) 11 | - trainer: ddp 12 | - paths: default 13 | - extras: default 14 | - hydra: default 15 | - exp_rlbench_act_policy: 16 | 17 | # optional local config for machine/user specific settings 18 | # it's optional since it doesn't need to exist and is excluded from version control 19 | - optional local: default 20 | 21 | # debugging config (enable through command line, e.g. `python train.py debug=default) 22 | - debug: 23 | 24 | # task name, determines output directory path 25 | task_name: "test" 26 | 27 | # tags to help you identify your experiments 28 | # you can overwrite this in experiment configs 29 | # overwrite from command line with `python train.py tags="[first_tag, second_tag]"` 30 | tags: ["test"] 31 | 32 | # simply provide checkpoint path to resume training 33 | ckpt_path: 34 | 35 | # seed for random number generators in pytorch, numpy and python.random 36 | seed: 1000 37 | 38 | # rlbench test options 39 | live_demos: false 40 | data_root: data/rlbench/raw/test 41 | headless: false 42 | episodes_num: 25 43 | max_steps: 200 44 | temporal_agg: 45 | apply: true 46 | action_dim: ${data.train.action_dim} 47 | chunk_size: ${data.train.chunk_size} 48 | k: 0.01 49 | 50 | result_path: logs/rlbench_act/${rlbench_task}/${task_name} 51 | result_file: test_result.txt 52 | result_name: ${ckpt_path} 53 | 54 | offline: false 55 | 56 | camera_view_test: 57 | apply: false 58 | rot_axis: y 59 | rot_angle: 0 60 | transl: [0, 0, 0] 61 | -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/point.frag: -------------------------------------------------------------------------------- 1 | #version 450 2 | #extension GL_ARB_separate_shader_objects : enable 3 | #extension GL_ARB_shading_language_420pack : enable 4 | 5 | layout (constant_id = 0) const int NUM_DIRECTIONAL_LIGHTS = 3; 6 | layout (constant_id = 1) const int NUM_POINT_LIGHTS = 10; 7 | layout (constant_id = 2) const int NUM_SPOT_LIGHTS = 10; 8 | layout (constant_id = 3) const float RESOLUTION_SCALE = 1.0; 9 | 10 | layout(location = 0) in vec4 inColor; 11 | layout(location = 1) in flat vec4 inCenterRadius; 12 | 13 | layout(location = 0) out vec4 outPointColor; 14 | layout(location = 1) out vec4 outPointDepthLinear; 15 | 16 | layout(set = 0, binding = 0) uniform CameraBuffer { 17 | mat4 viewMatrix; 18 | mat4 projectionMatrix; 19 | mat4 viewMatrixInverse; 20 | mat4 projectionMatrixInverse; 21 | mat4 prevViewMatrix; 22 | mat4 prevViewMatrixInverse; 23 | float width; 24 | float height; 25 | } cameraBuffer; 26 | 27 | void main() { 28 | vec3 center = inCenterRadius.xyz; 29 | float radius = inCenterRadius.w; 30 | 31 | vec2 res = vec2(cameraBuffer.width, cameraBuffer.height) * RESOLUTION_SCALE; 32 | vec2 pixelNdc = gl_FragCoord.xy / res * 2.0 - 1.0; 33 | 34 | vec4 dir4 = cameraBuffer.projectionMatrixInverse * vec4(pixelNdc, 1.0, 1.0); 35 | vec3 dir = normalize(dir4.xyz / dir4.w); 36 | 37 | float a = dot(dir, center); 38 | float N = a * a - (dot(center, center) - radius * radius); 39 | 40 | if (N < 0) { 41 | discard; 42 | } 43 | 44 | float d = dot(dir, center) - sqrt(N); 45 | 46 | if (d < 0) { 47 | discard; 48 | } 49 | 50 | vec4 csPosition = vec4(dir * d, 1.0); 51 | vec4 ndc = cameraBuffer.projectionMatrix * csPosition; 52 | ndc /= ndc.w; 53 | gl_FragDepth = ndc.z; 54 | 55 | outPointColor = inColor; 56 | outPointDepthLinear.x = csPosition.z; 57 | } 58 | -------------------------------------------------------------------------------- /configs/data/rlbench_diffusion_policy_pcd_dataset.yaml: -------------------------------------------------------------------------------- 1 | _target_: src.data.base_datamodule.BaseDataModule 2 | 3 | batch_size_train: 32 4 | batch_size_val: 1 5 | batch_size_test: 1 6 | num_workers: 16 7 | pin_memory: true 8 | train: 9 | _target_: src.data.components.rlbench.rlbench_single_task_dp.RLBenchSingleTaskDiffusionPolicyPCDDataset 10 | root: data/rlbench/processed/train/ 11 | task_names: 12 | - ${rlbench_task} 13 | chunk_size: 16 14 | max_episodes_per_task: 100 15 | cameras: 16 | - front 17 | action_dim: 11 18 | collision: true 19 | rot_type: 6d 20 | include_depth: null 21 | transform_pcd: 22 | - _target_: src.data.components.transformpcd.GridSamplePCD 23 | grid_size: 0.005 24 | hash_type: fnv 25 | mode: train 26 | return_grid_coord: true 27 | return_displacement: false 28 | keys: 29 | - coord 30 | - color 31 | - _target_: src.data.components.transformpcd.NormalizeColorPCD 32 | - _target_: src.data.components.transformpcd.ShufflePointPCD 33 | - _target_: src.data.components.transformpcd.ToTensorPCD 34 | - _target_: src.data.components.transformpcd.CollectPCD 35 | keys: 36 | - coord 37 | - grid_coord 38 | feat_keys: 39 | - color 40 | - coord 41 | loop: 1 42 | 43 | val: 44 | _target_: src.data.components.rlbench.rlbench_single_task_dp.RLBenchSingleTaskDiffusionPolicyPCDDataset 45 | root: data/rlbench/processed/val/ 46 | task_names: ${data.train.task_names} 47 | chunk_size: ${data.train.chunk_size} 48 | max_episodes_per_task: 25 49 | cameras: ${data.train.cameras} 50 | action_dim: ${data.train.action_dim} 51 | collision: ${data.train.collision} 52 | rot_type: ${data.train.rot_type} 53 | include_depth: ${data.train.include_depth} 54 | transform_pcd: ${data.train.transform_pcd} 55 | loop: 1 56 | -------------------------------------------------------------------------------- /src/utils/metrics.py: -------------------------------------------------------------------------------- 1 | from collections.abc import Sequence 2 | 3 | import torch 4 | import torch.nn as nn 5 | 6 | 7 | class Metrics(nn.Module): 8 | def __init__(self, metrics, input_keys, output_keys): 9 | super().__init__() 10 | self.metrics = metrics 11 | 12 | self.input_keys = input_keys 13 | if not isinstance(self.input_keys, Sequence): 14 | self.input_keys = [self.input_keys] * len(self.metrics) 15 | else: 16 | assert len(self.input_keys) == len(self.metrics) 17 | 18 | self.output_keys = output_keys 19 | if not isinstance(self.output_keys, Sequence): 20 | self.output_keys = [self.output_keys] * len(self.metrics) 21 | else: 22 | assert len(self.output_keys) == len(self.metrics) 23 | 24 | self.metrics = nn.ModuleList(self.metrics) 25 | 26 | def reset(self): 27 | for c in self.metrics: 28 | c.reset() 29 | 30 | def compute(self): 31 | metrics = dict() 32 | for c, out_k in zip(self.metrics, self.output_keys): 33 | metrics[out_k] = c.compute() 34 | return metrics 35 | 36 | def metrics_dict(self): 37 | metrics = dict() 38 | for c, out_k in zip(self.metrics, self.output_keys): 39 | metrics[out_k] = c 40 | return metrics 41 | 42 | @torch.inference_mode() 43 | def forward(self, data_dict): 44 | metrics = dict() 45 | for c, in_k, out_k in zip(self.metrics, self.input_keys, self.output_keys): 46 | if isinstance(in_k, Sequence): 47 | try: 48 | metrics[out_k] = c(*[data_dict[k_] for k_ in in_k]) 49 | except KeyError: 50 | metrics[out_k] = c(data_dict[in_k]) 51 | else: 52 | metrics[out_k] = c(data_dict[in_k]) 53 | return metrics 54 | -------------------------------------------------------------------------------- /src/utils/io.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | import os 4 | import pickle 5 | from io import BytesIO 6 | 7 | import cv2 8 | import h5py 9 | import imageio.v3 as iio 10 | import numpy as np 11 | import torch 12 | 13 | 14 | def load_bytes(filename): 15 | return open(filename, "rb").read() 16 | 17 | 18 | def load_text(filename): 19 | return open(filename, "r").read().splitlines() 20 | 21 | 22 | def load_numpy_text(filename): 23 | return np.loadtxt(filename) 24 | 25 | 26 | def listdir(dir): 27 | return os.listdir(dir) 28 | 29 | 30 | def load_json(filename): 31 | return json.load(open(filename)) 32 | 33 | 34 | def dump_json(filename, data): 35 | json.dump(data, open(filename, "w")) 36 | 37 | 38 | def load_numpy_pickle(filename): 39 | array = np.load(filename, allow_pickle=True) 40 | 41 | try: 42 | return array.item() 43 | except: 44 | return array 45 | 46 | 47 | def load_numpy(filename): 48 | return np.load(filename) 49 | 50 | 51 | def load_pickle(filename): 52 | return pickle.load(open(filename, "rb")) 53 | 54 | 55 | def load_image(filename): 56 | image = cv2.imread(filename, cv2.IMREAD_UNCHANGED) 57 | if ".jpg" in filename: 58 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 59 | return image 60 | 61 | 62 | def exists(path): 63 | return os.path.exists(path) 64 | 65 | 66 | def isdir(path): 67 | return os.path.isdir(path) 68 | 69 | 70 | def imwrite(img, path): 71 | iio.imwrite(path, img) 72 | 73 | 74 | def load_pth( 75 | path, 76 | map_location="cpu", 77 | ): 78 | return torch.load(path, map_location=map_location) 79 | 80 | 81 | def load_h5_data(data): 82 | out = dict() 83 | for k in data.keys(): 84 | if isinstance(data[k], h5py.Dataset): 85 | out[k] = data[k][:] 86 | else: 87 | out[k] = load_h5_data(data[k]) 88 | return out 89 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_spunet_pcd_wo_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_diffusion_policy_pcd_dataset 4 | - override /model: maniskill2_diffusion_policy_model 5 | - override /paths: default 6 | 7 | task_name: scratch_spunet_pcd_wo_rgb 8 | 9 | data: 10 | train: 11 | transform_pcd: 12 | - _targt_: src.data.components.transformpcd.GridSamplePCD 13 | grid_size: 0.005 14 | hash_type: fnv 15 | mode: train 16 | return_grid_coord: true 17 | return_displacement: false 18 | keys: 19 | - coord 20 | - color 21 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 22 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 23 | - _targt_: src.data.components.transformpcd.ToTensorPCD 24 | - _targt_: src.data.components.transformpcd.CollectPCD 25 | keys: 26 | - coord 27 | - grid_coord 28 | feat_keys: 29 | - coord 30 | 31 | model: 32 | obs_mode: pointcloud 33 | policy: 34 | shape_meta: 35 | obs: 36 | pcds: 37 | shape: [3] 38 | type: pcd 39 | qpos: 40 | shape: [7] 41 | type: low_dim 42 | obs_encoder: 43 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 44 | shape_meta: ${model.policy.shape_meta} 45 | pcd_model: 46 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 47 | in_channels: 3 48 | num_classes: 0 49 | pretrained_path: null 50 | share_pcd_model: true 51 | n_obs_step: ${model.policy.n_obs_steps} 52 | pcd_nsample: 16 53 | pcd_npoints: 2048 54 | use_mask: false 55 | bg_ratio: 0.0 56 | pcd_hidden_dim: 96 57 | projector_layers: 2 58 | projector_channels: [96, 128, 128] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_spunet_pcd_wo_xyz.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_diffusion_policy_pcd_dataset 4 | - override /model: maniskill2_diffusion_policy_model 5 | - override /paths: default 6 | 7 | task_name: scratch_spunet_pcd_wo_xyz 8 | 9 | data: 10 | train: 11 | transform_pcd: 12 | - _targt_: src.data.components.transformpcd.GridSamplePCD 13 | grid_size: 0.005 14 | hash_type: fnv 15 | mode: train 16 | return_grid_coord: true 17 | return_displacement: false 18 | keys: 19 | - coord 20 | - color 21 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 22 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 23 | - _targt_: src.data.components.transformpcd.ToTensorPCD 24 | - _targt_: src.data.components.transformpcd.CollectPCD 25 | keys: 26 | - coord 27 | - grid_coord 28 | feat_keys: 29 | - color 30 | 31 | model: 32 | obs_mode: pointcloud 33 | policy: 34 | shape_meta: 35 | obs: 36 | pcds: 37 | shape: [3] 38 | type: pcd 39 | qpos: 40 | shape: [7] 41 | type: low_dim 42 | obs_encoder: 43 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 44 | shape_meta: ${model.policy.shape_meta} 45 | pcd_model: 46 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 47 | in_channels: 3 48 | num_classes: 0 49 | pretrained_path: null 50 | share_pcd_model: true 51 | n_obs_step: ${model.policy.n_obs_steps} 52 | pcd_nsample: 16 53 | pcd_npoints: 2048 54 | use_mask: false 55 | bg_ratio: 0.0 56 | pcd_hidden_dim: 96 57 | projector_layers: 2 58 | projector_channels: [96, 128, 128] -------------------------------------------------------------------------------- /src/envs/custom_maniskill2/mpm/shader/point/gbuffer.vert: -------------------------------------------------------------------------------- 1 | #version 450 2 | #extension GL_ARB_separate_shader_objects : enable 3 | #extension GL_ARB_shading_language_420pack : enable 4 | 5 | layout(set = 0, binding = 0) uniform CameraBuffer { 6 | mat4 viewMatrix; 7 | mat4 projectionMatrix; 8 | mat4 viewMatrixInverse; 9 | mat4 projectionMatrixInverse; 10 | mat4 prevViewMatrix; 11 | mat4 prevViewMatrixInverse; 12 | float width; 13 | float height; 14 | } cameraBuffer; 15 | 16 | layout(set = 1, binding = 0) uniform ObjectBuffer { 17 | mat4 modelMatrix; 18 | mat4 prevModelMatrix; 19 | uvec4 segmentation; 20 | float transparency; 21 | int shadeFlat; 22 | } objectBuffer; 23 | 24 | layout(location = 0) in vec3 position; 25 | layout(location = 1) in vec3 normal; 26 | layout(location = 2) in vec2 uv; 27 | layout(location = 3) in vec3 tangent; 28 | layout(location = 4) in vec3 bitangent; 29 | 30 | layout(location = 0) out vec4 outPosition; 31 | layout(location = 1) out vec4 outPrevPosition; 32 | layout(location = 2) out vec2 outUV; 33 | layout(location = 3) out flat uvec4 outSegmentation; 34 | layout(location = 4) out vec3 objectCoord; 35 | layout(location = 5) out mat3 outTbn; 36 | 37 | void main() { 38 | outSegmentation = objectBuffer.segmentation; 39 | mat4 modelView = cameraBuffer.viewMatrix * objectBuffer.modelMatrix; 40 | mat3 normalMatrix = mat3(transpose(inverse(modelView))); 41 | objectCoord = position; 42 | outPosition = modelView * vec4(position, 1); 43 | outPrevPosition = cameraBuffer.prevViewMatrix * objectBuffer.prevModelMatrix * vec4(position, 1); 44 | outUV = uv; 45 | gl_Position = cameraBuffer.projectionMatrix * outPosition; 46 | vec3 outTangent = normalize(normalMatrix * tangent); 47 | vec3 outBitangent = normalize(normalMatrix * bitangent); 48 | vec3 outNormal = normalize(normalMatrix * normal); 49 | outTbn = mat3(outTangent, outBitangent, outNormal); 50 | } 51 | -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_pointnet_pcd_wo_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_diffusion_policy_pcd_dataset 4 | - override /model: maniskill2_diffusion_policy_model 5 | - override /paths: default 6 | 7 | task_name: scratch_pointnet_pcd_wo_rgb 8 | 9 | data: 10 | train: 11 | transform_pcd: 12 | - _targt_: src.data.components.transformpcd.GridSamplePCD 13 | grid_size: 0.005 14 | hash_type: fnv 15 | mode: train 16 | return_grid_coord: true 17 | return_displacement: false 18 | keys: 19 | - coord 20 | - color 21 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 22 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 23 | - _targt_: src.data.components.transformpcd.ToTensorPCD 24 | - _targt_: src.data.components.transformpcd.CollectPCD 25 | keys: 26 | - coord 27 | - grid_coord 28 | feat_keys: 29 | - coord 30 | 31 | model: 32 | obs_mode: pointcloud 33 | policy: 34 | shape_meta: 35 | obs: 36 | pcds: 37 | shape: [3] 38 | type: pcd 39 | qpos: 40 | shape: [7] 41 | type: low_dim 42 | obs_encoder: 43 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 44 | shape_meta: ${model.policy.shape_meta} 45 | pcd_model: 46 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 47 | in_channels: 3 48 | num_classes: 96 49 | pretrained_path: null 50 | share_pcd_model: true 51 | n_obs_step: ${model.policy.n_obs_steps} 52 | pcd_nsample: 16 53 | pcd_npoints: 2048 54 | use_mask: false 55 | bg_ratio: 0.0 56 | pcd_hidden_dim: 96 57 | projector_layers: 1 58 | projector_channels: [96, 128, 128] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_pointnet_pcd_wo_xyz.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_diffusion_policy_pcd_dataset 4 | - override /model: maniskill2_diffusion_policy_model 5 | - override /paths: default 6 | 7 | task_name: scratch_pointnet_pcd_wo_xyz 8 | 9 | data: 10 | train: 11 | transform_pcd: 12 | - _targt_: src.data.components.transformpcd.GridSamplePCD 13 | grid_size: 0.005 14 | hash_type: fnv 15 | mode: train 16 | return_grid_coord: true 17 | return_displacement: false 18 | keys: 19 | - coord 20 | - color 21 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 22 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 23 | - _targt_: src.data.components.transformpcd.ToTensorPCD 24 | - _targt_: src.data.components.transformpcd.CollectPCD 25 | keys: 26 | - coord 27 | - grid_coord 28 | feat_keys: 29 | - color 30 | 31 | model: 32 | obs_mode: pointcloud 33 | policy: 34 | shape_meta: 35 | obs: 36 | pcds: 37 | shape: [3] 38 | type: pcd 39 | qpos: 40 | shape: [7] 41 | type: low_dim 42 | obs_encoder: 43 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 44 | shape_meta: ${model.policy.shape_meta} 45 | pcd_model: 46 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 47 | in_channels: 3 48 | num_classes: 96 49 | pretrained_path: null 50 | share_pcd_model: true 51 | n_obs_step: ${model.policy.n_obs_steps} 52 | pcd_nsample: 16 53 | pcd_npoints: 2048 54 | use_mask: false 55 | bg_ratio: 0.0 56 | pcd_hidden_dim: 96 57 | projector_layers: 1 58 | projector_channels: [96, 128, 128] -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_pointnet_pcd_presample_wo_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_pointnet_pcd_presample_wo_rgb 7 | 8 | data: 9 | train: 10 | transform_pcd: 11 | - _targt_: src.data.components.transformpcd.GridSamplePCD 12 | grid_size: 0.005 13 | hash_type: fnv 14 | mode: train 15 | return_grid_coord: true 16 | return_displacement: false 17 | keys: 18 | - coord 19 | - color 20 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 21 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 22 | - _targt_: src.data.components.transformpcd.ToTensorPCD 23 | - _targt_: src.data.components.transformpcd.CollectPCD 24 | keys: 25 | - coord 26 | - grid_coord 27 | feat_keys: 28 | - coord 29 | 30 | model: 31 | obs_mode: pointcloud 32 | policy: 33 | shape_meta: 34 | obs: 35 | pcds: 36 | shape: [3] 37 | type: pcd 38 | qpos: 39 | shape: [7] 40 | type: low_dim 41 | obs_encoder: 42 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 43 | shape_meta: ${model.policy.shape_meta} 44 | pcd_model: 45 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 46 | in_channels: 3 47 | num_classes: 96 48 | pretrained_path: null 49 | share_pcd_model: true 50 | n_obs_step: ${model.policy.n_obs_steps} 51 | pcd_nsample: 16 52 | pcd_npoints: 2048 53 | use_mask: false 54 | bg_ratio: 0.0 55 | pcd_hidden_dim: 96 56 | projector_layers: 1 57 | projector_channels: [96, 128, 128] 58 | pre_sample: true 59 | in_channel: 3 -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_pointnet_pcd_presample_wo_xyz.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /model: maniskill2_diffusion_policy_model 4 | - override /paths: default 5 | 6 | task_name: scratch_pointnet_pcd_presample_wo_xyz 7 | 8 | data: 9 | train: 10 | transform_pcd: 11 | - _targt_: src.data.components.transformpcd.GridSamplePCD 12 | grid_size: 0.005 13 | hash_type: fnv 14 | mode: train 15 | return_grid_coord: true 16 | return_displacement: false 17 | keys: 18 | - coord 19 | - color 20 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 21 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 22 | - _targt_: src.data.components.transformpcd.ToTensorPCD 23 | - _targt_: src.data.components.transformpcd.CollectPCD 24 | keys: 25 | - coord 26 | - grid_coord 27 | feat_keys: 28 | - color 29 | 30 | model: 31 | obs_mode: pointcloud 32 | policy: 33 | shape_meta: 34 | obs: 35 | pcds: 36 | shape: [3] 37 | type: pcd 38 | qpos: 39 | shape: [7] 40 | type: low_dim 41 | obs_encoder: 42 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 43 | shape_meta: ${model.policy.shape_meta} 44 | pcd_model: 45 | _target_: src.models.components.pcd_encoder.pointnet.PointNet 46 | in_channels: 3 47 | num_classes: 96 48 | pretrained_path: null 49 | share_pcd_model: true 50 | n_obs_step: ${model.policy.n_obs_steps} 51 | pcd_nsample: 16 52 | pcd_npoints: 2048 53 | use_mask: false 54 | bg_ratio: 0.0 55 | pcd_hidden_dim: 96 56 | projector_layers: 1 57 | projector_channels: [96, 128, 128] 58 | pre_sample: true 59 | in_channel: 3 -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_spunet_pcd_presample_wo_rgb.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_diffusion_policy_pcd_dataset 4 | - override /model: maniskill2_diffusion_policy_model 5 | - override /paths: default 6 | 7 | task_name: scratch_spunet_pcd_presample_wo_rgb 8 | 9 | data: 10 | train: 11 | transform_pcd: 12 | - _targt_: src.data.components.transformpcd.GridSamplePCD 13 | grid_size: 0.005 14 | hash_type: fnv 15 | mode: train 16 | return_grid_coord: true 17 | return_displacement: false 18 | keys: 19 | - coord 20 | - color 21 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 22 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 23 | - _targt_: src.data.components.transformpcd.ToTensorPCD 24 | - _targt_: src.data.components.transformpcd.CollectPCD 25 | keys: 26 | - coord 27 | - grid_coord 28 | feat_keys: 29 | - coord 30 | 31 | model: 32 | obs_mode: pointcloud 33 | policy: 34 | shape_meta: 35 | obs: 36 | pcds: 37 | shape: [3] 38 | type: pcd 39 | qpos: 40 | shape: [7] 41 | type: low_dim 42 | obs_encoder: 43 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 44 | shape_meta: ${model.policy.shape_meta} 45 | pcd_model: 46 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 47 | in_channels: 3 48 | num_classes: 0 49 | pretrained_path: null 50 | share_pcd_model: true 51 | n_obs_step: ${model.policy.n_obs_steps} 52 | pcd_nsample: 16 53 | pcd_npoints: 2048 54 | use_mask: false 55 | bg_ratio: 0.0 56 | pcd_hidden_dim: 96 57 | projector_layers: 2 58 | projector_channels: [96, 128, 128] 59 | pre_sample: true 60 | in_channel: 3 -------------------------------------------------------------------------------- /configs/exp_maniskill2_diffusion_policy/maniskill2_model/scratch_spunet_pcd_presample_wo_xyz.yaml: -------------------------------------------------------------------------------- 1 | # @package _global_ 2 | defaults: 3 | - override /data: maniskill2_diffusion_policy_pcd_dataset 4 | - override /model: maniskill2_diffusion_policy_model 5 | - override /paths: default 6 | 7 | task_name: scratch_spunet_pcd_presample_wo_xyz 8 | 9 | data: 10 | train: 11 | transform_pcd: 12 | - _targt_: src.data.components.transformpcd.GridSamplePCD 13 | grid_size: 0.005 14 | hash_type: fnv 15 | mode: train 16 | return_grid_coord: true 17 | return_displacement: false 18 | keys: 19 | - coord 20 | - color 21 | - _targt_: src.data.components.transformpcd.NormalizeColorPCD 22 | - _targt_: src.data.components.transformpcd.ShufflePointPCD 23 | - _targt_: src.data.components.transformpcd.ToTensorPCD 24 | - _targt_: src.data.components.transformpcd.CollectPCD 25 | keys: 26 | - coord 27 | - grid_coord 28 | feat_keys: 29 | - color 30 | 31 | model: 32 | obs_mode: pointcloud 33 | policy: 34 | shape_meta: 35 | obs: 36 | pcds: 37 | shape: [3] 38 | type: pcd 39 | qpos: 40 | shape: [7] 41 | type: low_dim 42 | obs_encoder: 43 | _target_: src.models.components.diffusion_policy.vision.pcd_obs_encoder.PCDObsEncoder 44 | shape_meta: ${model.policy.shape_meta} 45 | pcd_model: 46 | _target_: src.models.components.pcd_encoder.spunet.SpUNet 47 | in_channels: 3 48 | num_classes: 0 49 | pretrained_path: null 50 | share_pcd_model: true 51 | n_obs_step: ${model.policy.n_obs_steps} 52 | pcd_nsample: 16 53 | pcd_npoints: 2048 54 | use_mask: false 55 | bg_ratio: 0.0 56 | pcd_hidden_dim: 96 57 | projector_layers: 2 58 | projector_channels: [96, 128, 128] 59 | pre_sample: true 60 | in_channel: 3 -------------------------------------------------------------------------------- /libs/pointops/functions/aggregation.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pointops._C import aggregation_backward_cuda, aggregation_forward_cuda 3 | from torch.autograd import Function 4 | 5 | 6 | class Aggregation(Function): 7 | @staticmethod 8 | def forward(ctx, input, position, weight, idx): 9 | """ 10 | input: input: (n, c), position: (n, nsample, c), weight : (n, nsample, c'), idx: (n, nsample) 11 | output: (n, c) 12 | """ 13 | assert ( 14 | input.is_contiguous() 15 | and position.is_contiguous() 16 | and weight.is_contiguous() 17 | ) 18 | n, nsample, c = position.shape 19 | w_c = weight.shape[-1] 20 | output = torch.cuda.FloatTensor(n, c).zero_() 21 | aggregation_forward_cuda( 22 | n, nsample, c, w_c, input, position, weight, idx, output 23 | ) 24 | ctx.save_for_backward(input, position, weight, idx) 25 | return output 26 | 27 | @staticmethod 28 | def backward(ctx, grad_output): 29 | """ 30 | input: grad_out: (n, c) 31 | output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight : (n, nsample, c') 32 | """ 33 | input, position, weight, idx = ctx.saved_tensors 34 | n, nsample, c = position.shape 35 | w_c = weight.shape[-1] 36 | grad_input = torch.cuda.FloatTensor(n, c).zero_() 37 | grad_position = torch.cuda.FloatTensor(n, nsample, c).zero_() 38 | grad_weight = torch.cuda.FloatTensor(n, nsample, w_c).zero_() 39 | aggregation_backward_cuda( 40 | n, 41 | nsample, 42 | c, 43 | w_c, 44 | input, 45 | position, 46 | weight, 47 | idx, 48 | grad_output, 49 | grad_input, 50 | grad_position, 51 | grad_weight, 52 | ) 53 | return grad_input, grad_position, grad_weight, None 54 | 55 | 56 | aggregation = Aggregation.apply 57 | -------------------------------------------------------------------------------- /src/utils/logging_utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from typing import Any, Dict 4 | 5 | from lightning_utilities.core.rank_zero import rank_zero_only 6 | from omegaconf import OmegaConf 7 | 8 | from src.utils import pylogger 9 | 10 | log = pylogger.RankedLogger(__name__, rank_zero_only=True) 11 | 12 | 13 | @rank_zero_only 14 | def log_hyperparameters(object_dict: dict[str, Any]) -> None: 15 | """Controls which config parts are saved by Lightning loggers. 16 | 17 | Additionally saves: 18 | - Number of model parameters 19 | 20 | :param object_dict: A dictionary containing the following objects: 21 | - `"cfg"`: A DictConfig object containing the main config. 22 | - `"model"`: The Lightning model. 23 | - `"trainer"`: The Lightning trainer. 24 | """ 25 | hparams = {} 26 | 27 | cfg = OmegaConf.to_container(object_dict["cfg"]) 28 | model = object_dict["model"] 29 | trainer = object_dict["trainer"] 30 | 31 | if not trainer.logger: 32 | log.warning("Logger not found! Skipping hyperparameter logging...") 33 | return 34 | 35 | hparams["model"] = cfg["model"] 36 | 37 | # save number of model parameters 38 | hparams["model/params/total"] = sum(p.numel() for p in model.parameters()) 39 | hparams["model/params/trainable"] = sum( 40 | p.numel() for p in model.parameters() if p.requires_grad 41 | ) 42 | hparams["model/params/non_trainable"] = sum( 43 | p.numel() for p in model.parameters() if not p.requires_grad 44 | ) 45 | 46 | hparams["data"] = cfg["data"] 47 | hparams["trainer"] = cfg["trainer"] 48 | 49 | hparams["callbacks"] = cfg.get("callbacks") 50 | hparams["extras"] = cfg.get("extras") 51 | 52 | hparams["task_name"] = cfg.get("task_name") 53 | hparams["tags"] = cfg.get("tags") 54 | hparams["ckpt_path"] = cfg.get("ckpt_path") 55 | hparams["seed"] = cfg.get("seed") 56 | 57 | # send hparams to all loggers 58 | for logger in trainer.loggers: 59 | logger.log_hyperparams(hparams) 60 | -------------------------------------------------------------------------------- /scripts/run_maniskill2_camera_view.sh: -------------------------------------------------------------------------------- 1 | ckpt_path=$1 2 | task=$2 3 | model=$3 4 | seed=$4 5 | 6 | task_name=${task%-v0} 7 | if [ "$task_name" = "PegInsertionSide" ];then 8 | model_env_id=${task_name}-3steps-MultiView 9 | else 10 | model_env_id=${task_name}-MultiView 11 | fi 12 | 13 | echo "Starting validation for task: ${task} with model: ${model} and checkpoint path: ${ckpt_path}" 14 | echo "Model environment ID: ${model_env_id}" 15 | echo "Seed: ${seed}" 16 | 17 | if [[ "$model" == *"pcd"* ]]; then 18 | echo "Running validation for PCD camera views..." 19 | # Change the camera view for PCD by specifying the `camera_ids` parameter 20 | # Mapping of camera names to IDs: 21 | # 4:left_camera_5 22 | # 5:down_camera_5 23 | # 6:left_camera_10 24 | # 7:down_camera_10 25 | for camera in 4 5 6 7 26 | do 27 | echo "Validating with camera ID: ${camera}" 28 | python src/validate.py exp_maniskill2_act_policy=base \ 29 | exp_maniskill2_act_policy/maniskill2_model@maniskill2_task=${task} \ 30 | exp_maniskill2_act_policy/maniskill2_model@maniskill2_model=${model} \ 31 | ckpt_path=${ckpt_path} \ 32 | model.env_id=${model_env_id} \ 33 | data.train.camera_ids=${camera} seed=${seed} 34 | done 35 | else 36 | echo "Running validation for RGB or RGBD camera views..." 37 | # Change the camera view for RGB and RGBD by specifying the `camera_names` parameter 38 | for camera in left_camera_5 left_camera_10 down_camera_5 down_camera_10 39 | do 40 | echo "Validating with camera view: ${camera}" 41 | python src/validate.py exp_maniskill2_act_policy=base \ 42 | exp_maniskill2_act_policy/maniskill2_model@maniskill2_task=${task} \ 43 | exp_maniskill2_act_policy/maniskill2_model@maniskill2_model=${model} \ 44 | ckpt_path=${ckpt_path} \ 45 | model.env_id=${model_env_id} \ 46 | data.train.camera_names=${camera} seed=${seed} 47 | done 48 | fi 49 | 50 | echo "Validation completed." 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | -------------------------------------------------------------------------------- /src/utils/instantiators.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from typing import List 4 | 5 | import hydra 6 | from lightning import Callback 7 | from lightning.pytorch.loggers import Logger 8 | from omegaconf import DictConfig 9 | 10 | from src.utils import pylogger 11 | 12 | log = pylogger.RankedLogger(__name__, rank_zero_only=True) 13 | 14 | 15 | def instantiate_callbacks(callbacks_cfg: DictConfig) -> list[Callback]: 16 | """Instantiates callbacks from config. 17 | 18 | :param callbacks_cfg: A DictConfig object containing callback configurations. 19 | :return: A list of instantiated callbacks. 20 | """ 21 | callbacks: list[Callback] = [] 22 | 23 | if not callbacks_cfg: 24 | log.warning("No callback configs found! Skipping..") 25 | return callbacks 26 | 27 | if not isinstance(callbacks_cfg, DictConfig): 28 | raise TypeError("Callbacks config must be a DictConfig!") 29 | 30 | for _, cb_conf in callbacks_cfg.items(): 31 | if isinstance(cb_conf, DictConfig) and "_target_" in cb_conf: 32 | log.info(f"Instantiating callback <{cb_conf._target_}>") 33 | callbacks.append(hydra.utils.instantiate(cb_conf)) 34 | 35 | return callbacks 36 | 37 | 38 | def instantiate_loggers(logger_cfg: DictConfig) -> list[Logger]: 39 | """Instantiates loggers from config. 40 | 41 | :param logger_cfg: A DictConfig object containing logger configurations. 42 | :return: A list of instantiated loggers. 43 | """ 44 | logger: list[Logger] = [] 45 | 46 | if not logger_cfg: 47 | log.warning("No logger configs found! Skipping...") 48 | return logger 49 | 50 | if not isinstance(logger_cfg, DictConfig): 51 | raise TypeError("Logger config must be a DictConfig!") 52 | 53 | for _, lg_conf in logger_cfg.items(): 54 | if isinstance(lg_conf, DictConfig) and "_target_" in lg_conf: 55 | log.info(f"Instantiating logger <{lg_conf._target_}>") 56 | logger.append(hydra.utils.instantiate(lg_conf)) 57 | 58 | return logger 59 | -------------------------------------------------------------------------------- /libs/pointops/src/grouping/grouping_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "grouping_cuda_kernel.h" 3 | 4 | 5 | __global__ void grouping_forward_cuda_kernel(int m, int nsample, int c, const float *__restrict__ input, const int *__restrict__ idx, float *__restrict__ output) { 6 | // input: input: (n, c), idx: (m, nsample), output: (m, nsample, c) 7 | int index = blockIdx.x * blockDim.x + threadIdx.x; 8 | if (index >= m * nsample * c) return; 9 | const int c_idx = index % c; 10 | const int nsample_idx = (index / c) % nsample; 11 | const int m_idx = index / nsample / c; 12 | const int input_idx = idx[m_idx * nsample + nsample_idx] * c + c_idx; 13 | output[index] = input[input_idx]; 14 | } 15 | 16 | __global__ void grouping_backward_cuda_kernel(int m, int nsample, int c, const float *__restrict__ grad_output, const int *__restrict__ idx, float *__restrict__ grad_input) { 17 | // input: grad_output: (m, nsample, c), idx: (m, nsample), output: grad_input: (n, c) 18 | int index = blockIdx.x * blockDim.x + threadIdx.x; 19 | if (index >= m * nsample * c) return; 20 | const int c_idx = index % c; 21 | const int nsample_idx = (index / c) % nsample; 22 | const int m_idx = index / nsample / c; 23 | const int input_idx = idx[m_idx * nsample + nsample_idx] * c + c_idx; 24 | atomicAdd(grad_input + input_idx, grad_output[index]); 25 | } 26 | 27 | void grouping_forward_cuda_launcher(int m, int nsample, int c, const float *input, const int *idx, float *output) { 28 | // input: input: (n, c), idx: (m, nsample), output: (m, nsample, c) 29 | dim3 blocks(DIVUP(m * nsample * c, THREADS_PER_BLOCK)); 30 | dim3 threads(THREADS_PER_BLOCK); 31 | grouping_forward_cuda_kernel<<>>(m, nsample, c, input, idx, output); 32 | } 33 | 34 | void grouping_backward_cuda_launcher(int m, int nsample, int c, const float *grad_output, const int *idx, float *grad_input) 35 | { 36 | // input: grad_output: (m, nsample, c), idx: (m, nsample), output: grad_input: (n, c) 37 | dim3 blocks(DIVUP(m * nsample * c, THREADS_PER_BLOCK)); 38 | dim3 threads(THREADS_PER_BLOCK); 39 | grouping_backward_cuda_kernel<<>>(m, nsample, c, grad_output, idx, grad_input); 40 | } 41 | --------------------------------------------------------------------------------