├── .gitignore ├── LICENSE ├── README.md ├── configs ├── brl │ └── bc │ │ ├── pointnet.py │ │ ├── pointnet_soft_body.py │ │ ├── rgbd.py │ │ └── rgbd_soft_body.py └── mfrl │ ├── dapg │ ├── maniskill2_pn.py │ ├── maniskill2_pn_frameminers.py │ └── maniskill2_rgbd.py │ ├── gail │ ├── maniskill2_pn.py │ └── maniskill2_pn_frameminers.py │ ├── ppo │ ├── maniskill2_pn.py │ ├── maniskill2_pn_frameminers.py │ ├── maniskill2_rgbd.py │ └── maniskill2_sparseconv.py │ └── sac │ ├── maniskill2_pn.py │ ├── maniskill2_rgbd.py │ └── maniskill2_state.py ├── maniskill2_learn ├── __init__.py ├── apis │ ├── __init__.py │ ├── render_traj.py │ ├── run_rl.py │ └── train_rl.py ├── env │ ├── __init__.py │ ├── action_space_utils.py │ ├── builder.py │ ├── env_utils.py │ ├── evaluation.py │ ├── observation_process.py │ ├── replay_buffer.py │ ├── rollout.py │ ├── sampling_strategy.py │ ├── vec_env.py │ └── wrappers.py ├── methods │ ├── __init__.py │ ├── brl │ │ ├── __init__.py │ │ └── bc.py │ ├── builder.py │ └── mfrl │ │ ├── __init__.py │ │ ├── gail.py │ │ ├── ppo.py │ │ └── sac.py ├── networks │ ├── __init__.py │ ├── applications │ │ ├── __init__.py │ │ └── actor_critic.py │ ├── backbones │ │ ├── __init__.py │ │ ├── mlp.py │ │ ├── pointnet.py │ │ ├── resnet.py │ │ ├── rl_cnn.py │ │ ├── sp_resnet.py │ │ ├── transformer.py │ │ └── visuomotor.py │ ├── builder.py │ ├── modules │ │ ├── __init__.py │ │ ├── activation.py │ │ ├── attention.py │ │ ├── block_utils.py │ │ ├── cnn_modules │ │ │ ├── __init__.py │ │ │ └── resnet_utils.py │ │ ├── conv.py │ │ ├── linear.py │ │ ├── norm.py │ │ ├── padding.py │ │ ├── pct_modules │ │ │ ├── __init__.py │ │ │ └── pointnet_util.py │ │ ├── plugin.py │ │ ├── pooling.py │ │ ├── spconv_modules │ │ │ ├── __init__.py │ │ │ ├── resnet_utils.py │ │ │ └── spconv_utils.py │ │ ├── test_norm.py │ │ └── weight_init.py │ ├── ops │ │ ├── __init__.py │ │ └── ops_3d │ │ │ ├── __init__.py │ │ │ ├── ball_query │ │ │ ├── __init__.py │ │ │ ├── ball_query.py │ │ │ └── src │ │ │ │ ├── ball_query.cpp │ │ │ │ └── ball_query_cuda.cu │ │ │ └── pcd_process │ │ │ ├── __init__.py │ │ │ ├── pcd_process.cpp │ │ │ ├── pcd_process.py │ │ │ └── pcd_process_cuda.cu │ ├── regression_heads │ │ ├── __init__.py │ │ ├── deterministic.py │ │ ├── gaussian.py │ │ └── regression_base.py │ └── utils.py ├── schedulers │ ├── __init__.py │ ├── custom_scheduler.py │ └── lr_scheduler.py ├── utils │ ├── __init__.py │ ├── data │ │ ├── __init__.py │ │ ├── array_ops.py │ │ ├── compression.py │ │ ├── converter.py │ │ ├── dict_array.py │ │ ├── dict_utils.py │ │ ├── filtering.py │ │ ├── misc.py │ │ ├── seq_utils.py │ │ ├── string_utils.py │ │ ├── type_utils.py │ │ └── wrappers.py │ ├── file │ │ ├── __init__.py │ │ ├── cache_utils.py │ │ ├── file_client.py │ │ ├── hash_utils.py │ │ ├── hdf5_utils.py │ │ ├── lmdb_utils.py │ │ ├── pandas_utils.py │ │ ├── record_utils.py │ │ ├── serialization │ │ │ ├── __init__.py │ │ │ ├── handlers │ │ │ │ ├── __init__.py │ │ │ │ ├── base.py │ │ │ │ ├── csv_handler.py │ │ │ │ ├── json_handler.py │ │ │ │ ├── pickle_handler.py │ │ │ │ ├── txt_handler.py │ │ │ │ └── yaml_handler.py │ │ │ ├── io.py │ │ │ └── utils.py │ │ └── zip_utils.py │ ├── image │ │ ├── __init__.py │ │ ├── colorspace.py │ │ ├── geometric.py │ │ ├── io.py │ │ ├── misc.py │ │ ├── photometric.py │ │ └── video_utils.py │ ├── lib3d │ │ ├── __init__.py │ │ ├── mani_skill2_contrib.py │ │ ├── o3d_utils.py │ │ ├── trimesh_utils.py │ │ └── utils.py │ ├── math │ │ ├── __init__.py │ │ ├── counting.py │ │ ├── running_stats.py │ │ ├── split_array.py │ │ └── trunc_normal.py │ ├── meta │ │ ├── __init__.py │ │ ├── collect_env.py │ │ ├── config.py │ │ ├── env_var.py │ │ ├── logger.py │ │ ├── magic_utils.py │ │ ├── module_utils.py │ │ ├── network.py │ │ ├── parallel_runner.py │ │ ├── path_utils.py │ │ ├── process_utils.py │ │ ├── progressbar.py │ │ ├── random_utils.py │ │ ├── registry.py │ │ ├── timer.py │ │ └── version_utils.py │ ├── torch │ │ ├── __init__.py │ │ ├── checkpoint_utils.py │ │ ├── cuda_utils.py │ │ ├── distributed_utils.py │ │ ├── distributions.py │ │ ├── flops_counter.py │ │ ├── freezer.py │ │ ├── logger │ │ │ ├── __init__.py │ │ │ ├── tensorboard_logger.py │ │ │ ├── tensorboard_utils.py │ │ │ └── wandb_logger.py │ │ ├── misc.py │ │ ├── module_utils.py │ │ ├── ops.py │ │ ├── optimizer_utils.py │ │ └── running_stats.py │ └── visualization │ │ ├── __init__.py │ │ └── o3d_utils.py └── version.py ├── requirements.txt ├── scripts ├── example_demo_conversion │ ├── general_rigid_body_multi_object_envs.sh │ ├── general_rigid_body_single_object_envs.sh │ ├── general_soft_body_envs.sh │ ├── maniskill1.sh │ ├── pick_place_multi_object_envs.sh │ └── pick_place_single_object_envs.sh └── example_training │ ├── pretrained_model │ ├── bc_soft_body_pointcloud.sh │ ├── bc_soft_body_rgbd.sh │ ├── dapg_pickcube_pointcloud.sh │ ├── dapg_pickcube_rgbd.sh │ ├── dapg_picksingleycb_pointcloud.sh │ ├── dapg_picksingleycb_rgbd.sh │ ├── dapg_stackcube_pointcloud.sh │ └── dapg_stackcube_rgbd.sh │ └── scratch_pointcloud_template │ ├── run_bc.sh │ ├── run_dapg.sh │ ├── run_gail.sh │ ├── run_ppo.sh │ └── run_sac.sh ├── setup.py └── tools ├── convert_state.py ├── merge_h5.py ├── merge_trajectory.py ├── permute_rgbd_agent_weights.py └── shuffle_demo.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | **/__pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Core dump file 10 | core 11 | **/core 12 | 13 | # Distribution / packaging 14 | .Python 15 | build/ 16 | test/ 17 | develop-eggs/ 18 | dist/ 19 | downloads/ 20 | eggs/ 21 | .eggs/ 22 | lib/ 23 | lib64/ 24 | parts/ 25 | sdist/ 26 | var/ 27 | wheels/ 28 | *.egg-info/ 29 | .installed.cfg 30 | *.egg 31 | MANIFEST 32 | 33 | # PyInstaller 34 | # Usually these files are written by a python script from a template 35 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 36 | *.manifest 37 | *.spec 38 | 39 | # Installer logs 40 | pip-log.txt 41 | pip-delete-this-directory.txt 42 | 43 | # Unit test / coverage reports 44 | htmlcov/ 45 | .tox/ 46 | .coverage 47 | .coverage.* 48 | .cache 49 | nosetests.xml 50 | coverage.xml 51 | *.cover 52 | .hypothesis/ 53 | .pytest_cache/ 54 | 55 | # Translations 56 | *.mo 57 | *.pot 58 | 59 | # Django stuff: 60 | *.log 61 | local_settings.py 62 | db.sqlite3 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | **/.ipynb_checkpoints 79 | *.ipynb 80 | *.ini 81 | 82 | # pyenv 83 | .python-version 84 | 85 | # celery beat schedule file 86 | celerybeat-schedule 87 | 88 | # SageMath parsed files 89 | *.sage.py 90 | 91 | # Environments 92 | .env 93 | .venv 94 | venv/ 95 | env.bak/ 96 | venv.bak/ 97 | 98 | # Spyder project settings 99 | .spyderproject 100 | .spyproject 101 | 102 | # Rope project settingspartnet-mobility-dataset 103 | .mypy_cache/ 104 | 105 | .vscode 106 | .idea 107 | 108 | # custom 109 | *.pkl 110 | *.pkl.json 111 | *.log.json 112 | *.zip 113 | 114 | # Ignore running logs 115 | data_analyze/ 116 | tmp/ 117 | work_dirs/ 118 | 119 | # Models, images and data 120 | *.pth 121 | *.ckpt 122 | *.h5 123 | *.tmp 124 | *.png 125 | *.svg 126 | *.jpg 127 | *.mp4 128 | *.prof 129 | *.cache 130 | 131 | .history 132 | **/.DS_Store 133 | workspace.code-workspace 134 | **/build/ 135 | wandb/ 136 | 137 | tools/torchsparse/ 138 | partnet-mobility-dataset 139 | 140 | tools/docker_file/corl/ 141 | -------------------------------------------------------------------------------- /configs/brl/bc/pointnet.py: -------------------------------------------------------------------------------- 1 | agent_cfg = dict( 2 | type="BC", 3 | batch_size=256, 4 | actor_cfg=dict( 5 | type="ContinuousActor", 6 | head_cfg=dict( 7 | type="TanhHead", 8 | noise_std=1e-5, 9 | ), 10 | nn_cfg=dict( 11 | type="Visuomotor", 12 | visual_nn_cfg=dict(type="PointNet", feat_dim="pcd_all_channel", mlp_spec=[64, 128, 512], feature_transform=[]), 13 | mlp_cfg=dict( 14 | type="LinearMLP", 15 | norm_cfg=None, 16 | mlp_spec=["512 + agent_shape", 256, "action_shape"], 17 | inactivated_output=True, 18 | zero_init_output=True, 19 | ), 20 | ), 21 | optim_cfg=dict(type="Adam", lr=3e-4), 22 | ), 23 | ) 24 | 25 | env_cfg = dict( 26 | type="gym", 27 | env_name="PickCube-v0", 28 | unwrapped=False, 29 | ) 30 | 31 | 32 | replay_cfg = dict( 33 | type="ReplayMemory", 34 | capacity=-1, 35 | num_samples=-1, 36 | keys=["obs", "actions", "dones", "episode_dones"], 37 | buffer_filenames=[ 38 | "SOME_DEMO_FILE", 39 | ], 40 | ) 41 | 42 | train_cfg = dict( 43 | on_policy=False, 44 | total_steps=50000, 45 | warm_steps=0, 46 | n_steps=0, 47 | n_updates=500, 48 | n_eval=50000, 49 | n_checkpoint=50000, 50 | ) 51 | 52 | eval_cfg = dict( 53 | type="Evaluation", 54 | num=10, 55 | num_procs=1, 56 | use_hidden_state=False, 57 | save_traj=False, 58 | save_video=True, 59 | use_log=False, 60 | ) 61 | -------------------------------------------------------------------------------- /configs/brl/bc/pointnet_soft_body.py: -------------------------------------------------------------------------------- 1 | agent_cfg = dict( 2 | type="BC", 3 | batch_size=256, 4 | actor_cfg=dict( 5 | type="ContinuousActor", 6 | head_cfg=dict( 7 | type="GaussianHead", 8 | init_log_std=-0.5, 9 | clip_return=True, 10 | predict_std=False 11 | ), 12 | nn_cfg=dict( 13 | type="Visuomotor", 14 | visual_nn_cfg=dict(type="PointNet", feat_dim="pcd_all_channel", mlp_spec=[64, 128, 512], feature_transform=[]), 15 | mlp_cfg=dict( 16 | type="LinearMLP", 17 | norm_cfg=None, 18 | mlp_spec=["512 + agent_shape", 256, 256, "action_shape"], 19 | inactivated_output=True, 20 | zero_init_output=True, 21 | ), 22 | ), 23 | optim_cfg=dict(type="Adam", lr=3e-4), 24 | ), 25 | ) 26 | 27 | env_cfg = dict( 28 | type="gym", 29 | env_name="Fill-v0", 30 | unwrapped=False, 31 | ) 32 | 33 | 34 | replay_cfg = dict( 35 | type="ReplayMemory", 36 | capacity=-1, 37 | num_samples=-1, 38 | keys=["obs", "actions", "dones", "episode_dones"], 39 | buffer_filenames=[ 40 | "SOME_DEMO_FILE", 41 | ], 42 | ) 43 | 44 | train_cfg = dict( 45 | on_policy=False, 46 | total_steps=50000, 47 | warm_steps=0, 48 | n_steps=0, 49 | n_updates=500, 50 | n_eval=50000, 51 | n_checkpoint=50000, 52 | ) 53 | 54 | eval_cfg = dict( 55 | type="Evaluation", 56 | num=10, 57 | num_procs=1, 58 | use_hidden_state=False, 59 | save_traj=False, 60 | save_video=True, 61 | use_log=False, 62 | ) 63 | -------------------------------------------------------------------------------- /configs/brl/bc/rgbd.py: -------------------------------------------------------------------------------- 1 | agent_cfg = dict( 2 | type="BC", 3 | batch_size=256, 4 | actor_cfg=dict( 5 | type="ContinuousActor", 6 | head_cfg=dict( 7 | type="TanhHead", 8 | noise_std=1e-5, 9 | ), 10 | nn_cfg=dict( 11 | type="Visuomotor", 12 | visual_nn_cfg=dict(type="IMPALA", in_channel="image_channels", image_size="image_size", out_feature_size=384), 13 | mlp_cfg=dict( 14 | type="LinearMLP", norm_cfg=None, mlp_spec=["384 + agent_shape", 256, 128, "action_shape"], bias=True, inactivated_output=True 15 | ), 16 | ), 17 | optim_cfg=dict(type="Adam", lr=3e-4), 18 | ), 19 | ) 20 | 21 | env_cfg = dict( 22 | type="gym", 23 | env_name="PickCube-v0", 24 | unwrapped=False, 25 | ) 26 | 27 | 28 | replay_cfg = dict( 29 | type="ReplayMemory", 30 | capacity=-1, 31 | num_samples=-1, 32 | keys=["obs", "actions", "dones", "episode_dones"], 33 | buffer_filenames=[ 34 | "SOME_DEMO_FILE", 35 | ], 36 | ) 37 | 38 | train_cfg = dict( 39 | on_policy=False, 40 | total_steps=50000, 41 | warm_steps=0, 42 | n_steps=0, 43 | n_updates=500, 44 | n_eval=50000, 45 | n_checkpoint=50000, 46 | ) 47 | 48 | eval_cfg = dict( 49 | type="Evaluation", 50 | num=10, 51 | num_procs=1, 52 | use_hidden_state=False, 53 | save_traj=False, 54 | save_video=True, 55 | use_log=False, 56 | ) 57 | -------------------------------------------------------------------------------- /configs/brl/bc/rgbd_soft_body.py: -------------------------------------------------------------------------------- 1 | agent_cfg = dict( 2 | type="BC", 3 | batch_size=256, 4 | actor_cfg=dict( 5 | type="ContinuousActor", 6 | head_cfg=dict( 7 | type="GaussianHead", 8 | init_log_std=-0.5, 9 | clip_return=True, 10 | predict_std=False 11 | ), 12 | nn_cfg=dict( 13 | type="Visuomotor", 14 | visual_nn_cfg=dict(type="IMPALA", in_channel="image_channels", image_size="image_size", out_feature_size=512), 15 | mlp_cfg=dict( 16 | type="LinearMLP", norm_cfg=None, mlp_spec=["512 + agent_shape", 256, 128, "action_shape"], bias=True, inactivated_output=True 17 | ), 18 | ), 19 | optim_cfg=dict(type="Adam", lr=3e-4), 20 | ), 21 | ) 22 | 23 | env_cfg = dict( 24 | type="gym", 25 | env_name="Fill-v0", 26 | unwrapped=False, 27 | ) 28 | 29 | 30 | replay_cfg = dict( 31 | type="ReplayMemory", 32 | capacity=-1, 33 | num_samples=-1, 34 | keys=["obs", "actions", "dones", "episode_dones"], 35 | buffer_filenames=[ 36 | "SOME_DEMO_FILE", 37 | ], 38 | ) 39 | 40 | train_cfg = dict( 41 | on_policy=False, 42 | total_steps=50000, 43 | warm_steps=0, 44 | n_steps=0, 45 | n_updates=500, 46 | n_eval=50000, 47 | n_checkpoint=50000, 48 | ) 49 | 50 | eval_cfg = dict( 51 | type="Evaluation", 52 | num=10, 53 | num_procs=1, 54 | use_hidden_state=False, 55 | save_traj=False, 56 | save_video=True, 57 | use_log=False, 58 | ) 59 | -------------------------------------------------------------------------------- /configs/mfrl/dapg/maniskill2_pn.py: -------------------------------------------------------------------------------- 1 | 2 | agent_cfg = dict( 3 | type="PPO", 4 | gamma=0.95, 5 | lmbda=0.95, 6 | critic_coeff=0.5, 7 | entropy_coeff=0, 8 | critic_clip=False, 9 | obs_norm=False, 10 | rew_norm=True, 11 | adv_norm=True, 12 | recompute_value=True, 13 | num_epoch=2, 14 | critic_warmup_epoch=4, 15 | batch_size=330, 16 | detach_actor_feature=False, 17 | max_grad_norm=0.5, 18 | eps_clip=0.2, 19 | max_kl=0.2, 20 | dual_clip=None, 21 | shared_backbone=True, 22 | ignore_dones=True, 23 | dapg_lambda=0.1, 24 | dapg_damping=0.995, 25 | actor_cfg=dict( 26 | type="ContinuousActor", 27 | head_cfg=dict( 28 | type="GaussianHead", 29 | init_log_std=-1, 30 | clip_return=True, 31 | predict_std=False, 32 | ), 33 | nn_cfg=dict( 34 | type="Visuomotor", 35 | visual_nn_cfg=dict(type="PointNet", feat_dim="pcd_all_channel", mlp_spec=[64, 128, 512], feature_transform=[]), 36 | mlp_cfg=dict( 37 | type="LinearMLP", 38 | norm_cfg=None, 39 | mlp_spec=["512 + agent_shape", 256, 256, "action_shape"], 40 | inactivated_output=True, 41 | zero_init_output=True, 42 | ), 43 | ), 44 | optim_cfg=dict(type="Adam", lr=3e-4, param_cfg={"(.*?)visual_nn(.*?)": None}), 45 | ), 46 | critic_cfg=dict( 47 | type="ContinuousCritic", 48 | nn_cfg=dict( 49 | type="Visuomotor", 50 | visual_nn_cfg=None, 51 | mlp_cfg=dict( 52 | type="LinearMLP", norm_cfg=None, mlp_spec=["512 + agent_shape", 256, 256, 1], inactivated_output=True, zero_init_output=True 53 | ), 54 | ), 55 | optim_cfg=dict(type="Adam", lr=3e-4), 56 | ), 57 | demo_replay_cfg=dict( 58 | type="ReplayMemory", 59 | capacity=-1, 60 | num_samples=-1, 61 | keys=["obs", "actions", "dones", "episode_dones"], 62 | buffer_filenames=[ 63 | "PATH_TO_DEMO.h5", 64 | ], 65 | ), 66 | ) 67 | 68 | 69 | train_cfg = dict( 70 | on_policy=True, 71 | total_steps=int(5e6), 72 | warm_steps=0, 73 | n_steps=int(2e4), 74 | n_updates=1, 75 | n_eval=int(5e6), 76 | n_checkpoint=int(1e6), 77 | ep_stats_cfg=dict( 78 | info_keys_mode=dict( 79 | success=[True, "max", "mean"], 80 | ) 81 | ), 82 | ) 83 | 84 | 85 | env_cfg = dict( 86 | type="gym", 87 | env_name="PickCube-v0", 88 | obs_mode='pointcloud', 89 | ignore_dones=True, 90 | ) 91 | 92 | 93 | rollout_cfg = dict( 94 | type="Rollout", 95 | num_procs=5, 96 | with_info=True, 97 | multi_thread=False, 98 | ) 99 | 100 | 101 | replay_cfg = dict( 102 | type="ReplayMemory", 103 | capacity=int(2e4), 104 | sampling_cfg=dict(type="OneStepTransition", with_replacement=False), 105 | ) 106 | 107 | 108 | eval_cfg = dict( 109 | type="Evaluation", 110 | num_procs=1, 111 | num=10, 112 | use_hidden_state=False, 113 | save_traj=False, 114 | save_video=True, 115 | log_every_step=False, 116 | env_cfg=dict(ignore_dones=False), 117 | ) 118 | -------------------------------------------------------------------------------- /configs/mfrl/dapg/maniskill2_pn_frameminers.py: -------------------------------------------------------------------------------- 1 | 2 | agent_cfg = dict( 3 | type="PPO", 4 | gamma=0.95, 5 | lmbda=0.95, 6 | critic_coeff=0.5, 7 | entropy_coeff=0, 8 | critic_clip=False, 9 | obs_norm=False, 10 | rew_norm=True, 11 | adv_norm=True, 12 | recompute_value=True, 13 | num_epoch=2, 14 | critic_warmup_epoch=4, 15 | batch_size=330, 16 | detach_actor_feature=False, 17 | max_grad_norm=0.5, 18 | eps_clip=0.2, 19 | max_kl=0.2, 20 | dual_clip=None, 21 | shared_backbone=True, 22 | ignore_dones=True, 23 | dapg_lambda=0.1, 24 | dapg_damping=0.995, 25 | actor_cfg=dict( 26 | type="ContinuousActor", 27 | head_cfg=dict( 28 | type="GaussianHead", 29 | init_log_std=-1, 30 | clip_return=True, 31 | predict_std=False, 32 | ), 33 | nn_cfg=dict( 34 | type="FrameMiners", 35 | visual_nn_cfg=dict(type="PointNet", feat_dim="pcd_all_channel", mlp_spec=[64, 128, 300], feature_transform=[]), 36 | num_frames=2, 37 | vis_feat_dim=300, 38 | action_dim='action_shape', 39 | robot_state_dim='agent_shape', 40 | ), 41 | optim_cfg=dict(type="Adam", lr=3e-4), 42 | ), 43 | critic_cfg=dict( 44 | type="ContinuousCritic", 45 | nn_cfg=dict( 46 | type="FrameMiners", 47 | visual_nn_cfg=None, 48 | num_frames=2, 49 | vis_feat_dim=300, 50 | action_dim='action_shape', 51 | robot_state_dim='agent_shape', 52 | is_critic=True, 53 | ), 54 | optim_cfg=dict(type="Adam", lr=3e-4), 55 | ), 56 | demo_replay_cfg=dict( 57 | type="ReplayMemory", 58 | capacity=-1, 59 | num_samples=-1, 60 | keys=["obs", "actions", "dones", "episode_dones"], 61 | buffer_filenames=[ 62 | "PATH_TO_DEMO.h5", 63 | ], 64 | ), 65 | ) 66 | 67 | 68 | train_cfg = dict( 69 | on_policy=True, 70 | total_steps=int(5e6), 71 | warm_steps=0, 72 | n_steps=int(2.5e4), # FrameMiners might need a bit larger n_steps (and replay_cfg.capacity) than regular PPO (not tested) 73 | n_updates=1, 74 | n_eval=int(5e6), 75 | n_checkpoint=int(1e6), 76 | ep_stats_cfg=dict( 77 | info_keys_mode=dict( 78 | success=[True, "max", "mean"], 79 | ) 80 | ), 81 | ) 82 | 83 | 84 | env_cfg = dict( 85 | type="gym", 86 | env_name="PickCube-v0", 87 | obs_mode='pointcloud', 88 | ignore_dones=True, 89 | ) 90 | 91 | 92 | rollout_cfg = dict( 93 | type="Rollout", 94 | num_procs=5, 95 | with_info=True, 96 | multi_thread=False, 97 | ) 98 | 99 | 100 | replay_cfg = dict( 101 | type="ReplayMemory", 102 | capacity=int(2.5e4), 103 | sampling_cfg=dict(type="OneStepTransition", with_replacement=False), 104 | ) 105 | 106 | 107 | eval_cfg = dict( 108 | type="Evaluation", 109 | num_procs=1, 110 | num=10, 111 | use_hidden_state=False, 112 | save_traj=False, 113 | save_video=True, 114 | log_every_step=False, 115 | env_cfg=dict(ignore_dones=False), 116 | ) 117 | -------------------------------------------------------------------------------- /configs/mfrl/dapg/maniskill2_rgbd.py: -------------------------------------------------------------------------------- 1 | 2 | agent_cfg = dict( 3 | type="PPO", 4 | gamma=0.95, 5 | lmbda=0.95, 6 | critic_coeff=0.5, 7 | entropy_coeff=0, 8 | critic_clip=False, 9 | obs_norm=False, 10 | rew_norm=True, 11 | adv_norm=True, 12 | recompute_value=True, 13 | detach_actor_feature=True, 14 | critic_warmup_epoch=4, 15 | num_epoch=2, 16 | batch_size=400, 17 | max_grad_norm=0.5, 18 | eps_clip=0.2, 19 | max_kl=0.1, 20 | dual_clip=None, 21 | shared_backbone=True, 22 | ignore_dones=True, 23 | dapg_lambda=0.1, 24 | dapg_damping=0.995, 25 | actor_cfg=dict( 26 | type="ContinuousActor", 27 | head_cfg=dict( 28 | type="GaussianHead", 29 | init_log_std=-0.5, 30 | clip_return=True, 31 | predict_std=False, 32 | ), 33 | nn_cfg=dict( 34 | type="Visuomotor", 35 | visual_nn_cfg=dict(type="IMPALA", in_channel="image_channels", image_size="image_size", out_feature_size=384), 36 | mlp_cfg=dict( 37 | type="LinearMLP", norm_cfg=None, mlp_spec=["384 + agent_shape", 256, 128, "action_shape"], bias=True, inactivated_output=True 38 | ), 39 | ), 40 | optim_cfg=dict(type="Adam", lr=3e-4, param_cfg={"(.*?)visual_nn(.*?)": None}), 41 | # For our PPO implementation, we need to prevent actor_optim from updating visual_nn as long as shared_backbone=True. 42 | # This is because we use (actor_loss + critic_loss).backward(), then apply actor_optim.step() and critic_optim.step() immediately afterwards. 43 | # Thus we need to prevent the visual backbone from being updated twice. 44 | ), 45 | critic_cfg=dict( 46 | type="ContinuousCritic", 47 | nn_cfg=dict( 48 | type="Visuomotor", 49 | visual_nn_cfg=None, 50 | mlp_cfg=dict(type="LinearMLP", norm_cfg=None, mlp_spec=["384 + agent_shape", 256, 128, 1], bias=True, inactivated_output=True), 51 | ), 52 | optim_cfg=dict(type="Adam", lr=3e-4), 53 | ), 54 | demo_replay_cfg=dict( 55 | type="ReplayMemory", 56 | capacity=-1, 57 | num_samples=-1, 58 | keys=["obs", "actions", "dones", "episode_dones"], 59 | buffer_filenames=[ 60 | "PATH_TO_DEMO.h5", 61 | ], 62 | ), 63 | ) 64 | 65 | 66 | train_cfg = dict( 67 | on_policy=True, 68 | total_steps=int(5e6), 69 | warm_steps=0, 70 | n_steps=int(2e4), 71 | n_updates=1, 72 | n_eval=int(5e6), 73 | n_checkpoint=int(1e6), 74 | ep_stats_cfg=dict( 75 | info_keys_mode=dict( 76 | success=[True, "max", "mean"], 77 | ) 78 | ), 79 | ) 80 | 81 | 82 | env_cfg = dict( 83 | type="gym", 84 | env_name="PickCube-v0", 85 | obs_mode='pointcloud', 86 | ignore_dones=True, 87 | ) 88 | 89 | 90 | rollout_cfg = dict( 91 | type="Rollout", 92 | num_procs=5, 93 | with_info=True, 94 | multi_thread=False, 95 | ) 96 | 97 | 98 | replay_cfg = dict( 99 | type="ReplayMemory", 100 | capacity=int(2e4), 101 | sampling_cfg=dict(type="OneStepTransition", with_replacement=False), 102 | ) 103 | 104 | 105 | eval_cfg = dict( 106 | type="Evaluation", 107 | num_procs=1, 108 | num=10, 109 | use_hidden_state=False, 110 | save_traj=False, 111 | save_video=True, 112 | log_every_step=False, 113 | env_cfg=dict(ignore_dones=False), 114 | ) 115 | -------------------------------------------------------------------------------- /configs/mfrl/ppo/maniskill2_pn.py: -------------------------------------------------------------------------------- 1 | 2 | agent_cfg = dict( 3 | type="PPO", 4 | gamma=0.95, 5 | lmbda=0.95, 6 | critic_coeff=0.5, 7 | entropy_coeff=0, 8 | critic_clip=False, 9 | obs_norm=False, 10 | rew_norm=True, 11 | adv_norm=True, 12 | recompute_value=True, 13 | num_epoch=2, 14 | critic_warmup_epoch=4, 15 | batch_size=330, 16 | detach_actor_feature=False, 17 | max_grad_norm=0.5, 18 | eps_clip=0.2, 19 | max_kl=0.2, 20 | dual_clip=None, 21 | shared_backbone=True, 22 | ignore_dones=True, 23 | actor_cfg=dict( 24 | type="ContinuousActor", 25 | head_cfg=dict( 26 | type="GaussianHead", 27 | init_log_std=-1, 28 | clip_return=True, 29 | predict_std=False, 30 | ), 31 | nn_cfg=dict( 32 | type="Visuomotor", 33 | visual_nn_cfg=dict(type="PointNet", feat_dim="pcd_all_channel", mlp_spec=[64, 128, 512], feature_transform=[]), 34 | mlp_cfg=dict( 35 | type="LinearMLP", 36 | norm_cfg=None, 37 | mlp_spec=["512 + agent_shape", 256, 256, "action_shape"], 38 | inactivated_output=True, 39 | zero_init_output=True, 40 | ), 41 | ), 42 | optim_cfg=dict(type="Adam", lr=3e-4, param_cfg={"(.*?)visual_nn(.*?)": None}), 43 | # For our PPO implementation, we need to prevent actor_optim from updating visual_nn as long as shared_backbone=True. 44 | # This is because we use (actor_loss + critic_loss).backward(), then apply actor_optim.step() and critic_optim.step() immediately afterwards. 45 | # Thus we need to prevent the visual backbone from being updated twice. 46 | ), 47 | critic_cfg=dict( 48 | type="ContinuousCritic", 49 | nn_cfg=dict( 50 | type="Visuomotor", 51 | visual_nn_cfg=None, 52 | mlp_cfg=dict( 53 | type="LinearMLP", norm_cfg=None, mlp_spec=["512 + agent_shape", 256, 256, 1], inactivated_output=True, zero_init_output=True 54 | ), 55 | ), 56 | optim_cfg=dict(type="Adam", lr=3e-4), 57 | ), 58 | ) 59 | 60 | 61 | train_cfg = dict( 62 | on_policy=True, 63 | total_steps=int(5e6), 64 | warm_steps=0, 65 | n_steps=int(2e4), 66 | n_updates=1, 67 | n_eval=int(5e6), 68 | n_checkpoint=int(1e6), 69 | ep_stats_cfg=dict( 70 | info_keys_mode=dict( 71 | success=[True, "max", "mean"], 72 | ) 73 | ), 74 | ) 75 | 76 | 77 | env_cfg = dict( 78 | type="gym", 79 | env_name="PickCube-v0", 80 | obs_mode='pointcloud', 81 | ignore_dones=True, 82 | ) 83 | 84 | 85 | replay_cfg = dict( 86 | type="ReplayMemory", 87 | capacity=int(2e4), 88 | ) 89 | 90 | rollout_cfg = dict( 91 | type="Rollout", 92 | num_procs=5, 93 | with_info=True, 94 | multi_thread=False, 95 | ) 96 | 97 | eval_cfg = dict( 98 | type="Evaluation", 99 | num_procs=1, 100 | num=10, 101 | use_hidden_state=False, 102 | save_traj=False, 103 | save_video=True, 104 | log_every_step=False, 105 | env_cfg=dict(ignore_dones=False), 106 | ) 107 | -------------------------------------------------------------------------------- /configs/mfrl/ppo/maniskill2_pn_frameminers.py: -------------------------------------------------------------------------------- 1 | 2 | agent_cfg = dict( 3 | type="PPO", 4 | gamma=0.95, 5 | lmbda=0.95, 6 | critic_coeff=0.5, 7 | entropy_coeff=0, 8 | critic_clip=False, 9 | obs_norm=False, 10 | rew_norm=True, 11 | adv_norm=True, 12 | recompute_value=True, 13 | num_epoch=2, 14 | critic_warmup_epoch=4, 15 | batch_size=330, 16 | detach_actor_feature=False, 17 | max_grad_norm=0.5, 18 | eps_clip=0.2, 19 | max_kl=0.2, 20 | dual_clip=None, 21 | shared_backbone=True, 22 | ignore_dones=True, 23 | actor_cfg=dict( 24 | type="ContinuousActor", 25 | head_cfg=dict( 26 | type="GaussianHead", 27 | init_log_std=-1, 28 | clip_return=True, 29 | predict_std=False, 30 | ), 31 | nn_cfg=dict( 32 | type="FrameMiners", 33 | visual_nn_cfg=dict(type="PointNet", feat_dim="pcd_all_channel", mlp_spec=[64, 128, 300], feature_transform=[]), 34 | num_frames=2, 35 | vis_feat_dim=300, 36 | action_dim='action_shape', 37 | robot_state_dim='agent_shape', 38 | ), 39 | optim_cfg=dict(type="Adam", lr=3e-4), 40 | ), 41 | critic_cfg=dict( 42 | type="ContinuousCritic", 43 | nn_cfg=dict( 44 | type="FrameMiners", 45 | visual_nn_cfg=None, 46 | num_frames=2, 47 | vis_feat_dim=300, 48 | action_dim='action_shape', 49 | robot_state_dim='agent_shape', 50 | is_critic=True, 51 | ), 52 | optim_cfg=dict(type="Adam", lr=3e-4), 53 | ), 54 | ) 55 | 56 | 57 | train_cfg = dict( 58 | on_policy=True, 59 | total_steps=int(5e6), 60 | warm_steps=0, 61 | n_steps=int(2.5e4), # FrameMiners might need a bit larger n_steps (and replay_cfg.capacity) than regular PPO (not tested) 62 | n_updates=1, 63 | n_eval=int(5e6), 64 | n_checkpoint=int(1e6), 65 | ep_stats_cfg=dict( 66 | info_keys_mode=dict( 67 | success=[True, "max", "mean"], 68 | ) 69 | ), 70 | ) 71 | 72 | 73 | env_cfg = dict( 74 | type="gym", 75 | env_name="PickCube-v0", 76 | obs_mode='pointcloud', 77 | ignore_dones=True, 78 | ) 79 | 80 | 81 | replay_cfg = dict( 82 | type="ReplayMemory", 83 | capacity=int(2.5e4), 84 | ) 85 | 86 | rollout_cfg = dict( 87 | type="Rollout", 88 | num_procs=5, 89 | with_info=True, 90 | multi_thread=False, 91 | ) 92 | 93 | eval_cfg = dict( 94 | type="Evaluation", 95 | num_procs=1, 96 | num=10, 97 | use_hidden_state=False, 98 | save_traj=False, 99 | save_video=True, 100 | log_every_step=False, 101 | env_cfg=dict(ignore_dones=False), 102 | ) 103 | -------------------------------------------------------------------------------- /configs/mfrl/ppo/maniskill2_rgbd.py: -------------------------------------------------------------------------------- 1 | 2 | agent_cfg = dict( 3 | type="PPO", 4 | gamma=0.95, 5 | lmbda=0.95, 6 | critic_coeff=0.5, 7 | entropy_coeff=0, 8 | critic_clip=False, 9 | obs_norm=False, 10 | rew_norm=True, 11 | adv_norm=True, 12 | recompute_value=True, 13 | detach_actor_feature=True, 14 | critic_warmup_epoch=4, 15 | num_epoch=2, 16 | batch_size=400, 17 | max_grad_norm=0.5, 18 | eps_clip=0.2, 19 | max_kl=0.1, 20 | dual_clip=None, 21 | shared_backbone=True, 22 | ignore_dones=True, 23 | actor_cfg=dict( 24 | type="ContinuousActor", 25 | head_cfg=dict( 26 | type="GaussianHead", 27 | init_log_std=-0.5, 28 | clip_return=True, 29 | predict_std=False, 30 | ), 31 | nn_cfg=dict( 32 | type="Visuomotor", 33 | visual_nn_cfg=dict(type="IMPALA", in_channel="image_channels", image_size="image_size", out_feature_size=384), 34 | mlp_cfg=dict( 35 | type="LinearMLP", norm_cfg=None, mlp_spec=["384 + agent_shape", 256, 128, "action_shape"], bias=True, inactivated_output=True 36 | ), 37 | ), 38 | optim_cfg=dict(type="Adam", lr=3e-4, param_cfg={"(.*?)visual_nn(.*?)": None}), 39 | # For our PPO implementation, we need to prevent actor_optim from updating visual_nn as long as shared_backbone=True. 40 | # This is because we use (actor_loss + critic_loss).backward(), then apply actor_optim.step() and critic_optim.step() immediately afterwards. 41 | # Thus we need to prevent the visual backbone from being updated twice. 42 | ), 43 | critic_cfg=dict( 44 | type="ContinuousCritic", 45 | nn_cfg=dict( 46 | type="Visuomotor", 47 | visual_nn_cfg=None, 48 | mlp_cfg=dict(type="LinearMLP", norm_cfg=None, mlp_spec=["384 + agent_shape", 256, 128, 1], bias=True, inactivated_output=True), 49 | ), 50 | optim_cfg=dict(type="Adam", lr=3e-4), 51 | ), 52 | ) 53 | 54 | train_cfg = dict( 55 | on_policy=True, 56 | total_steps=int(5e6), 57 | warm_steps=0, 58 | n_steps=int(2e4), 59 | n_updates=1, 60 | n_eval=int(5e6), 61 | n_checkpoint=int(1e6), 62 | ep_stats_cfg=dict( 63 | info_keys_mode=dict( 64 | success=[True, "max", "mean"], 65 | ) 66 | ), 67 | ) 68 | 69 | 70 | env_cfg = dict( 71 | type="gym", 72 | env_name="PickCube-v0", 73 | obs_mode='rgbd', 74 | ignore_dones=True, 75 | ) 76 | 77 | 78 | replay_cfg = dict( 79 | type="ReplayMemory", 80 | capacity=int(2e4), 81 | ) 82 | 83 | rollout_cfg = dict( 84 | type="Rollout", 85 | num_procs=5, 86 | with_info=True, 87 | multi_thread=False, 88 | ) 89 | 90 | eval_cfg = dict( 91 | type="Evaluation", 92 | num_procs=1, 93 | num=10, 94 | use_hidden_state=False, 95 | save_traj=False, 96 | save_video=True, 97 | log_every_step=False, 98 | env_cfg=dict(ignore_dones=False), 99 | ) 100 | -------------------------------------------------------------------------------- /configs/mfrl/ppo/maniskill2_sparseconv.py: -------------------------------------------------------------------------------- 1 | 2 | # torchsparse comes from torchsparse@git+https://github.com/lz1oceani/torchsparse.git , 3 | # which is modified from the original torchsparse with bug fix and additional normalization implementations 4 | agent_cfg = dict( 5 | type="PPO", 6 | gamma=0.95, 7 | lmbda=0.95, 8 | critic_coeff=0.5, 9 | entropy_coeff=0, 10 | critic_clip=False, 11 | obs_norm=False, 12 | rew_norm=True, 13 | adv_norm=True, 14 | recompute_value=True, 15 | num_epoch=2, 16 | critic_warmup_epoch=4, 17 | batch_size=400, 18 | detach_actor_feature=False, 19 | max_grad_norm=0.5, 20 | eps_clip=0.2, 21 | max_kl=0.2, 22 | dual_clip=None, 23 | shared_backbone=True, 24 | ignore_dones=True, 25 | actor_cfg=dict( 26 | type="ContinuousActor", 27 | head_cfg=dict( 28 | type="GaussianHead", 29 | init_log_std=-1, 30 | clip_return=True, 31 | predict_std=False, 32 | ), 33 | nn_cfg=dict( 34 | type="Visuomotor", 35 | visual_nn_cfg=dict(type="SparseResNet10", in_channel="pcd_all_channel", voxel_size=0.05, dropout=0.0), 36 | mlp_cfg=dict( 37 | type="LinearMLP", 38 | norm_cfg=None, 39 | mlp_spec=["512 + agent_shape", 256, 256, "action_shape"], 40 | inactivated_output=True, 41 | zero_init_output=True, 42 | ), 43 | ), 44 | optim_cfg=dict(type="Adam", lr=3e-4, param_cfg={"(.*?)visual_nn(.*?)": None}), 45 | ), 46 | critic_cfg=dict( 47 | type="ContinuousCritic", 48 | nn_cfg=dict( 49 | type="Visuomotor", 50 | visual_nn_cfg=None, 51 | mlp_cfg=dict( 52 | type="LinearMLP", norm_cfg=None, mlp_spec=["512 + agent_shape", 256, 256, 1], inactivated_output=True, zero_init_output=True 53 | ), 54 | ), 55 | optim_cfg=dict(type="Adam", lr=3e-4), 56 | ), 57 | ) 58 | 59 | 60 | train_cfg = dict( 61 | on_policy=True, 62 | total_steps=int(5e6), 63 | warm_steps=0, 64 | n_steps=int(2e4), 65 | n_updates=1, 66 | n_eval=int(5e6), 67 | n_checkpoint=int(1e6), 68 | ep_stats_cfg=dict( 69 | info_keys_mode=dict( 70 | success=[True, "max", "mean"], 71 | ) 72 | ), 73 | ) 74 | 75 | 76 | env_cfg = dict( 77 | type="gym", 78 | env_name="PickCube-v0", 79 | obs_mode='pointcloud', 80 | ignore_dones=True, 81 | ) 82 | 83 | 84 | replay_cfg = dict( 85 | type="ReplayMemory", 86 | capacity=int(2e4), 87 | ) 88 | 89 | rollout_cfg = dict( 90 | type="Rollout", 91 | num_procs=5, 92 | with_info=True, 93 | multi_thread=False, 94 | ) 95 | 96 | eval_cfg = dict( 97 | type="Evaluation", 98 | num_procs=1, 99 | num=10, 100 | use_hidden_state=False, 101 | save_traj=False, 102 | save_video=True, 103 | log_every_step=False, 104 | env_cfg=dict(ignore_dones=False), 105 | ) 106 | -------------------------------------------------------------------------------- /configs/mfrl/sac/maniskill2_pn.py: -------------------------------------------------------------------------------- 1 | 2 | agent_cfg = dict( 3 | type="SAC", 4 | batch_size=256, # Using multiple gpus leads to larger effective batch size, which can be crucial for SAC training 5 | gamma=0.95, 6 | update_coeff=0.005, 7 | alpha=0.2, 8 | target_update_interval=1, 9 | automatic_alpha_tuning=True, 10 | shared_backbone=True, 11 | detach_actor_feature=True, 12 | alpha_optim_cfg=dict(type="Adam", lr=3e-4), 13 | actor_cfg=dict( 14 | type="ContinuousActor", 15 | head_cfg=dict( 16 | type="TanhGaussianHead", 17 | log_std_bound=[-20, 2], 18 | ), 19 | nn_cfg=dict( 20 | type="Visuomotor", 21 | visual_nn_cfg=dict(type="PointNet", feat_dim="pcd_all_channel", mlp_spec=[64, 128, 512], feature_transform=[]), 22 | mlp_cfg=dict( 23 | type="LinearMLP", 24 | norm_cfg=None, 25 | mlp_spec=["512 + agent_shape", 256, 256, "action_shape * 2"], 26 | inactivated_output=True, 27 | zero_init_output=True, 28 | ), 29 | ), 30 | optim_cfg=dict(type="Adam", lr=3e-4, param_cfg={"(.*?)visual_nn(.*?)": None}), 31 | # *Above removes visual_nn from actor optimizer; should only do so if shared_backbone=True and detach_actor_feature=True 32 | # *If either of the config options is False, then param_cfg={} should be removed, i.e. actor should also update the visual backbone. 33 | # In addition, mlp_specs should be modified as well 34 | 35 | # *It is unknown if sharing backbone and detaching feature works well under the 3D setting. It is up to the users to figure this out. 36 | ), 37 | critic_cfg=dict( 38 | type="ContinuousCritic", 39 | num_heads=2, 40 | nn_cfg=dict( 41 | type="Visuomotor", 42 | visual_nn_cfg=None, 43 | mlp_cfg=dict( 44 | type="LinearMLP", norm_cfg=None, mlp_spec=["512 + agent_shape + action_shape", 256, 256, 1], inactivated_output=True, zero_init_output=True 45 | ), 46 | ), 47 | optim_cfg=dict(type="Adam", lr=3e-4), 48 | ), 49 | ) 50 | 51 | 52 | train_cfg = dict( 53 | on_policy=False, 54 | total_steps=20000000, 55 | warm_steps=8000, 56 | n_eval=20000000, 57 | n_checkpoint=2000000, 58 | n_steps=32, 59 | n_updates=4, 60 | ep_stats_cfg=dict( 61 | info_keys_mode=dict( 62 | success=[True, "max", "mean"], 63 | ) 64 | ), 65 | ) 66 | 67 | env_cfg = dict( 68 | type="gym", 69 | env_name="PickCube-v0", 70 | obs_mode='pointcloud', 71 | ignore_dones=True, 72 | ) 73 | 74 | 75 | replay_cfg = dict( 76 | type="ReplayMemory", 77 | capacity=400000, 78 | ) 79 | 80 | rollout_cfg = dict( 81 | type="Rollout", 82 | num_procs=4, 83 | with_info=True, 84 | multi_thread=False, 85 | ) 86 | 87 | eval_cfg = dict( 88 | type="Evaluation", 89 | num_procs=1, 90 | num=10, 91 | use_hidden_state=False, 92 | save_traj=False, 93 | save_video=True, 94 | log_every_step=False, 95 | env_cfg=dict(ignore_dones=False), 96 | ) -------------------------------------------------------------------------------- /configs/mfrl/sac/maniskill2_rgbd.py: -------------------------------------------------------------------------------- 1 | 2 | agent_cfg = dict( 3 | type="SAC", 4 | batch_size=256, # Using multiple gpus leads to larger effective batch size, which can be crucial for SAC training 5 | gamma=0.95, 6 | update_coeff=0.005, 7 | alpha=0.2, 8 | target_update_interval=1, 9 | automatic_alpha_tuning=True, 10 | shared_backbone=True, 11 | detach_actor_feature=True, 12 | alpha_optim_cfg=dict(type="Adam", lr=3e-4), 13 | actor_cfg=dict( 14 | type="ContinuousActor", 15 | head_cfg=dict( 16 | type="TanhGaussianHead", 17 | log_std_bound=[-20, 2], 18 | ), 19 | nn_cfg=dict( 20 | type="Visuomotor", 21 | visual_nn_cfg=dict(type="IMPALA", in_channel="image_channels", image_size="image_size", out_feature_size=384), 22 | mlp_cfg=dict( 23 | type="LinearMLP", norm_cfg=None, mlp_spec=["384 + agent_shape", 256, 128, "action_shape * 2"], bias=True, inactivated_output=True 24 | ), 25 | ), 26 | optim_cfg=dict(type="Adam", lr=3e-4, param_cfg={"(.*?)visual_nn(.*?)": None}), 27 | # *Above removes visual_nn from actor optimizer; should only do so if shared_backbone=True and detach_actor_feature=True 28 | # *If either of the config options is False, then param_cfg={} should be removed, i.e. actor should also update the visual backbone. 29 | # In addition, mlp_specs should be modified as well 30 | ), 31 | critic_cfg=dict( 32 | type="ContinuousCritic", 33 | num_heads=2, 34 | nn_cfg=dict( 35 | type="Visuomotor", 36 | visual_nn_cfg=None, 37 | mlp_cfg=dict(type="LinearMLP", norm_cfg=None, mlp_spec=["384 + agent_shape + action_shape", 256, 128, 1], bias=True, inactivated_output=True), 38 | ), 39 | optim_cfg=dict(type="Adam", lr=3e-4), 40 | ), 41 | ) 42 | 43 | 44 | train_cfg = dict( 45 | on_policy=False, 46 | total_steps=20000000, 47 | warm_steps=8000, 48 | n_eval=20000000, 49 | n_checkpoint=2000000, 50 | n_steps=32, 51 | n_updates=4, 52 | ep_stats_cfg=dict( 53 | info_keys_mode=dict( 54 | success=[True, "max", "mean"], 55 | ) 56 | ), 57 | ) 58 | 59 | env_cfg = dict( 60 | type="gym", 61 | env_name="PickCube-v0", 62 | obs_mode='rgbd', 63 | ignore_dones=True, 64 | ) 65 | 66 | 67 | replay_cfg = dict( 68 | type="ReplayMemory", 69 | capacity=400000, 70 | ) 71 | 72 | rollout_cfg = dict( 73 | type="Rollout", 74 | num_procs=4, 75 | with_info=True, 76 | multi_thread=False, 77 | ) 78 | 79 | eval_cfg = dict( 80 | type="Evaluation", 81 | num_procs=1, 82 | num=10, 83 | use_hidden_state=False, 84 | save_traj=False, 85 | save_video=True, 86 | log_every_step=False, 87 | env_cfg=dict(ignore_dones=False), 88 | ) -------------------------------------------------------------------------------- /configs/mfrl/sac/maniskill2_state.py: -------------------------------------------------------------------------------- 1 | 2 | agent_cfg = dict( 3 | type="SAC", 4 | batch_size=1024, 5 | gamma=0.95, 6 | update_coeff=0.005, 7 | alpha=0.2, 8 | target_update_interval=1, 9 | automatic_alpha_tuning=True, 10 | alpha_optim_cfg=dict(type="Adam", lr=3e-4), 11 | actor_cfg=dict( 12 | type="ContinuousActor", 13 | head_cfg=dict( 14 | type="TanhGaussianHead", 15 | log_std_bound=[-20, 2], 16 | ), 17 | nn_cfg=dict( 18 | type="LinearMLP", 19 | norm_cfg=None, 20 | mlp_spec=["obs_shape", 256, 256, "action_shape * 2"], 21 | bias="auto", 22 | inactivated_output=True, 23 | # zero_init_output=True, 24 | linear_init_cfg=dict( 25 | type="xavier_init", 26 | gain=1, 27 | bias=0, 28 | ), 29 | ), 30 | optim_cfg=dict(type="Adam", lr=3e-4), 31 | ), 32 | critic_cfg=dict( 33 | type="ContinuousCritic", 34 | num_heads=2, 35 | nn_cfg=dict( 36 | type="LinearMLP", 37 | norm_cfg=None, 38 | bias="auto", 39 | mlp_spec=["obs_shape + action_shape", 256, 256, 1], 40 | inactivated_output=True, 41 | # zero_init_output=True, 42 | linear_init_cfg=dict( 43 | type="xavier_init", 44 | gain=1, 45 | bias=0, 46 | ), 47 | ), 48 | optim_cfg=dict(type="Adam", lr=3e-4), 49 | ), 50 | ) 51 | 52 | 53 | train_cfg = dict( 54 | on_policy=False, 55 | total_steps=8000000, 56 | warm_steps=4000, 57 | n_eval=8000000, 58 | n_checkpoint=1000000, 59 | n_steps=8, 60 | n_updates=4, 61 | ep_stats_cfg=dict( 62 | info_keys_mode=dict( 63 | success=[True, "max", "mean"], 64 | ) 65 | ), 66 | ) 67 | 68 | env_cfg = dict( 69 | type="gym", 70 | env_name="PickCube-v0", 71 | obs_mode='state', 72 | ignore_dones=True, 73 | ) 74 | 75 | 76 | replay_cfg = dict( 77 | type="ReplayMemory", 78 | capacity=1000000, 79 | ) 80 | 81 | rollout_cfg = dict( 82 | type="Rollout", 83 | num_procs=8, 84 | with_info=True, 85 | multi_thread=False, 86 | ) 87 | 88 | eval_cfg = dict( 89 | type="Evaluation", 90 | num_procs=1, 91 | num=10, 92 | use_hidden_state=False, 93 | save_traj=False, 94 | save_video=True, 95 | log_every_step=False, 96 | env_cfg=dict(ignore_dones=False), 97 | ) -------------------------------------------------------------------------------- /maniskill2_learn/__init__.py: -------------------------------------------------------------------------------- 1 | from .version import * 2 | -------------------------------------------------------------------------------- /maniskill2_learn/apis/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/haosulab/ManiSkill2-Learn/e9860aff595da3ad34829d62c923be0b50001427/maniskill2_learn/apis/__init__.py -------------------------------------------------------------------------------- /maniskill2_learn/env/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import build_rollout, build_evaluation, build_replay 2 | from .rollout import Rollout 3 | from .replay_buffer import ReplayMemory 4 | from .sampling_strategy import OneStepTransition, TStepTransition 5 | from .evaluation import BatchEvaluation, Evaluation, save_eval_statistics 6 | from .observation_process import pcd_uniform_downsample 7 | from .env_utils import get_env_info, make_gym_env, build_vec_env, import_env, build_env 8 | from .vec_env import VectorEnv 9 | -------------------------------------------------------------------------------- /maniskill2_learn/env/action_space_utils.py: -------------------------------------------------------------------------------- 1 | from gymnasium.spaces import Box, Discrete, Dict, Space 2 | from maniskill2_learn.utils.data import repeat 3 | import numpy as np 4 | 5 | 6 | class StackedDiscrete(Space): 7 | def __init__(self, num_envs, n): 8 | assert n >= 0 9 | self.n = n 10 | self.num_envs = num_envs 11 | super(StackedDiscrete, self).__init__((), np.int64) 12 | 13 | def sample(self): 14 | return self.np_random.randint(self.n, size=[self.num_envs]) 15 | 16 | def contains(self, x): 17 | return np.logical_and(0 <= x, x < self.n) 18 | 19 | def __repr__(self): 20 | return f"StackedDiscrete(n={self.n}, size={self.num_envs})" 21 | 22 | def __eq__(self, other): 23 | return isinstance(other, StackedDiscrete) and self.n == other.n and self.num_envs == other.num_envs 24 | 25 | 26 | def stack_action_space(action_space, num): 27 | if isinstance(action_space, Box): 28 | low, high = action_space.low, action_space.high 29 | return Box(low=repeat(low[None], num, 0), high=repeat(high[None], num, 0)) 30 | elif isinstance(action_space, Discrete): 31 | return StackedDiscrete(num, action_space.n) 32 | else: 33 | print("Unknown action space:", action_space, type(action_space)) 34 | raise NotImplementedError() 35 | 36 | 37 | def unstack_action_space(action_space): 38 | if isinstance(action_space, Box): 39 | low, high = action_space.low[0], action_space.high[0] 40 | return Box(low=low, high=high) 41 | elif isinstance(action_space, StackedDiscrete): 42 | return Discrete(action_space.n) 43 | else: 44 | raise NotImplementedError() 45 | -------------------------------------------------------------------------------- /maniskill2_learn/env/builder.py: -------------------------------------------------------------------------------- 1 | from maniskill2_learn.utils.meta import Registry, build_from_cfg 2 | 3 | 4 | ROLLOUTS = Registry("rollout") 5 | EVALUATIONS = Registry("evaluation") 6 | 7 | REPLAYS = Registry("replay") 8 | SAMPLING = Registry("sampling") 9 | 10 | 11 | def build_rollout(cfg, default_args=None): 12 | # cfg.type = 'Rollout' 13 | # elif cfg.get("type", "Rollout") == 'BatchRollout': 14 | # print("Although we use only one thread, you still want to use BatchRollout!") 15 | return build_from_cfg(cfg, ROLLOUTS, default_args) 16 | 17 | 18 | def build_evaluation(cfg, default_args=None): 19 | if cfg.get("num_procs", 1) > 1 and cfg.type == "Evaluation": 20 | cfg.type = "BatchEvaluation" 21 | elif cfg.get("type", "Evaluation") == "BatchEvaluation": 22 | print("Although we use only one thread, you still want to use BatchEvaluation!") 23 | return build_from_cfg(cfg, EVALUATIONS, default_args) 24 | 25 | 26 | def build_replay(cfg, default_args=None): 27 | return build_from_cfg(cfg, REPLAYS, default_args) 28 | 29 | 30 | def build_sampling(cfg, default_args=None): 31 | return build_from_cfg(cfg, SAMPLING, default_args) 32 | -------------------------------------------------------------------------------- /maniskill2_learn/env/observation_process.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from maniskill2_learn.utils.data import float_to_int, as_dtype, GDict, sample_and_pad, is_np 3 | from maniskill2_learn.utils.meta import get_logger 4 | from maniskill2_learn.version import __version__ 5 | import deprecation 6 | 7 | 8 | def select_mask(obs, key, mask): 9 | if key in obs: 10 | obs[key] = obs[key][mask] 11 | 12 | 13 | def pcd_filter_ground(pcd, eps=1e-3): 14 | return pcd["xyz"][..., 2] > eps 15 | 16 | 17 | def pcd_filter_with_mask(obs, mask, env=None): 18 | assert isinstance(obs, dict), f"{type(obs)}" 19 | for key in ["xyz", "rgb", "seg", "visual_seg", "robot_seg"]: 20 | select_mask(obs, key, mask) 21 | 22 | 23 | def pcd_uniform_downsample(obs, env=None, ground_eps=1e-3, num=1200): 24 | obs_mode = env.obs_mode 25 | assert obs_mode in ["pointcloud"] 26 | 27 | if ground_eps is not None: 28 | pcd_filter_with_mask(obs, pcd_filter_ground(obs, eps=ground_eps), env) 29 | pcd_filter_with_mask(obs, sample_and_pad(obs["xyz"].shape[0], num), env) 30 | return obs -------------------------------------------------------------------------------- /maniskill2_learn/methods/__init__.py: -------------------------------------------------------------------------------- 1 | from .mfrl import * 2 | from .brl import * 3 | -------------------------------------------------------------------------------- /maniskill2_learn/methods/brl/__init__.py: -------------------------------------------------------------------------------- 1 | from .bc import BC -------------------------------------------------------------------------------- /maniskill2_learn/methods/builder.py: -------------------------------------------------------------------------------- 1 | from ..utils.meta import Registry, build_from_cfg 2 | 3 | 4 | MPC = Registry("mpc") # Model predictive control 5 | MFRL = Registry("mfrl") # Model free RL 6 | BRL = Registry("brl") # Offline RL / Batch RL 7 | 8 | 9 | def build_agent(cfg, default_args=None): 10 | for agent_type in [MPC, MFRL, BRL]: 11 | if cfg["type"] in agent_type: 12 | return build_from_cfg(cfg, agent_type, default_args) 13 | return None 14 | -------------------------------------------------------------------------------- /maniskill2_learn/methods/mfrl/__init__.py: -------------------------------------------------------------------------------- 1 | from .sac import SAC 2 | from .ppo import PPO 3 | from .gail import GAIL -------------------------------------------------------------------------------- /maniskill2_learn/networks/__init__.py: -------------------------------------------------------------------------------- 1 | from .modules import * 2 | from .backbones import * 3 | from .applications import * 4 | from .regression_heads import * 5 | 6 | from .builder import build_backbone, build_model, build_reg_head, build_actor_critic 7 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/applications/__init__.py: -------------------------------------------------------------------------------- 1 | from .actor_critic import ContinuousActor, ContinuousCritic, DiscreteActor, DiscreteCritic 2 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/backbones/__init__.py: -------------------------------------------------------------------------------- 1 | from .mlp import LinearMLP, ConvMLP 2 | from .visuomotor import Visuomotor 3 | from .pointnet import PointNet 4 | 5 | from .transformer import TransformerEncoder 6 | from .resnet import ResNet, ResNetV1c, ResNetV1d 7 | from .visuomotor import Visuomotor 8 | from .rl_cnn import IMPALA, NatureCNN 9 | 10 | try: 11 | from .sp_resnet import SparseResNet10, SparseResNet18, SparseResNet34, SparseResNet50, SparseResNet101 12 | except ImportError as e: 13 | print("SparseConv is not supported", flush=True) 14 | print(e, flush=True) 15 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/builder.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | from ..utils.meta import Registry, build_from_cfg 4 | 5 | BACKBONES = Registry("backbone") 6 | APPLICATIONS = Registry("applications") 7 | REGHEADS = Registry("regression_head") 8 | 9 | POLICYNETWORKS = Registry("policy_network") 10 | VALUENETWORKS = Registry("value_network") 11 | MODELNETWORKS = Registry("model_network") 12 | 13 | 14 | def build(cfg, registry, default_args=None): 15 | if cfg is None: 16 | return None 17 | elif isinstance(cfg, (list, tuple)): 18 | modules = [build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg] 19 | return modules # nn.Sequential(*modules) 20 | else: 21 | return build_from_cfg(cfg, registry, default_args) 22 | 23 | 24 | def build_reg_head(cfg): 25 | return build(cfg, REGHEADS) 26 | 27 | 28 | def build_backbone(cfg): 29 | return build(cfg, BACKBONES) 30 | 31 | 32 | def build_model(cfg, default_args=None): 33 | if cfg is None: 34 | return None 35 | for model_type in [BACKBONES, POLICYNETWORKS, VALUENETWORKS, MODELNETWORKS]: 36 | if cfg["type"] in model_type.module_dict: 37 | return build(cfg, model_type, default_args) 38 | raise RuntimeError(f"This model type:{cfg['type']} does not exist!") 39 | 40 | 41 | def build_actor_critic(actor_cfg, critic_cfg, shared_backbone=False): 42 | if shared_backbone: 43 | assert ( 44 | actor_cfg["nn_cfg"]["type"] in ["Visuomotor", "FrameMiners"] 45 | or "Visuomotor" in actor_cfg["nn_cfg"]["type"] 46 | ), f"Only Visuomotor models can share visual backbone. Your model has type {actor_cfg['nn_cfg']['type']}!" 47 | actor = build_model(actor_cfg) 48 | 49 | if getattr(actor.backbone, "visual_nn", None) is not None: 50 | critic_cfg.nn_cfg.visual_nn_cfg = None 51 | critic_cfg.nn_cfg.visual_nn = actor.backbone.visual_nn 52 | 53 | critic = build_model(critic_cfg) 54 | return actor, critic 55 | else: 56 | return build_model(actor_cfg), build_model(critic_cfg) 57 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .conv import CONV_LAYERS, build_conv_layer 2 | from .linear import LINEAR_LAYERS, build_linear_layer 3 | from .activation import ACTIVATION_LAYERS, build_activation_layer 4 | from .norm import NORM_LAYERS, build_norm_layer, need_bias 5 | from .padding import PADDING_LAYERS, build_padding_layer 6 | from .weight_init import constant_init, normal_init, kaiming_init, uniform_init, build_init, delta_orthogonal_init 7 | 8 | # from .conv_module import PLUGIN_LAYERS, ConvModule 9 | from .block_utils import NN_BLOCKS, build_nn_block, BasicBlock, FlexibleBasicBlock, LinearModule, ConvModule, MLP, SharedMLP 10 | 11 | try: 12 | from .pn2_modules import * 13 | except ImportError as e: 14 | print("Pointnet++ is not compiled") 15 | print(e) 16 | 17 | from .attention import AttentionPooling, MultiHeadSelfAttention, MultiHeadAttention, ATTENTION_LAYERS, build_attention_layer 18 | from .plugin import PLUGIN_LAYERS, build_plugin_layer 19 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/modules/activation.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from maniskill2_learn.utils.meta import build_from_cfg, Registry 5 | 6 | ACTIVATION_LAYERS = Registry("activation layer") 7 | for module in [ 8 | nn.ELU, 9 | nn.Hardshrink, 10 | nn.Hardsigmoid, 11 | nn.Hardtanh, 12 | nn.Hardswish, 13 | nn.LeakyReLU, 14 | nn.LogSigmoid, 15 | nn.MultiheadAttention, 16 | nn.PReLU, 17 | nn.ReLU, 18 | nn.ReLU6, 19 | nn.RReLU, 20 | nn.SELU, 21 | nn.CELU, 22 | nn.GELU, 23 | nn.Sigmoid, 24 | nn.SiLU, 25 | # nn.Mish, 26 | nn.Softplus, 27 | nn.Softshrink, 28 | nn.Softsign, 29 | nn.Tanh, 30 | nn.Tanhshrink, 31 | nn.Threshold, 32 | nn.Softmin, 33 | nn.Softmax, 34 | nn.Softmax2d, 35 | nn.LogSoftmax, 36 | nn.AdaptiveLogSoftmaxWithLoss, 37 | ]: 38 | ACTIVATION_LAYERS.register_module(module=module) 39 | 40 | 41 | @ACTIVATION_LAYERS.register_module(name="Clip") 42 | @ACTIVATION_LAYERS.register_module() 43 | class Clamp(nn.Module): 44 | def __init__(self, min=-1.0, max=1.0): 45 | super(Clamp, self).__init__() 46 | self.min = min 47 | self.max = max 48 | 49 | def forward(self, x): 50 | return torch.clamp(x, min=self.min, max=self.max) 51 | 52 | 53 | try: 54 | import torchsparse.nn as spnn 55 | 56 | ACTIVATION_LAYERS.register_module(name="SparseReLU", module=spnn.ReLU) 57 | ACTIVATION_LAYERS.register_module(name="SparseLeakyReLU", module=spnn.LeakyReLU) 58 | except ImportError as e: 59 | print("Spconv is not installed!") 60 | print(e) 61 | 62 | 63 | def build_activation_layer(cfg, default_args=None): 64 | return build_from_cfg(cfg, ACTIVATION_LAYERS, default_args) 65 | 66 | 67 | INPLACE_ACTIVATIONS = ["ELU", "Hardsigmoid", "Hardtanh", "Hardswish", "ReLU", "LeakyReLU", "ReLU6", "RReLU", "SELU", "CELU", "SiLU", "Threshold"] 68 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/modules/cnn_modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .resnet_utils import BasicBlock, Bottleneck, ResLayer 2 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/modules/conv.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from maniskill2_learn.utils.meta import Registry, build_from_cfg 5 | 6 | 7 | CONV_LAYERS = Registry("conv layer") 8 | for module in [ 9 | nn.Conv1d, 10 | nn.Conv2d, 11 | nn.Conv3d, 12 | nn.ConvTranspose1d, 13 | nn.ConvTranspose2d, 14 | nn.ConvTranspose3d, 15 | nn.LazyConv1d, 16 | nn.LazyConv2d, 17 | nn.LazyConv3d, 18 | nn.LazyConvTranspose1d, 19 | nn.LazyConvTranspose2d, 20 | nn.LazyConvTranspose3d, 21 | nn.Unfold, 22 | nn.Fold, 23 | ]: 24 | CONV_LAYERS.register_module(module=module) 25 | CONV_LAYERS.register_module("Conv", module=nn.Conv2d) 26 | CONV_LAYERS.register_module("Deconv", module=nn.ConvTranspose2d) 27 | 28 | # SparseConv 29 | SPARSE_CONV_LAYERS = Registry("sparse conv layer") 30 | 31 | 32 | @CONV_LAYERS.register_module() 33 | class Conv2dAdaptivePadding(nn.Conv2d): 34 | """ 35 | Implementation of 2D convolution in tensorflow with `padding` as "same", 36 | which applies padding to input so that input image gets fully covered by filter and stride you specified. 37 | For example: 38 | With stride 1, this will ensure that output image size is same as input. 39 | With stride 2, output dimensions will be half. 40 | Args: 41 | in_channels (int): Number of channels in the input image 42 | out_channels (int): Number of channels produced by the convolution 43 | kernel_size (int or tuple): Size of the convolving kernel 44 | stride (int or tuple, optional): Stride of the convolution. Default: 1 45 | padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0 46 | dilation (int or tuple, optional): Spacing between kernel elements. Default: 1 47 | groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1 48 | bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True`` 49 | """ 50 | 51 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True): 52 | super().__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias) 53 | 54 | def forward(self, x): 55 | img_h, img_w = x.size()[-2:] 56 | kernel_h, kernel_w = self.weight.size()[-2:] 57 | stride_h, stride_w = self.stride 58 | output_h = math.ceil(img_h / stride_h) 59 | output_w = math.ceil(img_w / stride_w) 60 | pad_h = max((output_h - 1) * self.stride[0] + (kernel_h - 1) * self.dilation[0] + 1 - img_h, 0) 61 | pad_w = max((output_w - 1) * self.stride[1] + (kernel_w - 1) * self.dilation[1] + 1 - img_w, 0) 62 | if pad_h > 0 or pad_w > 0: 63 | x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]) 64 | return F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups) 65 | 66 | 67 | def build_conv_layer(cfg, *args, **kwargs): 68 | if cfg is None: 69 | cfg_ = dict(type="Conv2d") 70 | else: 71 | if not isinstance(cfg, dict): 72 | raise TypeError("cfg must be a dict") 73 | if "type" not in cfg: 74 | raise KeyError('the cfg dict must contain the key "type"') 75 | cfg_ = cfg.copy() 76 | layer_type = cfg_.pop("type") 77 | if layer_type not in CONV_LAYERS: 78 | raise KeyError(f"Unrecognized norm type {layer_type}") 79 | else: 80 | conv_layer = CONV_LAYERS.get(layer_type) 81 | layer = conv_layer(*args, **kwargs, **cfg_) 82 | return layer 83 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/modules/padding.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from maniskill2_learn.utils.meta import Registry, build_from_cfg 3 | 4 | PADDING_LAYERS = Registry("padding layer") 5 | 6 | for module in [ 7 | nn.ReflectionPad1d, 8 | nn.ReflectionPad2d, 9 | nn.ReplicationPad1d, 10 | nn.ReplicationPad2d, 11 | nn.ReplicationPad3d, 12 | nn.ZeroPad2d, 13 | nn.ConstantPad1d, 14 | nn.ConstantPad2d, 15 | nn.ConstantPad3d, 16 | ]: 17 | PADDING_LAYERS.register_module(module=module) 18 | 19 | 20 | def build_padding_layer(cfg, default_args=None): 21 | return build_from_cfg(cfg, PADDING_LAYERS, default_args) 22 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/modules/pct_modules/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/haosulab/ManiSkill2-Learn/e9860aff595da3ad34829d62c923be0b50001427/maniskill2_learn/networks/modules/pct_modules/__init__.py -------------------------------------------------------------------------------- /maniskill2_learn/networks/modules/plugin.py: -------------------------------------------------------------------------------- 1 | import inspect, platform 2 | from maniskill2_learn.utils.meta import Registry, build_from_cfg 3 | 4 | PLUGIN_LAYERS = Registry("plugin layer") 5 | 6 | if platform.system() == "Windows": 7 | import regex as re 8 | else: 9 | import re 10 | 11 | 12 | def infer_abbr(class_type): 13 | """Infer abbreviation from the class name. 14 | 15 | This method will infer the abbreviation to map class types to 16 | abbreviations. 17 | 18 | Rule 1: If the class has the property "abbr", return the property. 19 | Rule 2: Otherwise, the abbreviation falls back to snake case of class 20 | name, e.g. the abbreviation of ``FancyBlock`` will be ``fancy_block``. 21 | 22 | Args: 23 | class_type (type): The norm layer type. 24 | 25 | Returns: 26 | str: The inferred abbreviation. 27 | """ 28 | 29 | def camel2snack(word): 30 | """Convert camel case word into snack case. 31 | 32 | Modified from `inflection lib 33 | `_. 34 | 35 | Example:: 36 | 37 | >>> camel2snack("FancyBlock") 38 | 'fancy_block' 39 | """ 40 | 41 | word = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", word) 42 | word = re.sub(r"([a-z\d])([A-Z])", r"\1_\2", word) 43 | word = word.replace("-", "_") 44 | return word.lower() 45 | 46 | if not inspect.isclass(class_type): 47 | raise TypeError(f"class_type must be a type, but got {type(class_type)}") 48 | if hasattr(class_type, "_abbr_"): 49 | return class_type._abbr_ 50 | else: 51 | return camel2snack(class_type.__name__) 52 | 53 | 54 | def build_plugin_layer(cfg, postfix="", **kwargs): 55 | """Build plugin layer. 56 | 57 | Args: 58 | cfg (None or dict): cfg should contain: 59 | type (str): identify plugin layer type. 60 | layer args: args needed to instantiate a plugin layer. 61 | postfix (int, str): appended into norm abbreviation to 62 | create named layer. Default: ''. 63 | 64 | Returns: 65 | tuple[str, nn.Module]: 66 | name (str): abbreviation + postfix 67 | layer (nn.Module): created plugin layer 68 | """ 69 | if not isinstance(cfg, dict): 70 | raise TypeError("cfg must be a dict") 71 | if "type" not in cfg: 72 | raise KeyError('the cfg dict must contain the key "type"') 73 | cfg_ = cfg.copy() 74 | 75 | layer_type = cfg_.pop("type") 76 | if layer_type not in PLUGIN_LAYERS: 77 | raise KeyError(f"Unrecognized plugin type {layer_type}") 78 | 79 | plugin_layer = PLUGIN_LAYERS.get(layer_type) 80 | abbr = infer_abbr(plugin_layer) 81 | 82 | assert isinstance(postfix, (int, str)) 83 | name = abbr + str(postfix) 84 | 85 | layer = plugin_layer(**kwargs, **cfg_) 86 | 87 | return name, layer 88 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/modules/pooling.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from maniskill2_learn.utils.meta import Registry, build_from_cfg 5 | 6 | 7 | CONV_LAYERS = Registry("conv layer") 8 | for module in [ 9 | nn.Conv1d, 10 | nn.Conv2d, 11 | nn.Conv3d, 12 | nn.ConvTranspose1d, 13 | nn.ConvTranspose2d, 14 | nn.ConvTranspose3d, 15 | nn.LazyConv1d, 16 | nn.LazyConv2d, 17 | nn.LazyConv3d, 18 | nn.LazyConvTranspose1d, 19 | nn.LazyConvTranspose2d, 20 | nn.LazyConvTranspose3d, 21 | nn.Unfold, 22 | nn.Fold, 23 | ]: 24 | CONV_LAYERS.register_module(module=module) 25 | CONV_LAYERS.register_module("Conv", module=nn.Conv2d) 26 | CONV_LAYERS.register_module("Deconv", module=nn.ConvTranspose2d) 27 | 28 | # SparseConv 29 | SPARSE_CONV_LAYERS = Registry("sparse conv layer") 30 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/modules/spconv_modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .spconv_utils import initial_voxelize, point_to_voxel, voxel_to_point, build_points 2 | from .resnet_utils import BasicConvolutionBlock, BasicDeconvolutionBlock, ResidualBlock, Bottleneck, build_sparse_norm 3 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/modules/spconv_modules/resnet_utils.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torchsparse 3 | import torchsparse.nn as spnn 4 | 5 | 6 | def build_sparse_norm(channels, use_ln=True): 7 | return spnn.LayerNorm(channels, eps=1e-6) if use_ln else spnn.BatchNorm(channels) 8 | 9 | 10 | class BasicConvolutionBlock(nn.Module): 11 | def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, use_ln=False): 12 | super().__init__() 13 | self.net = nn.Sequential( 14 | spnn.Conv3d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, transposed=False), 15 | build_sparse_norm(out_channels, use_ln), 16 | spnn.ReLU(True), 17 | ) 18 | 19 | def forward(self, x): 20 | return self.net(x) 21 | 22 | 23 | class BasicDeconvolutionBlock(nn.Module): 24 | def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, use_ln=False): 25 | super().__init__() 26 | self.net = nn.Sequential( 27 | spnn.Conv3d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, transposed=True), 28 | build_sparse_norm(out_channels, use_ln), 29 | spnn.ReLU(True), 30 | ) 31 | 32 | def forward(self, x): 33 | return self.net(x) 34 | 35 | 36 | class ResidualBlock(nn.Module): 37 | expansion = 1 38 | 39 | def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, use_ln=False): 40 | super().__init__() 41 | self.net = nn.Sequential( 42 | spnn.Conv3d(in_channels, out_channels, kernel_size=kernel_size, dilation=dilation, stride=stride), 43 | build_sparse_norm(out_channels, use_ln), 44 | spnn.ReLU(True), 45 | spnn.Conv3d(out_channels, out_channels, kernel_size=kernel_size, dilation=dilation, stride=1), 46 | build_sparse_norm(out_channels, use_ln), 47 | ) 48 | 49 | if in_channels == out_channels * self.expansion and stride == 1: 50 | self.downsample = nn.Sequential() 51 | else: 52 | if stride == 1: 53 | self.downsample = nn.Sequential( 54 | spnn.Conv3d(in_channels, out_channels, kernel_size=1, dilation=1, stride=stride), 55 | build_sparse_norm(out_channels, use_ln), 56 | ) 57 | else: 58 | self.downsample = nn.Sequential( 59 | spnn.Conv3d(in_channels, out_channels, kernel_size=kernel_size, dilation=1, stride=stride), 60 | build_sparse_norm(out_channels, use_ln), 61 | ) 62 | 63 | self.relu = spnn.ReLU(True) 64 | 65 | def forward(self, x): 66 | return self.relu(self.net(x) + self.downsample(x)) 67 | 68 | 69 | class Bottleneck(nn.Module): 70 | expansion = 4 71 | 72 | def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, use_ln=False): 73 | super(Bottleneck, self).__init__() 74 | 75 | self.net = nn.Sequential( 76 | spnn.Conv3d(in_channels, out_channels, kernel_size=1), 77 | build_sparse_norm(out_channels, use_ln), 78 | spnn.ReLU(True), 79 | spnn.Conv3d(out_channels, out_channels, kernel_size=kernel_size, dilation=dilation, stride=stride), 80 | build_sparse_norm(out_channels, use_ln), 81 | spnn.ReLU(True), 82 | spnn.Conv3d(out_channels, out_channels * self.expansion, kernel_size=1), 83 | build_sparse_norm(out_channels * self.expansion, use_ln), 84 | ) 85 | 86 | if in_channels == out_channels * self.expansion and stride == 1: 87 | self.downsample = nn.Sequential() 88 | else: 89 | if stride == 1: 90 | self.downsample = nn.Sequential( 91 | spnn.Conv3d(in_channels, out_channels * self.expansion, kernel_size=1, dilation=1, stride=stride), 92 | build_sparse_norm(out_channels * self.expansion, use_ln), 93 | ) 94 | else: 95 | self.downsample = nn.Sequential( 96 | spnn.Conv3d(in_channels, out_channels * self.expansion, kernel_size=kernel_size, dilation=1, stride=stride), 97 | build_sparse_norm(out_channels * self.expansion, use_ln), 98 | ) 99 | 100 | self.relu = spnn.ReLU(True) 101 | 102 | def forward(self, x): 103 | return self.relu(self.net(x) + self.downsample(x)) 104 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/modules/test_norm.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | 5 | import time 6 | 7 | ln_time = 0.0 8 | ln_ln_time = 0.0 9 | ln_permute1_time = 0.0 10 | ln_permute2_time = 0.0 11 | custom_ln_time = 0.0 12 | bn_time = 0.0 13 | 14 | T = 8 15 | device = 'cuda:0' 16 | size=(16,256,1024,64) 17 | weight = torch.rand(size[1]).to(device) 18 | bias = torch.randn(size[1]).to(device) 19 | 20 | class LayerNorm(nn.Module): 21 | r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 22 | The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 23 | shape (batch_size, height, width, channels) while channels_first corresponds to inputs 24 | with shape (batch_size, channels, height, width). 25 | """ 26 | def __init__(self, normalized_shape, eps=1e-8, data_format="channels_last"): 27 | super().__init__() 28 | self.weight = nn.Parameter(weight.clone()) 29 | self.bias = nn.Parameter(bias.clone()) 30 | self.eps = eps 31 | self.data_format = data_format 32 | if self.data_format not in ["channels_last", "channels_first"]: 33 | raise NotImplementedError 34 | self.normalized_shape = (normalized_shape, ) 35 | 36 | def forward(self, x): 37 | if self.data_format == "channels_last": 38 | return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) 39 | elif self.data_format == "channels_first": 40 | u = x.mean(1, keepdim=True) 41 | s = (x - u).pow(2).mean(1, keepdim=True) 42 | x = (x - u) / torch.sqrt(s + self.eps) 43 | x = self.weight[:, None, None] * x + self.bias[:, None, None] 44 | return x 45 | 46 | 47 | ln = nn.LayerNorm(size[1], eps=1e-8).to(device) 48 | ln.weight = nn.Parameter(weight.clone()) 49 | ln.bias = nn.Parameter(bias.clone()) 50 | ln_custom = LayerNorm(size[1], data_format="channels_first").to(device) 51 | 52 | a = torch.randn(size).to(device) 53 | b = a.permute(0,2,3,1).contiguous() 54 | ln_a = ln(b) 55 | ln_a = ln_a.permute(0,3,1,2).contiguous() 56 | ln_custom_a = ln_custom(a) 57 | print("|ln - ln_custom| max", torch.max(torch.abs(ln_a - ln_custom_a))) 58 | 59 | del a,b 60 | 61 | 62 | 63 | ln_custom = LayerNorm(size[1], data_format="channels_first").to(device) 64 | for t in range(T): 65 | print(t) 66 | a = torch.randn(size).to(device) 67 | torch.cuda.synchronize() 68 | tt = time.time() 69 | b=ln_custom(a) 70 | torch.cuda.synchronize() 71 | custom_ln_time += time.time() - tt 72 | del b 73 | print("custom_ln_time", custom_ln_time / T) 74 | 75 | ln = nn.LayerNorm(size[1]).to(device) 76 | ln.weight = nn.Parameter(weight.clone()) 77 | ln.bias = nn.Parameter(bias.clone()) 78 | for t in range(T): 79 | print(t) 80 | a = torch.randn(size).to(device) 81 | torch.cuda.synchronize() 82 | 83 | tt = time.time() 84 | b=a.permute(0,2,3,1) 85 | torch.cuda.synchronize() 86 | ln_permute1_time += time.time() - tt 87 | 88 | tt=time.time() 89 | b=ln(b) 90 | torch.cuda.synchronize() 91 | ln_ln_time += time.time() - tt 92 | 93 | tt=time.time() 94 | b=b.permute(0,3,1,2) 95 | torch.cuda.synchronize() 96 | ln_permute2_time += time.time() - tt 97 | del b 98 | 99 | print("pytorch ln_ln_time", ln_ln_time / T, "ln_permute1_time", ln_permute1_time / T, "ln_permute2_time", ln_permute2_time / T) 100 | print("pytorch ln_total_time", (ln_ln_time + ln_permute1_time + ln_permute2_time) / T) 101 | 102 | # bn = nn.BatchNorm2d(size[1]).to(device) 103 | # for t in range(T): 104 | # print(t) 105 | # a = torch.randn(size).to(device) 106 | # torch.cuda.synchronize() 107 | # tt = time.time() 108 | # b=bn(a) 109 | # torch.cuda.synchronize() 110 | # bn_time += time.time() - tt 111 | # del b 112 | # print("bn_time", bn_time / T) 113 | 114 | 115 | 116 | 117 | 118 | 119 | 120 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/ops/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/haosulab/ManiSkill2-Learn/e9860aff595da3ad34829d62c923be0b50001427/maniskill2_learn/networks/ops/__init__.py -------------------------------------------------------------------------------- /maniskill2_learn/networks/ops/ops_3d/__init__.py: -------------------------------------------------------------------------------- 1 | from .ball_query import ball_query 2 | from .pcd_process import downsample_pcd 3 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/ops/ops_3d/ball_query/__init__.py: -------------------------------------------------------------------------------- 1 | from .ball_query import ball_query 2 | 3 | __all__ = ['ball_query'] 4 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/ops/ops_3d/ball_query/ball_query.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | from . import ball_query_ext 4 | 5 | 6 | class BallQuery(Function): 7 | """Ball Query. 8 | 9 | Find nearby points in spherical space. 10 | """ 11 | 12 | @staticmethod 13 | def forward(ctx, min_radius: float, max_radius: float, sample_num: int, 14 | xyz: torch.Tensor, center_xyz: torch.Tensor) -> torch.Tensor: 15 | """forward. 16 | Args: 17 | min_radius (float): minimum radius of the balls. 18 | max_radius (float): maximum radius of the balls. 19 | sample_num (int): maximum number of features in the balls. 20 | If sammple_num=-1, it will outputs all points in the range. 21 | xyz (Tensor): (B, N, 3) xyz coordinates of the features. 22 | center_xyz (Tensor): (B, npoint, 3) centers of the ball query. 23 | Returns: 24 | Tensor: (B, npoint, nsample) tensor with the indicies of the features that form the query balls. 25 | """ 26 | assert center_xyz.is_contiguous() and center_xyz.shape[-1] == 3 27 | assert xyz.is_contiguous() and xyz.shape[-1] == 3 28 | assert min_radius < max_radius 29 | 30 | if sample_num == -1: 31 | distance = torch.linalg.norm(xyz[:, None] - center_xyz[:, :, None], dim=-1) 32 | mask = torch.logical_and(min_radius <= distance, distance <= max_radius) 33 | sample_num = mask.long().sum(-1).max() 34 | B, N, _ = xyz.size() 35 | npoint = center_xyz.size(1) 36 | idx = torch.cuda.IntTensor(B, npoint, sample_num).zero_() 37 | 38 | ball_query_ext.ball_query_wrapper(B, N, npoint, min_radius, max_radius, 39 | sample_num, center_xyz, xyz, idx) 40 | ctx.mark_non_differentiable(idx) 41 | return idx 42 | 43 | @staticmethod 44 | def backward(ctx, a=None): 45 | return None, None, None, None 46 | 47 | 48 | ball_query = BallQuery.apply 49 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/ops/ops_3d/ball_query/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/ball_query.cpp 3 | 4 | // #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | 14 | #define CHECK_CUDA(x) \ 15 | TORCH_CHECK(x.device().is_cuda(), #x, " must be a CUDAtensor ") 16 | #define CHECK_CONTIGUOUS(x) \ 17 | TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ") 18 | #define CHECK_INPUT(x) \ 19 | CHECK_CUDA(x); \ 20 | CHECK_CONTIGUOUS(x) 21 | 22 | int ball_query_wrapper(int b, int n, int m, float min_radius, float max_radius, int nsample, 23 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, 24 | at::Tensor idx_tensor); 25 | 26 | void ball_query_kernel_launcher(int b, int n, int m, float min_radius, float max_radius, 27 | int nsample, const float *xyz, const float *new_xyz, 28 | int *idx, cudaStream_t stream); 29 | 30 | int ball_query_wrapper(int b, int n, int m, float min_radius, float max_radius, int nsample, 31 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, 32 | at::Tensor idx_tensor) { 33 | CHECK_INPUT(new_xyz_tensor); 34 | CHECK_INPUT(xyz_tensor); 35 | const float *new_xyz = new_xyz_tensor.data_ptr(); 36 | const float *xyz = xyz_tensor.data_ptr(); 37 | int *idx = idx_tensor.data_ptr(); 38 | 39 | cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); 40 | ball_query_kernel_launcher(b, n, m, min_radius, max_radius, 41 | nsample, new_xyz, xyz, idx, stream); 42 | return 1; 43 | } 44 | 45 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 46 | m.def("ball_query_wrapper", &ball_query_wrapper, "ball_query_wrapper"); 47 | } 48 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/ops/ops_3d/ball_query/src/ball_query_cuda.cu: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/ball_query_gpu.cu 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #define THREADS_PER_BLOCK 256 9 | #define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0)) 10 | 11 | __global__ void ball_query_kernel(int b, int n, int m, 12 | float min_radius, 13 | float max_radius, 14 | int nsample, 15 | const float *__restrict__ new_xyz, 16 | const float *__restrict__ xyz, 17 | int *__restrict__ idx) { 18 | // new_xyz: (B, M, 3) 19 | // xyz: (B, N, 3) 20 | // output: 21 | // idx: (B, M, nsample) 22 | int bs_idx = blockIdx.y; 23 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 24 | if (bs_idx >= b || pt_idx >= m) return; 25 | 26 | new_xyz += bs_idx * m * 3 + pt_idx * 3; 27 | xyz += bs_idx * n * 3; 28 | idx += bs_idx * m * nsample + pt_idx * nsample; 29 | 30 | float max_radius2 = max_radius * max_radius; 31 | float min_radius2 = min_radius * min_radius; 32 | float new_x = new_xyz[0]; 33 | float new_y = new_xyz[1]; 34 | float new_z = new_xyz[2]; 35 | 36 | int cnt = 0; 37 | for (int k = 0; k < n; ++k) { 38 | float x = xyz[k * 3 + 0]; 39 | float y = xyz[k * 3 + 1]; 40 | float z = xyz[k * 3 + 2]; 41 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 42 | (new_z - z) * (new_z - z); 43 | if (d2 == 0 || (d2 >= min_radius2 && d2 < max_radius2)) { 44 | if (cnt == 0) { 45 | for (int l = 0; l < nsample; ++l) { 46 | idx[l] = k; 47 | } 48 | } 49 | idx[cnt] = k; 50 | ++cnt; 51 | if (cnt >= nsample) break; 52 | } 53 | } 54 | } 55 | 56 | void ball_query_kernel_launcher(int b, int n, int m, float min_radius, float max_radius, 57 | int nsample, const float *new_xyz, const float *xyz, 58 | int *idx, cudaStream_t stream) { 59 | // new_xyz: (B, M, 3) 60 | // xyz: (B, N, 3) 61 | // output: 62 | // idx: (B, M, nsample) 63 | 64 | cudaError_t err; 65 | 66 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), 67 | b); // blockIdx.x(col), blockIdx.y(row) 68 | dim3 threads(THREADS_PER_BLOCK); 69 | 70 | ball_query_kernel<<>>(b, n, m, min_radius, max_radius, 71 | nsample, new_xyz, xyz, idx); 72 | // cudaDeviceSynchronize(); // for using printf in kernel function 73 | err = cudaGetLastError(); 74 | if (cudaSuccess != err) { 75 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 76 | exit(-1); 77 | } 78 | } 79 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/ops/ops_3d/pcd_process/__init__.py: -------------------------------------------------------------------------------- 1 | from .pcd_process import downsample_pcd 2 | 3 | __all__ = ["downsample_pcd"] 4 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/ops/ops_3d/pcd_process/pcd_process.cpp: -------------------------------------------------------------------------------- 1 | /* Modified from [Use double in some computation to improve numerical stability 2 | * https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/ball_query.cpp 3 | */ 4 | 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | 12 | // at::Tensor VoxelDownSample(const at::Tensor &xyz, const float &z_min, const int &num); 13 | 14 | // at::Tensor ManiSkillDownSample(const at::Tensor &xyz, const float &min_z, const int &num); 15 | 16 | std::tuple UniformDownSample(const at::Tensor &xyz, const float &min_z, const int &num); 17 | 18 | std::tuple ManiSkillDownSample(const at::Tensor &xyz, const at::Tensor &seg, const float &min_z, const int &num, const int &num_min, 19 | const int &num_fg); 20 | 21 | // at::Tensor cumsum(const at::Tensor &mask); 22 | 23 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 24 | // m.def("voxel_downsample", &VoxelDownsample, "VoxelDownsample"); 25 | m.def("uniform_downsample", &UniformDownSample, "UniformDownSample", py::arg("xyz"), py::arg("min_z"), py::arg("num")); 26 | m.def("maniskill_downsample", &ManiSkillDownSample, "ManiSkillDownSample", py::arg("xyz"), py::arg("seg"), py::arg("min_z"), py::arg("num"), 27 | py::arg("num_min"), py::arg("num_fg")); 28 | } 29 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/ops/ops_3d/pcd_process/pcd_process.py: -------------------------------------------------------------------------------- 1 | import torch, torch.nn.functional as F 2 | import torch.utils.cpp_extension 3 | from torch import nn as nn 4 | from maniskill2_learn.utils.torch import no_grad 5 | import os.path as osp, time 6 | from . import pcd_process_ext 7 | 8 | 9 | __folder__ = osp.dirname(__file__) 10 | 11 | 12 | # pcd_process_ext = torch.utils.cpp_extension.load( 13 | # name="pcd_process_ext", 14 | # sources=[osp.join(__folder__, "pcd_process.cpp"), osp.join(__folder__, "pcd_process.cu")], 15 | # extra_cflags=["-O3", "-std=c++17"], 16 | # verbose=True, 17 | # ) 18 | 19 | 20 | def ravel_multi_index(multi_index: torch.Tensor): 21 | """ 22 | Args: 23 | multi_index: [N, D] tensor. 24 | Returns: 25 | raveled_indices: [N] tensor. 26 | """ 27 | min_index, max_index = multi_index.min(dim=0).values, multi_index.max(dim=0).values 28 | grid_size, multi_index = max_index - min_index + 1, multi_index - min_index 29 | 30 | grid_coef = torch.cumprod(grid_size, 0) 31 | grid_coef = torch.flip(grid_coef, [0]) 32 | grid_coef = F.pad(grid_coef, [0, 1], value=1)[1:] 33 | print(grid_size, grid_coef) 34 | raveled_indices = (multi_index * grid_coef).sum(dim=-1) 35 | return raveled_indices 36 | 37 | 38 | def unique(x, dim=None): 39 | """Unique elements of x and indices of those unique elements 40 | https://github.com/pytorch/pytorch/issues/36748#issuecomment-619514810 41 | 42 | e.g. 43 | 44 | unique(tensor([ 45 | [1, 2, 3], 46 | [1, 2, 4], 47 | [1, 2, 3], 48 | [1, 2, 5] 49 | ]), dim=0) 50 | => (tensor([[1, 2, 3], 51 | [1, 2, 4], 52 | [1, 2, 5]]), 53 | tensor([0, 1, 3])) 54 | """ 55 | unique, inverse = torch.unique(x, sorted=True, return_inverse=True, dim=dim) 56 | perm = torch.arange(inverse.size(0), dtype=inverse.dtype, device=inverse.device) 57 | inverse, perm = inverse.flip([0]), perm.flip([0]) 58 | return unique, inverse.new_empty(unique.size(0)).scatter_(0, inverse, perm) 59 | 60 | 61 | @no_grad 62 | def downsample_pcd(xyz, mode="uniform", **kwargs): 63 | if mode == "maniskill": 64 | index, mask = pcd_process_ext.maniskill_downsample(xyz, **kwargs) 65 | elif mode == "voxel": 66 | index, mask = pcd_process_ext.voxel_downsample(xyz, **kwargs) 67 | elif mode == "uniform": 68 | index, mask = pcd_process_ext.uniform_downsample(xyz, **kwargs) 69 | """ 70 | index = torch.randperm(xyz.shape[1], device=xyz.device) 71 | ret_index = torch.zeros([xyz.shape[0], num + 1], device=xyz.device, dtype=torch.long) 72 | ret_mask = torch.zeros([xyz.shape[0], num + 1], device=xyz.device, dtype=torch.bool) 73 | mask = xyz[:, index, -1] > min_z 74 | cum_idx = torch.cumsum(mask.long(), dim=-1) 75 | # mask = torch.logical_and(mask, cum_idx <= num) 76 | # final_idx = torch.where(mask, cum_idx - 1, num) 77 | # ret_index.scatter_(-1, final_idx, index.expand_as(xyz[..., 0])) 78 | # ret_mask.scatter_(-1, final_idx, mask) 79 | # ret_index = ret_index[:, :-1] 80 | # ret_mask = ret_mask[:, :-1] 81 | torch.cuda.synchronize() 82 | """ 83 | return index, mask 84 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/regression_heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .deterministic import TanhHead, BasicHead 2 | from .regression_base import DiscreteBaseHead 3 | from .gaussian import GaussianHead, TanhGaussianHead, SoftplusGaussianHead 4 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/regression_heads/deterministic.py: -------------------------------------------------------------------------------- 1 | from maniskill2_learn.utils.torch.distributions import ScaledNormal 2 | import torch.nn as nn, torch, numpy as np 3 | from torch.nn import Parameter 4 | from ..builder import REGHEADS 5 | from maniskill2_learn.utils.data import is_num 6 | from torch.distributions import Normal, Categorical, MixtureSameFamily 7 | from maniskill2_learn.utils.torch import ExtendedModule, CustomIndependent 8 | from .regression_base import ContinuousBaseHead 9 | 10 | 11 | class DeterministicHead(ContinuousBaseHead): 12 | def __init__(self, bound=None, dim_output=None, clip_return=False, num_heads=1, nn_cfg=None, noise_std=0.1): 13 | super(DeterministicHead, self).__init__(bound=bound, dim_output=dim_output, clip_return=clip_return, num_heads=num_heads, nn_cfg=nn_cfg) 14 | self.dim_feature = self.dim_output if self.num_heads == 1 else (self.dim_output + 1) * self.num_heads 15 | if dim_output is not None: 16 | if is_num(noise_std): 17 | noise_std = np.ones(self.dim_output) * noise_std 18 | assert noise_std.shape[-1] == dim_output 19 | # The noise is the Gaussian noise on normalized action space for exploration. 20 | self.noise_std = Parameter(self.scale.data * torch.tensor(noise_std)) 21 | 22 | def split_feature(self, feature, num_actions=1): 23 | assert feature.shape[-1] == self.dim_feature 24 | feature = feature.repeat_interleave(num_actions, dim=0) 25 | 26 | if self.num_heads > 1: 27 | logits = feature[..., : self.num_heads] 28 | feature = feature[..., self.num_heads :] 29 | pred_shape = list(feature.shape) 30 | pred_dim = pred_shape[-1] // self.num_heads 31 | mean_shape = pred_shape[:-1] + [self.num_heads, pred_dim] 32 | mean = feature.reshape(*mean_shape) 33 | else: 34 | logits = None 35 | mean = feature 36 | std = self.noise_std.data.expand_as(mean) 37 | return logits, mean, std 38 | 39 | def return_with_mean_std(self, mean, std, mode, logits=None): 40 | if self.num_heads > 1: 41 | logits_max = logits.argmax(-1) 42 | logits_max = logits_max[..., None, None].repeat_interleave(mean.shape[-1], dim=-1) 43 | 44 | if mode == "mean" or mode == "eval": 45 | ret = mean if logits is None else torch.gather(mean, -2, logits_max).squeeze(-2) 46 | return ret * self.scale + self.bias 47 | 48 | dist = CustomIndependent(ScaledNormal(mean, std, self.scale, self.bias), 1) 49 | mean_ret = dist.mean 50 | std_ret = dist.stddev 51 | if self.num_heads > 1: 52 | mixture_distribution = Categorical(logits=logits) 53 | dist = MixtureSameFamily(mixture_distribution, dist) 54 | mean_ret = torch.gather(mean_ret, -2, logits_max).squeeze(-2) 55 | std_ret = torch.gather(std_ret, -2, logits_max).squeeze(-2) 56 | 57 | sample = self.clamp(dist.rsample() if dist.has_rsample else dist.sample()) 58 | if mode == "explore" or mode == "sample": 59 | return sample 60 | elif mode == "dist_mean": 61 | return dist, mean_ret 62 | elif mode == "all": 63 | log_prob = dist.log_prob(sample) 64 | return sample, log_prob[..., None], mean_ret, std_ret, dist 65 | else: 66 | raise ValueError(f"Unsupported mode {mode}!!") 67 | 68 | 69 | @REGHEADS.register_module() 70 | class BasicHead(DeterministicHead): 71 | def forward(self, feature, num_actions=1, mode="explore", **kwargs): 72 | feature = super(BasicHead, self).forward(feature) 73 | logits, mean, std = self.split_feature(feature, num_actions) 74 | return self.return_with_mean_std(mean, std, mode, logits) 75 | 76 | 77 | @REGHEADS.register_module() 78 | class TanhHead(DeterministicHead): 79 | def forward(self, feature, num_actions=1, mode="explore", **kwargs): 80 | feature = super(TanhHead, self).forward(feature) 81 | logits, mean, std = self.split_feature(feature, num_actions) 82 | mean = torch.tanh(mean) 83 | return self.return_with_mean_std(mean, std, mode, logits) 84 | -------------------------------------------------------------------------------- /maniskill2_learn/networks/regression_heads/regression_base.py: -------------------------------------------------------------------------------- 1 | import torch, torch.nn.functional as F 2 | from torch.nn import Parameter 3 | import numpy as np 4 | 5 | from maniskill2_learn.utils.data import is_num, is_not_null, to_np 6 | from maniskill2_learn.utils.torch import ExtendedModule, CustomCategorical 7 | from ..builder import build_backbone, REGHEADS 8 | 9 | 10 | class ContinuousBaseHead(ExtendedModule): 11 | def __init__(self, bound=None, dim_output=None, nn_cfg=None, clip_return=False, num_heads=1): 12 | super(ContinuousBaseHead, self).__init__() 13 | self.bound = bound 14 | self.net = build_backbone(nn_cfg) 15 | self.clip_return = clip_return and is_not_null(bound) 16 | 17 | if is_not_null(bound): 18 | if is_num(bound[0]): 19 | bound[0] = np.ones(dim_output) * bound[0] 20 | if is_num(bound[1]): 21 | bound[1] = np.ones(dim_output) * bound[1] 22 | assert (to_np(bound[0].shape) == to_np(bound[1].shape)).all() 23 | assert dim_output is None or bound[0].shape[-1] == dim_output 24 | dim_output = bound[0].shape[-1] 25 | if bound[0].ndim > 1: 26 | assert bound[0].ndim == 2 and bound[0].shape[0] == num_heads and num_heads > 1 27 | self.lb, self.ub = [Parameter(torch.tensor(bound[i]), requires_grad=False) for i in [0, 1]] 28 | self.log_uniform_prob = torch.log(1.0 / ((self.ub - self.lb).data)).sum().item() 29 | self.scale = Parameter(torch.tensor(bound[1] - bound[0]) / 2, requires_grad=False) 30 | self.bias = Parameter(torch.tensor(bound[0] + bound[1]) / 2, requires_grad=False) 31 | else: 32 | self.scale, self.bias = 1, 0 33 | 34 | self.dim_output = dim_output 35 | self.num_heads = num_heads 36 | self.dim_feature = None 37 | 38 | def uniform(self, sample_shape): 39 | r = torch.rand(sample_shape, self.dim_output, device=self.device) 40 | return (r * self.ub + (1 - r) * self.lb), torch.ones(sample_shape, device=self.device) * self.log_uniform_prob 41 | 42 | def clamp(self, x): 43 | if self.clip_return: 44 | x = torch.clamp(x, min=self.lb, max=self.ub) 45 | return x 46 | 47 | def forward(self, feature, **kwargs): 48 | return self.net(feature) if self.net is not None else feature 49 | 50 | 51 | @REGHEADS.register_module() 52 | class DiscreteBaseHead(ExtendedModule): 53 | def __init__(self, num_choices, num_heads=1, **kwargs): 54 | super(DiscreteBaseHead, self).__init__() 55 | assert num_heads == 1, "We only support one head recently." 56 | self.num_choices = num_choices 57 | self.num_heads = num_heads 58 | 59 | def forward(self, feature, num_actions=1, mode="explore", **kwargs): 60 | assert feature.shape[-1] == self.num_choices * self.num_heads 61 | feature = feature.repeat_interleave(num_actions, dim=0) 62 | dist = CustomCategorical(logits=feature) 63 | 64 | if "std" in mode or mode == "all": 65 | raise ValueError("Discrete action do not support log std..") 66 | 67 | if mode == "mean" or mode == "eval": 68 | return feature.argmax(dim=-1, keepdim=True) 69 | elif mode == "dist": 70 | return dist 71 | elif mode == "dist_mean": 72 | return dist, feature.argmax(dim=-1, keepdim=True) 73 | elif mode == "explore" or mode == "sample": 74 | sample = dist.rsample() if dist.has_rsample else dist.sample() 75 | return sample[..., None] 76 | elif mode == "all_discrete": 77 | # For SAC only, num_heads > 1 are not supported recently.. 78 | 79 | # sample = dist.rsample() if dist.has_rsample else dist.sample( ) 80 | # log_p = dist.log_prob(sample) 81 | # sample, log_p = dist.rsample_with_log_prob() 82 | # log_p = log_p[..., None] 83 | return feature, dist.probs, feature.argmax(dim=-1, keepdim=True), None, dist 84 | else: 85 | raise ValueError(f"Unsupported mode {mode}!!") 86 | 87 | def extra_repr(self) -> str: 88 | return f"num_actions={self.num_choices}, num_head={self.num_heads}" 89 | -------------------------------------------------------------------------------- /maniskill2_learn/schedulers/__init__.py: -------------------------------------------------------------------------------- 1 | from .custom_scheduler import build_scheduler, SCHEDULERS 2 | from .lr_scheduler import build_lr_scheduler, LRSCHEDULERS 3 | -------------------------------------------------------------------------------- /maniskill2_learn/schedulers/lr_scheduler.py: -------------------------------------------------------------------------------- 1 | from maniskill2_learn.utils.meta import Registry, build_from_cfg 2 | import torch.optim.lr_scheduler as lr_scheduler 3 | 4 | LRSCHEDULERS = Registry("scheduler of pytorch learning rate") 5 | 6 | 7 | for scheduler in [ 8 | lr_scheduler.LambdaLR, 9 | lr_scheduler.MultiplicativeLR, 10 | lr_scheduler.StepLR, 11 | lr_scheduler.MultiStepLR, 12 | lr_scheduler.ConstantLR, 13 | lr_scheduler.LinearLR, 14 | lr_scheduler.ExponentialLR, 15 | lr_scheduler.CosineAnnealingLR, 16 | lr_scheduler.ChainedScheduler, 17 | lr_scheduler.SequentialLR, 18 | lr_scheduler.ReduceLROnPlateau, 19 | lr_scheduler.CyclicLR, 20 | lr_scheduler.OneCycleLR, 21 | lr_scheduler.CosineAnnealingWarmRestarts, 22 | ]: 23 | LRSCHEDULERS.register_module(module=scheduler) 24 | 25 | 26 | def build_lr_scheduler(cfg, default_args=None): 27 | if cfg.get("type", None) == "LambdaLR": 28 | assert cfg.get("lr_lambda") is not None 29 | if isinstance(cfg["lr_lambda"], str): 30 | cfg["lr_lambda"] = eval(cfg["lr_lambda"]) 31 | 32 | return build_from_cfg(cfg, LRSCHEDULERS, default_args) 33 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/haosulab/ManiSkill2-Learn/e9860aff595da3ad34829d62c923be0b50001427/maniskill2_learn/utils/__init__.py -------------------------------------------------------------------------------- /maniskill2_learn/utils/data/__init__.py: -------------------------------------------------------------------------------- 1 | from .array_ops import ( 2 | unsqueeze, 3 | squeeze, 4 | zeros_like, 5 | ones_like, 6 | repeat, 7 | tile, 8 | shuffle, 9 | take, 10 | concat, 11 | stack, 12 | share_memory, 13 | to_item, 14 | select_with_mask, 15 | recover_with_mask, 16 | slice_to_range, 17 | detach, 18 | split, 19 | norm, 20 | normalize, 21 | clip, 22 | arr_sum, 23 | is_pcd, 24 | arr_min, 25 | arr_max, 26 | arr_mean, 27 | batch_shuffle, 28 | batch_perm, 29 | pad_item, 30 | pad_clip, 31 | clip_item, 32 | to_gc, 33 | to_nc, 34 | encode_np, 35 | decode_np, 36 | gather, 37 | to_two_dims, 38 | reshape, 39 | transpose, 40 | contiguous, 41 | split_dim, 42 | slice_item, 43 | sample_and_pad, 44 | batch_index_select, 45 | einsum, 46 | broadcast_to, 47 | deepcopy, 48 | to_float 49 | ) 50 | from .compression import float_to_int, int_to_float, f64_to_f32, to_f32, to_f16, DataCoder 51 | from .converter import as_dtype, to_np, to_torch, to_array, dict_to_str, list_to_str, dict_to_seq, seq_to_dict, slice_to_range, range_to_slice, index_to_slice 52 | from .dict_array import GDict, DictArray, SharedGDict, SharedDictArray, create_smm, delete_smm 53 | from .filtering import filter_none, filter_with_regex 54 | from .misc import equal, SLICE_ALL 55 | from .seq_utils import ( 56 | concat_seq, 57 | concat_list, 58 | concat_tuple, 59 | auto_pad_seq, 60 | flatten_seq, 61 | split_list_of_parameters, 62 | select_by_index, 63 | random_pad_clip_list, 64 | ) 65 | from .string_utils import regex_match, custom_format, prefix_match, num_to_str, float_str, regex_replace, any_string, is_regex 66 | from .type_utils import ( 67 | is_str, 68 | is_dict, 69 | is_num, 70 | is_integer, 71 | is_type, 72 | is_seq_of, 73 | is_list_of, 74 | is_tuple_of, 75 | is_iterable, 76 | get_dtype, 77 | is_np, 78 | is_np_arr, 79 | is_torch, 80 | is_arr, 81 | is_slice, 82 | is_torch_distribution, 83 | is_h5, 84 | is_null, 85 | is_not_null, 86 | ) 87 | from .dict_utils import update_dict_with_begin_keys, first_dict_key, map_dict_keys, update_dict 88 | from .wrappers import process_input, process_output, seq_to_np 89 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/data/dict_utils.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | from .type_utils import is_dict, is_seq_of 3 | 4 | 5 | def update_dict(x, y): 6 | """ 7 | Update x with y 8 | """ 9 | assert type(x) == type(y), f"{type(x), type(y)}" 10 | if is_dict(x): 11 | ret = deepcopy(x) 12 | for key in y: 13 | if key in x: 14 | ret[key] = update_dict(x[key], y[key]) 15 | else: 16 | ret[key] = deepcopy(y[key]) 17 | else: 18 | ret = deepcopy(y) 19 | return ret 20 | 21 | 22 | def update_dict_with_begin_keys(x, y, keys, begin=False, history_key=()): 23 | if len(keys) == 0: 24 | if type(x) == type(y): 25 | return update_dict(x, y) 26 | elif is_seq_of(x, dict) and is_dict(y): 27 | return [update_dict(_, y) for _ in x] 28 | else: 29 | raise NotImplementedError() 30 | if not is_dict(x): 31 | return deepcopy(x) 32 | 33 | ret = {} 34 | for key in x: 35 | if key == keys[0]: 36 | ret[key] = update_dict_with_begin_keys(x[key], y, keys[1:], True, history_key + (key,)) 37 | elif not begin: 38 | ret[key] = update_dict_with_begin_keys(x[key], y, keys, False, history_key + (key,)) 39 | else: 40 | ret[key] = deepcopy(x[key]) 41 | return ret 42 | 43 | 44 | def first_dict_key(item): 45 | return sorted(item.keys())[0] 46 | 47 | 48 | def map_dict_keys(inputs, keys_map, logger_print=None): 49 | from .string_utils import regex_replace, regex_match, is_regex 50 | import re 51 | 52 | outputs = {} 53 | for key, value in inputs.items(): 54 | new_key = key 55 | for in_pattern, out_pattern in keys_map.items(): 56 | if regex_match(key, in_pattern): 57 | new_key = regex_replace(key, in_pattern, out_pattern) 58 | break 59 | if new_key == "None" or new_key is None: 60 | if logger_print is not None: 61 | logger_print(f"Delete {key}!") 62 | continue 63 | if new_key != key and logger_print is not None: 64 | logger_print(f"Change {key} to {new_key}.") 65 | outputs[new_key] = value 66 | return outputs 67 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/data/filtering.py: -------------------------------------------------------------------------------- 1 | from .string_utils import regex_match 2 | from .type_utils import is_dict, is_tuple_of, is_list_of 3 | 4 | 5 | def custom_filter(item, func, value=True): 6 | """ 7 | Recursively filter all elements with function func. 8 | Assumptions: 9 | None means the item does not pass func. 10 | """ 11 | if is_tuple_of(item): 12 | item = list(item) 13 | if is_list_of(item): 14 | ret = [] 15 | for i in range(len(item)): 16 | x = custom_filter(item[i], func, value) 17 | if x is not None: 18 | ret.append(x) 19 | item = ret 20 | elif is_dict(item): 21 | ret = {} 22 | for key in item: 23 | x = custom_filter(item[key], func, value) 24 | if x is not None: 25 | ret[key] = x 26 | item = ret 27 | return item if not value or (item is not None and func(item)) else None 28 | 29 | 30 | def filter_none(x): 31 | func = lambda _: _ is not None 32 | return custom_filter(x, func, True) 33 | 34 | 35 | def filter_with_regex(x, regex, value=True): 36 | func = lambda _: _ is not None and regex_match(_, regex) 37 | return custom_filter(x, func, value) 38 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/data/misc.py: -------------------------------------------------------------------------------- 1 | def equal(x, y): 2 | return True if x is None or y is None else x == y 3 | 4 | 5 | SLICE_ALL = slice(None, None, None) 6 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/data/seq_utils.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | from copy import deepcopy 3 | from random import shuffle 4 | from .type_utils import is_seq_of 5 | 6 | 7 | def concat_seq(in_list, dtype): 8 | assert dtype in [list, tuple] 9 | return dtype(itertools.chain(*in_list)) 10 | 11 | 12 | def concat_list(in_list): 13 | return concat_seq(in_list, list) 14 | 15 | 16 | def concat_tuple(in_list): 17 | return concat_seq(in_list, tuple) 18 | 19 | 20 | def auto_pad_seq(a, b): 21 | """ 22 | Input two sequence, then output two list of objects with the same size. 23 | """ 24 | a = list(a) if isinstance(a, (list, tuple)) else [a] 25 | b = list(b) if isinstance(b, (list, tuple)) else [b] 26 | if len(a) > len(b): 27 | for i in range(len(a) - len(b)): 28 | b.append(a[0]) 29 | elif len(a) < len(b): 30 | for i in range(len(b) - len(a)): 31 | a.append(b[0]) 32 | return a, b 33 | 34 | 35 | def flatten_seq(x, dtype=list): 36 | if not is_seq_of(x, (tuple, list)): 37 | return x 38 | return dtype(concat_list([flatten_seq(_) for _ in x])) 39 | 40 | 41 | def split_list_of_parameters(num_procsess, *args, **kwargs): 42 | from ..math import split_num 43 | 44 | args = [_ for _ in args if _ is not None] 45 | kwargs = {_: __ for _, __ in kwargs.items() if __ is not None} 46 | assert len(args) > 0 or len(kwargs) > 0 47 | first_item = args[0] if len(args) > 0 else kwargs[list(kwargs.keys())[0]] 48 | n, running_steps = split_num(len(first_item), num_procsess) 49 | start_idx = 0 50 | paras = [] 51 | for i in range(n): 52 | slice_i = slice(start_idx, start_idx + running_steps[i]) 53 | start_idx += running_steps[i] 54 | args_i = list([_[slice_i] for _ in args]) 55 | kwargs_i = {_: kwargs[_][slice_i] for _ in kwargs} 56 | paras.append([args_i, kwargs_i]) 57 | return paras 58 | 59 | 60 | def select_by_index(files, indices): 61 | return [files[i] for i in indices] 62 | 63 | 64 | def random_pad_clip_list(x, num): 65 | x = deepcopy(list(x)) 66 | if len(x) > num: 67 | shuffle(x) 68 | return x[:num] 69 | else: 70 | ret = [] 71 | for i in range(num // len(x)): 72 | shuffle(x) 73 | ret = ret + x 74 | ret = ret + x[: num - len(ret)] 75 | return ret 76 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/data/string_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Useful regex expression 3 | 1. nothing else classifier: '^((?!classifier).)*$' 4 | 2. any string: '(.*?)' 5 | """ 6 | 7 | import re 8 | 9 | 10 | any_string = r"(.*?)" 11 | 12 | 13 | def custom_format(template_string, **kwargs): 14 | template_string = template_string.replace("{", "{{") 15 | template_string = template_string.replace("}", "}}") 16 | template_string = template_string.replace("&lformat ", "{") 17 | template_string = template_string.replace(" &rformat", "}") 18 | return template_string.format_map(kwargs) 19 | 20 | 21 | def regex_match(string, pattern): 22 | return re.match(pattern, string) is not None 23 | 24 | 25 | def regex_replace(string, pattern, new_pattern): 26 | return re.sub(pattern, new_pattern, string) 27 | 28 | 29 | def prefix_match(string, prefix=None): 30 | """Check if the string matches the given prefix""" 31 | if prefix is None or len(prefix) == 0: 32 | return True 33 | return re.match(f"({prefix})+(.*?)", string) is not None 34 | 35 | 36 | def is_regex(s): 37 | try: 38 | re.compile(s) 39 | return True 40 | except: 41 | return False 42 | 43 | 44 | def float_str(num, precision): 45 | format_str = "%.{0}f".format(precision) 46 | return format_str % num 47 | 48 | 49 | def num_to_str(num, unit=None, precision=2, number_only=False, auto_select_unit=False): 50 | unit_list = ["K", "M", "G", "T", "P"] 51 | if auto_select_unit and unit is None: 52 | for i, tmp in enumerate(unit_list): 53 | unit_num = 1024 ** (i + 1) 54 | if num < unit_num: 55 | break 56 | unit = tmp 57 | if unit is not None: 58 | unit_num = 1024 ** (unit_list.index(unit) + 1) 59 | num = num * 1.0 / unit_num 60 | else: 61 | unit = "" 62 | if number_only: 63 | return num 64 | else: 65 | return float_str(num, precision) + unit 66 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/data/type_utils.py: -------------------------------------------------------------------------------- 1 | from collections.abc import Sequence 2 | from numbers import Number 3 | import numpy as np 4 | 5 | 6 | """ For python basic type """ 7 | 8 | 9 | def is_null(item): 10 | return item is None 11 | 12 | 13 | def is_not_null(item): 14 | return item is not None 15 | 16 | 17 | def is_slice(item): 18 | return isinstance(item, slice) 19 | 20 | 21 | def is_str(item): 22 | return isinstance(item, str) 23 | 24 | 25 | def is_dict(item): 26 | return isinstance(item, dict) 27 | 28 | 29 | def is_num(item): 30 | return isinstance(item, Number) 31 | 32 | 33 | def is_integer(item): 34 | return isinstance(item, (int, np.integer)) 35 | 36 | 37 | def is_type(item): 38 | return isinstance(item, type) 39 | 40 | 41 | def is_seq_of(seq, expected_type=None, seq_type=None): 42 | if seq_type is None: 43 | exp_seq_type = Sequence 44 | else: 45 | assert isinstance(seq_type, type) 46 | exp_seq_type = seq_type 47 | if not isinstance(seq, exp_seq_type): 48 | return False 49 | if expected_type: 50 | for item in seq: 51 | if not isinstance(item, expected_type): 52 | return False 53 | return True 54 | 55 | 56 | def is_list_of(seq, expected_type=None): 57 | return is_seq_of(seq, expected_type, seq_type=list) 58 | 59 | 60 | def is_tuple_of(seq, expected_type=None): 61 | return is_seq_of(seq, expected_type, seq_type=tuple) 62 | 63 | 64 | def is_iterable(item): 65 | return isinstance(item, (dict, tuple, list)) 66 | 67 | 68 | """ For numpy and torch type """ 69 | 70 | 71 | def get_dtype(item): 72 | if isinstance(item, (list, tuple)): 73 | item = item[0] 74 | if hasattr(item, "dtype"): 75 | return str(item.dtype).split(".")[-1] 76 | elif isinstance(item, (int, float, bytes, str)): 77 | return type(item) 78 | else: 79 | return None 80 | 81 | 82 | def is_np(item): 83 | return isinstance(item, np.ndarray) or is_num(item) 84 | 85 | 86 | def is_np_arr(item): 87 | return isinstance(item, np.ndarray) 88 | 89 | 90 | def is_torch(item): 91 | import torch 92 | 93 | return isinstance(item, torch.Tensor) 94 | 95 | 96 | def is_torch_distribution(item): 97 | import torch 98 | 99 | return isinstance(item, torch.distributions.Distribution) 100 | 101 | 102 | def is_arr(item, arr_type=None): 103 | if is_num(item): 104 | return False 105 | if arr_type is not None: 106 | assert arr_type in ["np", "torch"] 107 | return eval(f"is_{arr_type}")(item) 108 | elif is_np(item): 109 | return True 110 | else: 111 | # Torch as the last option to reduce memory usage 112 | return is_torch(item) 113 | 114 | 115 | """ For HDF5 type """ 116 | 117 | 118 | def is_h5(item): 119 | from h5py import File, Group, Dataset 120 | 121 | return isinstance(item, (File, Group, Dataset)) 122 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/data/wrappers.py: -------------------------------------------------------------------------------- 1 | from functools import wraps 2 | 3 | from numbers import Number 4 | from .converter import to_np, to_torch, to_array 5 | 6 | 7 | def seq_to_np(to_arr=True): 8 | def decorator(func): 9 | @wraps(func) 10 | def wrapper(item, *args, **kwargs): 11 | if isinstance(item, (list, tuple, Number)): 12 | item = to_np(item) 13 | if to_arr: 14 | item = to_array(item) 15 | return func(item, *args, **kwargs) 16 | 17 | return wrapper 18 | 19 | return decorator 20 | 21 | 22 | def check_consistent(keys, dtypes): 23 | if dtypes: 24 | dtypes = (dtypes,) if not isinstance(dtypes, (list, tuple)) else dtypes 25 | if keys: 26 | keys = (keys,) if not isinstance(keys, (list, tuple)) else keys 27 | if dtypes: 28 | assert len(keys) == len(dtypes) 29 | else: 30 | if dtypes: 31 | assert len(dtypes) == 1 32 | 33 | 34 | def apply_func(func, x): 35 | if isinstance(x, (list, tuple, set)): 36 | return type(x)(map(func, x)) 37 | elif isinstance(x, dict): 38 | for k in x: 39 | x[k] = func(x[k]) 40 | return x 41 | else: 42 | return func(x) 43 | 44 | 45 | def change_dtype(x, keys=None, dtypes=None, np=False): 46 | if dtypes is None: 47 | return x 48 | processor = to_np if np else to_torch 49 | if not isinstance(dtypes, (list, tuple)): 50 | dtypes = [dtypes] 51 | 52 | if not isinstance(x, (tuple, list, dict)) or keys is None: 53 | assert len(dtypes) == 1 54 | return processor(x, dtypes[0]) 55 | 56 | if not isinstance(keys, (list, tuple)): 57 | keys = [keys] 58 | # key and dtypes are list or tuple, dtypes is a list, x is a list, tuple or dict 59 | 60 | ret = list(x) if isinstance(x, (list, tuple)) else x 61 | if len(dtypes) == 1: 62 | dtypes = [dtypes[0] for i in range(len(keys))] 63 | for k, dtype in enumerate(keys, dtypes): 64 | ret[k] = processor(ret[k], dtype) 65 | return type(x)(ret) 66 | 67 | 68 | def process_output(keys=None, dtypes=None, np=True): 69 | check_consistent(keys, dtypes) 70 | 71 | def decorator(func): 72 | wraps(func) 73 | 74 | def wrapper(*args, **kwargs): 75 | ret = func(*args, **kwargs) 76 | return change_dtype(ret, keys, dtypes, np) 77 | 78 | return wrapper 79 | 80 | return decorator 81 | 82 | 83 | def process_input(keys=None, dtypes=None, np=True): 84 | check_consistent(keys, dtypes) 85 | 86 | def decorator(func): 87 | wraps(func) 88 | 89 | def wrapper(*args, **kwargs): 90 | args = list(args) 91 | kwargs = dict(kwargs) 92 | args = change_dtype(args, keys, dtypes, np) 93 | kwargs = change_dtype(kwargs, keys, dtypes, np) 94 | return func(*args, **kwargs) 95 | 96 | return wrapper 97 | 98 | return decorator 99 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/file/__init__.py: -------------------------------------------------------------------------------- 1 | from .file_client import BaseStorageBackend, FileClient 2 | 3 | # (is_saved_with_pandas, , load_h5_as_dict_array, 4 | # load_h5s_as_list_dict_array, convert_h5_trajectory_to_pandas, DataEpisode, 5 | # generate_chunked_h5_replay) 6 | from .hash_utils import md5sum, check_md5sum 7 | from .lmdb_utils import LMDBFile 8 | 9 | # from .pandas_utils import (convert_hdf_with_pickle_4, load_hdf, hdf_to_dict_list, merge_hdf_trajectory, 10 | # save_hdf_trajectory, concat_hdf_trajectory, try_to_open_hdf_trajectory) 11 | from .serialization import * 12 | from .zip_utils import extract_files 13 | from .record_utils import ( 14 | output_record, 15 | generate_index_from_record, 16 | shuffle_merge_records, 17 | shuffle_reocrd, 18 | get_index_filenames, 19 | read_record, 20 | merge_h5_trajectory, 21 | convert_h5_trajectory_to_record, 22 | load_items_from_record, 23 | load_record_indices, 24 | train_test_split, 25 | convert_h5_trajectories_to_shard, 26 | do_train_test_split, 27 | ) 28 | 29 | from .hdf5_utils import load_hdf5, dump_hdf5 30 | from .cache_utils import get_total_size, FileCache, is_h5_traj, decode_items 31 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/file/hash_utils.py: -------------------------------------------------------------------------------- 1 | import hashlib, numpy as np, struct 2 | 3 | 4 | def md5sum(filename, block_size=None): 5 | if block_size is None: 6 | block_size = 65536 7 | hash_res = hashlib.md5() 8 | with open(filename, "rb") as f: 9 | for block in iter(lambda: f.read(block_size), b""): 10 | hash_res.update(block) 11 | return hash_res.hexdigest() 12 | 13 | 14 | def check_md5sum(filename, md5, block_size=None): 15 | if not (isinstance(md5, str) and len(md5) == 32): 16 | raise ValueError(f"MD5 must be 32 chars: {md5}") 17 | md5_actual = md5sum(filename, block_size=block_size) 18 | if md5_actual == md5: 19 | return True 20 | else: 21 | print(f"MD5 does not match!: {filename} has md5 {md5_actual}, target md5 is {md5}") 22 | return False 23 | 24 | 25 | def masked_crc(data: bytes) -> bytes: 26 | try: 27 | from crc32c import crc32c 28 | except ImportError: 29 | print("Cannot import crc32c, please install it!") 30 | exit(0) 31 | """CRC checksum.""" 32 | mask = 0xA282EAD8 33 | crc = crc32c(data) 34 | masked = ((crc >> 15) | (crc << 17)) + mask 35 | masked = np.uint32(masked & np.iinfo(np.uint32).max) 36 | masked_bytes = struct.pack("= 0 and cnt >= max_num: 20 | break 21 | item_list.append(prefix + line.rstrip("\n")) 22 | cnt += 1 23 | return item_list 24 | 25 | 26 | def dict_from_file(filename, key_type=str, offset=0, max_num=-1): 27 | mapping = {} 28 | cnt = 0 29 | with open(filename, "r") as f: 30 | for _ in range(offset): 31 | f.readline() 32 | for line in f: 33 | if max_num >= 0 and cnt >= max_num: 34 | break 35 | items = line.rstrip("\n").split() 36 | assert len(items) >= 2 37 | key = key_type(items[0]) 38 | val = items[1:] if len(items) > 2 else items[1] 39 | mapping[key] = val 40 | cnt += 1 41 | return mapping 42 | 43 | 44 | def dict_to_csv_table(x): 45 | ret = [] 46 | for key in x.keys(): 47 | ret.append([key, x[key]]) 48 | return ret 49 | 50 | 51 | def csv_table_to_dict(x): 52 | for y in x: 53 | assert len(y) == 2 54 | ret = {} 55 | for y in x: 56 | ret[y[0]] = y[1] 57 | return ret 58 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/image/__init__.py: -------------------------------------------------------------------------------- 1 | from .colorspace import ( 2 | bgr2gray, 3 | bgr2hls, 4 | bgr2hsv, 5 | bgr2rgb, 6 | bgr2ycbcr, 7 | gray2bgr, 8 | gray2rgb, 9 | hls2bgr, 10 | hsv2bgr, 11 | imconvert, 12 | rgb2bgr, 13 | rgb2gray, 14 | rgb2ycbcr, 15 | ycbcr2bgr, 16 | ycbcr2rgb, 17 | ) 18 | from .geometric import ( 19 | imcrop, 20 | imflip, 21 | imflip_, 22 | impad, 23 | impad_to_multiple, 24 | imrescale, 25 | imresize, 26 | imresize_like, 27 | imrotate, 28 | imshear, 29 | imtranslate, 30 | rescale_size, 31 | ) 32 | from .io import imfrombytes, imread, imwrite, supported_backends, use_backend, imencode, imdecode 33 | from .misc import tensor2imgs 34 | from .photometric import ( 35 | adjust_brightness, 36 | adjust_color, 37 | adjust_contrast, 38 | clahe, 39 | imdenormalize, 40 | imequalize, 41 | iminvert, 42 | imnormalize, 43 | imnormalize_, 44 | lut_transform, 45 | posterize, 46 | solarize, 47 | ) 48 | from .video_utils import concat_videos, put_names_on_image, grid_images, video_to_frames, images_to_video 49 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/image/misc.py: -------------------------------------------------------------------------------- 1 | from typing import List, Dict 2 | import numpy as np, cv2 3 | from .photometric import imdenormalize 4 | 5 | try: 6 | import torch 7 | except ImportError: 8 | torch = None 9 | 10 | 11 | def tensor2imgs(tensor, mean=(0, 0, 0), std=(1, 1, 1), to_rgb=True): 12 | """Convert tensor to 3-channel images. 13 | Args: 14 | tensor (torch.Tensor): Tensor that contains multiple images, shape (N, C, H, W). 15 | mean (tuple[float], optional): Mean of images. Defaults to (0, 0, 0). 16 | std (tuple[float], optional): Standard deviation of images. Defaults to (1, 1, 1). 17 | to_rgb (bool, optional): Whether the tensor was converted to RGB format in the first place. 18 | If so, convert it back to BGR. Defaults to True. 19 | Returns: 20 | list[np.ndarray]: A list that contains multiple images. 21 | """ 22 | 23 | if torch is None: 24 | raise RuntimeError("pytorch is not installed") 25 | assert torch.is_tensor(tensor) and tensor.ndim == 4 26 | assert len(mean) == 3 27 | assert len(std) == 3 28 | 29 | num_imgs = tensor.size(0) 30 | mean = np.array(mean, dtype=np.float32) 31 | std = np.array(std, dtype=np.float32) 32 | imgs = [] 33 | for img_id in range(num_imgs): 34 | img = tensor[img_id, ...].cpu().numpy().transpose(1, 2, 0) 35 | img = imdenormalize(img, mean, std, to_bgr=to_rgb).astype(np.uint8) 36 | imgs.append(np.ascontiguousarray(img)) 37 | return imgs 38 | 39 | 40 | def put_text_on_image(image: np.ndarray, lines: List[str]): 41 | assert image.dtype == np.uint8, image.dtype 42 | image = image.copy() 43 | 44 | font_size = 0.5 45 | font_thickness = 1 46 | font = cv2.FONT_HERSHEY_SIMPLEX 47 | 48 | y = 0 49 | for line in lines: 50 | if isinstance(line, np.ndarray): 51 | continue 52 | text_size = cv2.getTextSize(line, font, font_size, font_thickness)[0] 53 | y += text_size[1] + 10 54 | x = 10 55 | cv2.putText( 56 | image, 57 | line, 58 | (x, y), 59 | font, 60 | font_size, 61 | (0, 255, 0), 62 | font_thickness, 63 | lineType=cv2.LINE_AA, 64 | ) 65 | return image 66 | 67 | 68 | def append_text_to_image(image: np.ndarray, lines: List[str]): 69 | r"""Appends text left to an image of size (height, width, channels). 70 | The returned image has white text on a black background. 71 | 72 | Args: 73 | image: the image to put text 74 | text: a string to display 75 | 76 | Returns: 77 | A new image with text inserted left to the input image 78 | 79 | See also: 80 | habitat.utils.visualization.utils 81 | """ 82 | # h, w, c = image.shape 83 | font_size = 0.5 84 | font_thickness = 1 85 | font = cv2.FONT_HERSHEY_SIMPLEX 86 | blank_image = np.zeros(image.shape, dtype=np.uint8) 87 | 88 | y = 0 89 | for line in lines: 90 | if isinstance(line, np.ndarray): 91 | continue 92 | text_size = cv2.getTextSize(line, font, font_size, font_thickness)[0] 93 | y += text_size[1] + 10 94 | x = 10 95 | cv2.putText( 96 | blank_image, 97 | line, 98 | (x, y), 99 | font, 100 | font_size, 101 | (255, 255, 255), 102 | font_thickness, 103 | lineType=cv2.LINE_AA, 104 | ) 105 | # text_image = blank_image[0 : y + 10, 0:w] 106 | # final = np.concatenate((image, text_image), axis=0) 107 | final = np.concatenate((blank_image, image), axis=1) 108 | return final 109 | 110 | 111 | def put_info_on_image(image, info: Dict[str, float], extras=None, overlay=True): 112 | lines = [f"{k}: {v:.3f}" for k, v in info.items() if not isinstance(v, np.ndarray)] 113 | if extras is not None: 114 | lines.extend(extras) 115 | if overlay: 116 | return put_text_on_image(image, lines) 117 | else: 118 | return append_text_to_image(image, lines) 119 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/image/video_utils.py: -------------------------------------------------------------------------------- 1 | import cv2, numpy as np, imageio 2 | import os.path as osp 3 | 4 | 5 | def video_to_frames(filename, output_dir=None): 6 | video = cv2.VideoCapture(filename) 7 | images = [] 8 | count = 0 9 | success = True 10 | while success: 11 | success, image = video.read() 12 | if success: 13 | if output_dir is None: 14 | images.append(image) 15 | else: 16 | cv2.imwrite(osp.join(output_dir, f"frame_{count}.jpg"), image) 17 | count += 1 18 | return np.stack(images, axis=0)[..., ::-1] 19 | 20 | 21 | def grid_images(images, max_col_num=4): 22 | total_num = ((len(images) + max_col_num - 1) // max_col_num) * max_col_num 23 | images = list(images) + [np.zeros_like(images[0]) for i in range(total_num - len(images))] 24 | # print(len(images), max_col_num, total_num) 25 | ret = [] 26 | for i in range(total_num // max_col_num): 27 | row_i = [] 28 | for j in range(max_col_num): 29 | k = i * max_col_num + j 30 | row_i.append(images[k]) 31 | ret.append(np.concatenate(row_i, axis=2)) 32 | ret = np.concatenate(ret, axis=1) 33 | return ret 34 | 35 | 36 | def put_names_on_image(images, names): 37 | font = cv2.FONT_HERSHEY_SIMPLEX 38 | for image, name in zip(images, names): 39 | for i in range(len(image)): 40 | # print(image.shape, name, font) 41 | cv2.putText(image[i], name, (10, image.shape[1] - 10), font, 1, (0, 255, 0), 1, cv2.LINE_AA) 42 | # import matplotlib.pyplot as plt 43 | # plt.imshow(image[i]) 44 | # plt.show() 45 | # cv2.imwrite(path + 'pillar_text.jpg', im) 46 | 47 | 48 | def concat_videos(filenames, output_filename, names=None, max_col_num=4, fps=10): 49 | images = [video_to_frames(filename) for filename in filenames] 50 | num = np.max([image.shape[0] for image in images]) 51 | images = [ 52 | np.concatenate( 53 | [ 54 | image, 55 | ] 56 | + [image[-1:] for i in range(num - image.shape[0])], 57 | axis=0, 58 | ) 59 | for image in images 60 | ] 61 | 62 | if names is not None: 63 | assert len(names) == len(images) 64 | put_names_on_image(images, names) 65 | 66 | images = grid_images(images, max_col_num) 67 | images_to_video(images, output_filename, fps=fps) 68 | 69 | 70 | def images_to_video(images, filebane, fps=10): 71 | writer = imageio.get_writer(filebane, fps=fps) 72 | for im in images: 73 | writer.append_data(im) 74 | writer.close() 75 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/lib3d/__init__.py: -------------------------------------------------------------------------------- 1 | from .o3d_utils import ( 2 | to_o3d, 3 | np2mesh, 4 | merge_mesh, 5 | np2pcd, 6 | one_point_vis, 7 | create_aabb, 8 | create_obb, 9 | create_aabb_from_pcd, 10 | create_obb_from_pcd, 11 | create_aabb_from_mesh, 12 | create_obb_from_mesh, 13 | ) 14 | from .trimesh_utils import to_trimesh 15 | from .utils import convex_hull, angle, check_coplanar, apply_pose, mesh_to_pcd 16 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/lib3d/trimesh_utils.py: -------------------------------------------------------------------------------- 1 | import trimesh, numpy as np, open3d as o3d 2 | from maniskill2_learn.utils.data import is_pcd 3 | 4 | 5 | def to_trimesh(x): 6 | if is_trimesh(x): 7 | return x 8 | elif isinstance(x, np.ndarray): 9 | assert is_pcd(x) 10 | return trimesh.points.PointCloud(x) 11 | elif isinstance(x, o3d.geometry.TriangleMesh): 12 | vertices = np.asarray(x.vertices) 13 | faces = np.asarray(x.triangles) 14 | return trimesh.Trimesh(vertices=vertices, faces=faces) 15 | elif isinstance(x, o3d.geometry.PointCloud): 16 | points = np.asarray(x.points) 17 | return trimesh.points.PointCloud(vertices=points) 18 | else: 19 | print(type(x)) 20 | raise NotImplementedError() 21 | 22 | 23 | def is_trimesh(x): 24 | return isinstance(x, (trimesh.Trimesh, trimesh.points.PointCloud)) 25 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/lib3d/utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np, trimesh, sklearn 2 | from sapien.core import Pose 3 | from .o3d_utils import is_o3d, to_o3d, np2mesh 4 | from .trimesh_utils import is_trimesh, to_trimesh 5 | from maniskill2_learn.utils.data import normalize, to_gc, to_nc 6 | 7 | 8 | def convex_hull(x, o3d=True): 9 | x = to_trimesh(x) 10 | x = trimesh.convex.convex_hull(x) 11 | if o3d: 12 | x = to_o3d(x) 13 | return x 14 | 15 | 16 | def angle(x1, x2): 17 | if isinstance(x1, np.ndarray): 18 | x1 = normalize(x1) 19 | x2 = normalize(x2) 20 | return np.arccos(np.dot(x1, x2)) 21 | 22 | 23 | def mesh_to_pcd(x, num_points, o3d=True): 24 | x = to_trimesh(x) 25 | x = x.sample(num_points) 26 | if o3d: 27 | x = to_o3d(x) 28 | return x 29 | 30 | 31 | def apply_pose(pose, x): 32 | if x is None: 33 | return x 34 | if isinstance(pose, np.ndarray): 35 | pose = Pose.from_transformation_matrix(pose) 36 | assert isinstance(pose, Pose) 37 | if isinstance(x, Pose): 38 | return pose * x 39 | elif isinstance(x, np.ndarray): 40 | return to_nc(to_gc(x, dim=3) @ pose.to_transformation_matrix().T, dim=3) 41 | elif is_trimesh(x) or is_o3d(x): 42 | sign = is_o3d(x) 43 | x = to_trimesh(x) 44 | if isinstance(x, trimesh.Trimesh): 45 | vertices = x.vertices 46 | faces = x.faces 47 | vertices = apply_pose(pose, vertices) 48 | x = trimesh.Trimesh(vertices=vertices, faces=faces) 49 | elif isinstance(x, trimesh.points.PointCloud): 50 | vertices = x.vertices 51 | vertices = apply_pose(pose, vertices) 52 | x = trimesh.points.PointCloud(vertices=vertices) 53 | if sign: 54 | x = to_o3d(x) 55 | return x 56 | else: 57 | print(x, type(x)) 58 | raise NotImplementedError("") 59 | 60 | 61 | def check_coplanar(vertices): 62 | pca = sklearn.decomposition.PCA(n_components=3) 63 | pca.fit(vertices) 64 | return pca.singular_values_[-1] < 1e-3 -------------------------------------------------------------------------------- /maniskill2_learn/utils/math/__init__.py: -------------------------------------------------------------------------------- 1 | from .running_stats import MovingAverage, RunningMeanStd 2 | from .trunc_normal import trunc_normal 3 | from .split_array import split_num 4 | from .counting import EveryNSteps 5 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/math/counting.py: -------------------------------------------------------------------------------- 1 | class EveryNSteps: 2 | def __init__(self, interval=None): 3 | self.interval = interval 4 | self.next_value = interval 5 | 6 | def reset(self): 7 | self.next_value = self.interval 8 | 9 | def check(self, x): 10 | if self.interval is None: 11 | return False 12 | sign = False 13 | while x >= self.next_value: 14 | self.next_value += self.interval 15 | sign = True 16 | return sign 17 | 18 | def standard(self, x): 19 | return int(x // self.interval) * self.interval 20 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/math/running_stats.py: -------------------------------------------------------------------------------- 1 | from collections import deque 2 | 3 | import numpy as np 4 | 5 | 6 | class MovingAverage: 7 | def __init__(self, size=10): 8 | self.size = size 9 | self.memory = deque(maxlen=size) 10 | 11 | def add(self, x): 12 | if np.isscalar(x): 13 | x = [x] 14 | for i in x: 15 | if i not in [np.inf, np.nan, -np.inf]: 16 | self.memory.append(i) 17 | 18 | def mean(self): 19 | return np.mean(self.memory) if len(self.memory) > 0 else 0 20 | 21 | def std(self): 22 | return np.std(self.memory) if len(self.memory) > 0 else 0 23 | 24 | 25 | class RunningMeanStd: 26 | """ 27 | Calculates the running mean and std of a data stream: 28 | https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm 29 | """ 30 | 31 | def __init__(self, mean=0, std=0): 32 | self.mean, self.var = mean, std**2 33 | self.n = 0 34 | 35 | @property 36 | def std(self): 37 | return np.maximum(np.sqrt(self.var), 1e-8) 38 | 39 | def add(self, x): 40 | """ 41 | Parallel algorithm: a is the history, b is the new one. 42 | m = n * var 43 | """ 44 | new_mean, new_var = np.mean(x, axis=0), np.var(x, axis=0) 45 | new_n = len(x) 46 | 47 | n = self.n + new_n 48 | delta = new_mean - self.mean 49 | m2_a = self.n * self.var 50 | m2_b = new_n * new_var 51 | m2 = m2_a + m2_b + delta**2 * (self.n * new_n) / n 52 | self.mean = self.mean + delta * new_n / n 53 | self.var = m2 / n 54 | self.n = n 55 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/math/split_array.py: -------------------------------------------------------------------------------- 1 | def split_num(num, n): 2 | """ 3 | Divide num into m=min(n, num) elements x_1, ...., x_n, where x_1, ..., x_n >= 1 and max_{i,j} |x_i - x_j| <= 1 4 | """ 5 | n = min(num, n) 6 | min_steps = num // n 7 | splits = [] 8 | for i in range(n): 9 | if i < num - min_steps * n: 10 | splits.append(min_steps + 1) 11 | else: 12 | splits.append(min_steps) 13 | assert sum(splits) == num 14 | return n, splits 15 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/math/trunc_normal.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.stats import norm 3 | 4 | 5 | def trunc_normal(shape, a=-2, b=2): 6 | """ 7 | :param shape: The shape of the trunc normal 8 | :param a, b: Sample between [a, b] with i.i.d. normal distribution 9 | :return: samples 10 | Gaussian density of N(mu, sigma^2): exp(-((x - mu) / sigma)^2 / 2) / (sqrt(2 * pi) * sigma) 11 | """ 12 | a_cdf = norm.cdf(a) 13 | b_cdf = norm.cdf(b) 14 | p = a_cdf + (b_cdf - a_cdf) * np.random.rand(*shape) 15 | return np.clip(norm.ppf(p), a, b) 16 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/meta/__init__.py: -------------------------------------------------------------------------------- 1 | from .config import ConfigDict, Config, DictAction, merge_a_to_b 2 | from .collect_env import collect_env, log_meta_info, get_meta_info 3 | from .logger import get_logger, print_log, flush_print, get_logger_name, TqdmToLogger, flush_logger 4 | from .magic_utils import * 5 | from .module_utils import import_modules_from_strings, check_prerequisites, requires_package, requires_executable, deprecated_api_warning 6 | from .path_utils import ( 7 | is_filepath, 8 | fopen, 9 | check_files_exist, 10 | mkdir_or_exist, 11 | parse_files, 12 | symlink, 13 | scandir, 14 | find_vcs_root, 15 | get_filename, 16 | get_filename_suffix, 17 | copy_folder, 18 | copy_folders, 19 | add_suffix_to_filename, 20 | get_dirname, 21 | to_abspath, 22 | replace_suffix, 23 | ) 24 | from .process_utils import get_total_memory, get_memory_list, get_subprocess_ids, get_memory_dict 25 | from .progressbar import ProgressBar, track_progress, track_iter_progress, track_parallel_progress 26 | from .random_utils import RandomWrapper, get_random_generator, set_random_seed, random_id_generator 27 | from .registry import Registry, build_from_cfg 28 | from .timer import Timer, TimerError, check_time, get_time_stamp, td_format, get_today 29 | from .version_utils import digit_version 30 | from .env_var import add_env_var, add_dist_var, get_world_rank, get_world_size, is_debug_mode, get_dist_info 31 | from .parallel_runner import Worker 32 | from .network import is_port_in_use 33 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/meta/env_var.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | 4 | def add_env_var(): 5 | default_values = {"NUMEXPR_MAX_THREADS": "1", "MKL_NUM_THREADS": "1", "OMP_NUM_THREADS": "1", "CUDA_DEVICE_ORDER": "PCI_BUS_ID", "DISPLAY": "0", "MUJOCO_GL": "egl"} 6 | for key, value in default_values.items(): 7 | os.environ[key] = os.environ.get(key, value) 8 | 9 | 10 | def add_dist_var(rank, world_size): 11 | os.environ["PYRL_RANK"] = f"{rank}" 12 | os.environ["PYRL_WORLD_SIZE"] = f"{world_size}" 13 | os.environ["MASTER_ADDR"] = "localhost" 14 | 15 | def find_free_port(port): 16 | from .network import is_port_in_use 17 | 18 | while is_port_in_use(port): 19 | port += 1 20 | return port 21 | 22 | os.environ["MASTER_PORT"] = str(find_free_port(12355)) 23 | os.environ["PYRL_TCP_PORT"] = str(find_free_port(15015)) 24 | 25 | 26 | def get_world_rank(): 27 | if "PYRL_RANK" not in os.environ: 28 | return 0 29 | return eval(os.environ["PYRL_RANK"]) 30 | 31 | 32 | def get_world_size(): 33 | if "PYRL_WORLD_SIZE" not in os.environ: 34 | return 1 35 | return eval(os.environ["PYRL_WORLD_SIZE"]) 36 | 37 | 38 | def get_dist_info(): 39 | return get_world_rank(), get_world_size() 40 | 41 | 42 | def is_debug_mode(): 43 | if "PYRL_DEBUG" not in os.environ: 44 | return 0 45 | return eval(os.environ["PYRL_DEBUG"]) 46 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/meta/magic_utils.py: -------------------------------------------------------------------------------- 1 | # https://github.com/alexmojaki/sorcery 2 | from sorcery import assigned_names, unpack_keys, unpack_attrs, dict_of, print_args, call_with_name, delegate_to_attr, maybe, select_from 3 | 4 | # https://github.com/gruns/icecream 5 | # from icecream import ic 6 | 7 | 8 | def colored_print(output_string, level="info", logger=print): 9 | from termcolor import colored 10 | import sys 11 | 12 | if level.lower() in ["warning", "error"]: 13 | level = colored(level.upper(), "red") 14 | output_string = colored(output_string, "cyan") 15 | logger(f"{level}: {output_string}") 16 | else: 17 | logger(output_string) 18 | 19 | 20 | def empty_print(*args, **kwargs): 21 | pass 22 | 23 | 24 | def custom_assert(pause, output_string, logger=None): 25 | if logger is not None: 26 | logger = logger.log 27 | else: 28 | logger = print 29 | import sys 30 | 31 | if not pause: 32 | from termcolor import colored 33 | file_name = colored(sys._getframe().f_code.co_filename, "red") 34 | line_number = colored(sys._getframe().f_back.f_lineno, "cyan") 35 | output_string = colored(output_string, "red") 36 | logger(f"Assert Error at {file_name}, line {line_number}") 37 | logger(f"Output: {output_string}") 38 | 39 | 40 | class SlicePrinter: 41 | def __getitem__(self, index): 42 | print(index) 43 | 44 | slice_printer = SlicePrinter() 45 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/meta/network.py: -------------------------------------------------------------------------------- 1 | def is_port_in_use(port: int) -> bool: 2 | import socket 3 | 4 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: 5 | return s.connect_ex(("localhost", port)) == 0 6 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/meta/process_utils.py: -------------------------------------------------------------------------------- 1 | import psutil, os 2 | 3 | 4 | def format_memory_str(x, unit, number_only=False): 5 | unit_list = ["K", "M", "G", "T"] 6 | assert unit in unit_list 7 | unit_num = 1024 ** (unit_list.index(unit) + 1) 8 | if number_only: 9 | return x * 1.0 / unit_num 10 | else: 11 | return f"{x * 1.0 / unit_num:.2f}{unit}" 12 | 13 | 14 | def get_total_memory(unit="G", number_only=False, init_pid=None): 15 | from maniskill2_learn.utils.data import num_to_str 16 | 17 | if init_pid is None: 18 | init_pid = os.getpid() 19 | process = psutil.Process(init_pid) 20 | ret = process.memory_full_info().uss 21 | for proc in process.children(): 22 | process_info = proc.memory_full_info() 23 | ret += process_info.uss 24 | return num_to_str(ret, unit, number_only=number_only) 25 | 26 | 27 | def get_memory_list(unit="G", number_only=False, init_pid=None): 28 | from maniskill2_learn.utils.data import num_to_str 29 | 30 | if init_pid is None: 31 | init_pid = os.getpid() 32 | process = psutil.Process(init_pid) 33 | ret = [ 34 | num_to_str(process.memory_full_info().uss, unit, number_only=number_only), 35 | ] 36 | for proc in process.children(): 37 | process_info = proc.memory_full_info() 38 | ret.append(num_to_str(process_info.uss, unit, number_only=number_only)) 39 | return ret 40 | 41 | 42 | def get_memory_dict(unit="G", number_only=False, init_pid=None): 43 | from maniskill2_learn.utils.data import num_to_str 44 | 45 | if init_pid is None: 46 | init_pid = os.getpid() 47 | process = psutil.Process(init_pid) 48 | ret = {init_pid: num_to_str(process.memory_full_info().uss, unit, number_only=number_only)} 49 | for i, proc in enumerate(process.children()): 50 | process_info = proc.memory_full_info() 51 | ret[proc.pid] = num_to_str(process_info.uss, unit, number_only=number_only) 52 | return ret 53 | 54 | 55 | def get_subprocess_ids(init_pid=None): 56 | if init_pid is None: 57 | init_pid = os.getpid() 58 | ret = [init_pid] 59 | process = psutil.Process(os.getpid()) 60 | for proc in process.children(): 61 | ret.append(proc.pid) 62 | return ret 63 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/meta/random_utils.py: -------------------------------------------------------------------------------- 1 | import random, string, time, numpy as np 2 | 3 | 4 | def random_id_generator(size=6, chars=string.ascii_uppercase): 5 | timestamp = time.strftime("%Y%m%d_%H%M%S", time.localtime()) 6 | return timestamp + "." + "".join(random.choice(chars) for _ in range(size)) 7 | 8 | 9 | class RandomWrapper(object): 10 | def __init__(self, seed): 11 | self.seed = seed 12 | self.state = None 13 | 14 | def __enter__(self): 15 | self.state = np.random.get_state() 16 | np.random.seed(self.seed) 17 | return self.state 18 | 19 | def __exit__(self, exc_type, exc_val, exc_tb): 20 | np.random.set_state(self.state) 21 | 22 | 23 | def get_random_generator(seed): 24 | return np.random.RandomState(seed) 25 | 26 | 27 | def set_random_seed(seed): 28 | if seed is not None: 29 | random.seed(seed) 30 | np.random.seed(seed) 31 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/meta/timer.py: -------------------------------------------------------------------------------- 1 | from time import strftime, localtime 2 | from datetime import date 3 | import time 4 | 5 | 6 | def td_format(td_object): 7 | seconds = int(td_object.total_seconds()) 8 | periods = [("y", 60 * 60 * 24 * 365), ("m", 60 * 60 * 24 * 30), ("d", 60 * 60 * 24), ("h", 60 * 60), ("m", 60), ("s", 1)] 9 | ret = "" 10 | for period_name, period_seconds in periods: 11 | if seconds > period_seconds: 12 | period_value, seconds = divmod(seconds, period_seconds) 13 | ret += f"{period_value}{period_name}" 14 | return ret 15 | 16 | 17 | class TimerError(Exception): 18 | def __init__(self, message): 19 | self.message = message 20 | super(TimerError, self).__init__(message) 21 | 22 | 23 | class Timer: 24 | """A flexible Timer class. 25 | 26 | :Example: 27 | 28 | >>> import time 29 | >>> import maniskill2_learn 30 | >>> with maniskill2_learn.Timer(): 31 | >>> # simulate a code block that will run for 1s 32 | >>> time.sleep(1) 33 | 1.000 34 | >>> with maniskill2_learn.Timer(print_tmpl='it takes {:.1f} seconds'): 35 | >>> # simulate a code block that will run for 1s 36 | >>> time.sleep(1) 37 | it takes 1.0 seconds 38 | >>> timer = maniskill2_learn.Timer() 39 | >>> time.sleep(0.5) 40 | >>> print(timer.since_start()) 41 | 0.500 42 | >>> time.sleep(0.5) 43 | >>> print(timer.since_last_check()) 44 | 0.500 45 | >>> print(timer.since_start()) 46 | 1.000 47 | """ 48 | 49 | def __init__(self, start=True, print_tmpl="{:.3f}"): 50 | self._is_running = False 51 | self.print_tmpl = print_tmpl 52 | if start: 53 | self.start() 54 | 55 | @property 56 | def is_running(self): 57 | """bool: indicate whether the timer is running""" 58 | return self._is_running 59 | 60 | def __enter__(self): 61 | self.start() 62 | return self 63 | 64 | def __exit__(self, type, value, traceback): 65 | print(self.print_tmpl.format(self.since_last_check())) 66 | self._is_running = False 67 | 68 | def start(self): 69 | """Start the timer.""" 70 | if not self._is_running: 71 | self._t_start = time.time() 72 | self._is_running = True 73 | self._t_last = time.time() 74 | 75 | def since_start(self): 76 | """Total time since the timer is started. 77 | Returns (float): Time in seconds. 78 | """ 79 | if not self._is_running: 80 | raise TimerError("timer is not running") 81 | self._t_last = time.time() 82 | return self._t_last - self._t_start 83 | 84 | def since_last_check(self): 85 | """Time since the last checking. 86 | Either :func:`since_start` or :func:`since_last_check` is a checking operation. 87 | Returns (float): Time in seconds. 88 | """ 89 | if not self._is_running: 90 | raise TimerError("timer is not running") 91 | dur = time.time() - self._t_last 92 | self._t_last = time.time() 93 | return dur 94 | 95 | 96 | global_timers = {} 97 | 98 | 99 | def check_time(timer_id): 100 | """Add check points in a single line. 101 | This method is suitable for running a task on a list of items. A timer will be registered when the method is called 102 | for the first time. 103 | 104 | :Example: 105 | 106 | >>> import time 107 | >>> import maniskill2_learn 108 | >>> for i in range(1, 6): 109 | >>> # simulate a code block 110 | >>> time.sleep(i) 111 | >>> maniskill2_learn.check_time('task1') 112 | 2.000 113 | 3.000 114 | 4.000 115 | 5.000 116 | 117 | Args: 118 | timer_id (str): Timer identifier. 119 | """ 120 | if timer_id not in global_timers: 121 | global_timers[timer_id] = Timer() 122 | return 0 123 | else: 124 | return global_timers[timer_id].since_last_check() 125 | 126 | 127 | def get_time_stamp(): 128 | return strftime("%Y%m%d_%H%M%S", localtime()) 129 | 130 | 131 | def get_today(): 132 | return date.today() 133 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/meta/version_utils.py: -------------------------------------------------------------------------------- 1 | def digit_version(version_str): 2 | ret = [] 3 | for x in version_str.split("."): 4 | if x.isdigit(): 5 | ret.append(int(x)) 6 | elif x.find("rc") != -1: 7 | ret = x.split("rc") 8 | ret.append(int(patch_version[0]) - 1) 9 | ret.append(int(patch_version[1])) 10 | return tuple(ret) 11 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/torch/__init__.py: -------------------------------------------------------------------------------- 1 | from .checkpoint_utils import load_checkpoint, save_checkpoint, load_state_dict, get_state_dict 2 | 3 | try: 4 | from .cuda_utils import ( 5 | get_cuda_info, 6 | get_gpu_utilization, 7 | get_gpu_memory_usage_by_process, 8 | get_gpu_memory_usage_by_current_program, 9 | get_device, 10 | get_one_device, 11 | ) 12 | except: 13 | print(f"Not support gpu usage printing") 14 | 15 | from .misc import no_grad, disable_gradients, run_with_mini_batch, mini_batch 16 | from .ops import ( 17 | set_flat_params, 18 | get_flat_params, 19 | get_flat_grads, 20 | set_flat_grads, 21 | batch_random_perm, 22 | masked_average, 23 | masked_max, 24 | smooth_cross_entropy, 25 | batch_rot_with_axis, 26 | soft_update, 27 | hard_update, 28 | avg_grad, 29 | ) 30 | from .logger import * 31 | from .running_stats import RunningMeanStdTorch, MovingMeanStdTorch, RunningSecondMomentumTorch 32 | from .module_utils import BaseAgent, ExtendedModule, ExtendedModuleList, ExtendedDDP, async_no_grad_pi, ExtendedSequential 33 | from .distributions import ScaledTanhNormal, CustomIndependent, ScaledNormal, CustomCategorical 34 | from .optimizer_utils import get_mean_lr, build_optimizer 35 | from .distributed_utils import init_dist, cleanup_dist, master_only, allreduce_params, allreduce_grads, barrier, build_dist_var, get_dist_info 36 | from .freezer import freeze_modules, freeze_params, freeze_bn, unfreeze_modules, unfreeze_params 37 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/torch/cuda_utils.py: -------------------------------------------------------------------------------- 1 | from torch.cuda import _get_device_index as get_device_index 2 | from ..meta.process_utils import get_subprocess_ids 3 | from maniskill2_learn.utils.data import is_dict, is_seq_of, num_to_str 4 | 5 | 6 | try: 7 | import pynvml 8 | from pynvml import NVMLError_DriverNotLoaded 9 | except ModuleNotFoundError: 10 | print("pynvml module not found, please install pynvml") 11 | exit(0) 12 | 13 | try: 14 | pynvml.nvmlInit() 15 | except NVMLError_DriverNotLoaded: 16 | print("cuda driver can't be loaded, is cuda enabled?") 17 | exit(0) 18 | 19 | 20 | def get_gpu_memory_info(device, unit="G", number_only=False): 21 | device = get_device_index(device, optional=True) 22 | handler = pynvml.nvmlDeviceGetHandleByIndex(device) 23 | meminfo = pynvml.nvmlDeviceGetMemoryInfo(handler) 24 | total = num_to_str(meminfo.total, unit, number_only=number_only) 25 | used = num_to_str(meminfo.used, unit, number_only=number_only) 26 | free = num_to_str(meminfo.free, unit, number_only=number_only) 27 | ratio = meminfo.used / meminfo.total 28 | ratio = ratio * 100 if number_only else f"{ratio * 100:.1f}%" 29 | return total, used, free, ratio 30 | 31 | 32 | def get_gpu_memory_usage_by_process(process, device=None, unit="G", number_only=False): 33 | if not isinstance(process, (list, tuple)): 34 | process = [process] 35 | device = get_device_index(device, optional=True) 36 | handler = pynvml.nvmlDeviceGetHandleByIndex(device) 37 | procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handler) 38 | mem = 0 39 | for p in procs: 40 | if p.pid in process: 41 | mem += p.usedGpuMemory 42 | return num_to_str(mem, unit, number_only=number_only) 43 | 44 | 45 | def get_gpu_memory_usage_by_current_program(device=None, unit="G", number_only=False): 46 | proc_in_current_program = get_subprocess_ids() 47 | return get_gpu_memory_usage_by_process(proc_in_current_program) 48 | 49 | 50 | def get_gpu_utilization(device=None): 51 | device = get_device_index(device, optional=True) 52 | handler = pynvml.nvmlDeviceGetHandleByIndex(device) 53 | return pynvml.nvmlDeviceGetUtilizationRates(handler).gpu 54 | 55 | 56 | def get_cuda_info(device=None, unit="G", number_only=True): 57 | current_mem = get_gpu_memory_usage_by_current_program(device, unit, number_only) 58 | all_mem, used, _, ratio = get_gpu_memory_info(device, unit, number_only) 59 | utilization = get_gpu_utilization(device) 60 | return {"gpu_mem_ratio": ratio, "gpu_mem": used, "gpu_mem_this": current_mem, "gpu_util": utilization if number_only else f"{utilization}%"} 61 | 62 | 63 | def get_one_device(x): 64 | if is_dict(x): 65 | return get_one_device(x[list(x.keys())[0]]) 66 | elif is_seq_of(x): 67 | return get_one_device(x[0]) 68 | else: 69 | assert hasattr(x, "device"), type(x) 70 | return x.device 71 | 72 | 73 | def get_device(x): 74 | if is_dict(x): 75 | return {k: get_device(x[k]) for k in x} 76 | elif is_seq_of(x): 77 | return type(x)([get_device(y) for y in x]) 78 | else: 79 | assert hasattr(x, "device"), type(x) 80 | return x.device 81 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/torch/freezer.py: -------------------------------------------------------------------------------- 1 | """ 2 | From Jiayuan 3 | 4 | Helpers for operating modules/parameters 5 | """ 6 | 7 | import re 8 | import torch.nn as nn 9 | from ..data.string_utils import any_string, regex_match 10 | 11 | 12 | def get_frozen_params(module): 13 | return [name for name, params in module.named_parameters() if not params.requires_grad] 14 | 15 | 16 | def get_frozen_modules(module): 17 | return [name for name, m in module.named_modules() if not m.training] 18 | 19 | 20 | def print_frozen_modules_and_params(module, logger=None): 21 | _print = print if logger is None else logger.info 22 | for name in get_frozen_modules(module): 23 | _print("Module {} is frozen.".format(name)) 24 | for name in get_frozen_params(module): 25 | _print("Params {} is frozen.".format(name)) 26 | 27 | 28 | def apply_params(module, patterns=any_string, requires_grad=False): 29 | """Apply freeze/unfreeze on parameters 30 | 31 | Args: 32 | module (torch.nn.Module): the module to apply 33 | patterns (sequence of str): strings which define all the patterns of interests 34 | requires_grad (bool, optional): whether to freeze params 35 | 36 | """ 37 | if isinstance(patterns, str): 38 | patterns = [patterns] 39 | for name, params in module.named_parameters(): 40 | for pattern in patterns: 41 | assert isinstance(pattern, str) 42 | if regex_match(name, pattern): 43 | params.requires_grad = requires_grad 44 | 45 | 46 | def apply_modules(module, patterns=any_string, mode=False, prefix=""): 47 | """Apply train/eval on modules 48 | 49 | Args: 50 | module (torch.nn.Module): the module to apply 51 | patterns (sequence of str): strings which define all the patterns of interests 52 | mode (bool, optional): whether to set the module training mode 53 | prefix (str, optional) 54 | 55 | """ 56 | if isinstance(patterns, str): 57 | patterns = [patterns] 58 | for name, m in module._modules.items(): 59 | for pattern in patterns: 60 | assert isinstance(pattern, str) 61 | full_name = prefix + ("." if prefix else "") + name 62 | if regex_match(full_name, pattern): 63 | # avoid redundant call 64 | print(name) 65 | m.train(mode) 66 | else: 67 | apply_modules(m, patterns, mode=mode, prefix=full_name) 68 | 69 | 70 | def freeze_modules(module, patterns=any_string): 71 | """Freeze modules by matching patterns""" 72 | apply_modules(module, patterns, mode=False) 73 | 74 | 75 | def unfreeze_modules(module, patterns=any_string): 76 | """Unfreeze module by matching patterns""" 77 | apply_modules(module, patterns, mode=True) 78 | 79 | 80 | def freeze_params(module, patterns=any_string): 81 | """Freeze parameters by matching patterns""" 82 | apply_params(module, patterns, requires_grad=False) 83 | 84 | 85 | def unfreeze_params(module, patterns=any_string): 86 | """Unfreeze parameters by matching patterns""" 87 | apply_params(module, patterns, requires_grad=True) 88 | 89 | 90 | def apply_bn(module, mode, requires_grad): 91 | """Modify batch normalization in the module 92 | 93 | Args: 94 | module (nn.Module): the module to operate 95 | mode (bool): train/eval mode 96 | requires_grad (bool): whether parameters require gradients 97 | 98 | Notes: 99 | Note that the difference between the behaviors of BatchNorm.eval() and BatchNorm(track_running_stats=False) 100 | 101 | """ 102 | for m in module.modules(): 103 | if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)): 104 | m.train(mode) 105 | for params in m.parameters(): 106 | params.requires_grad = requires_grad 107 | 108 | 109 | def freeze_bn(module): 110 | apply_bn(module, mode=False, requires_grad=False) 111 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/torch/logger/__init__.py: -------------------------------------------------------------------------------- 1 | from .tensorboard_logger import TensorboardLogger 2 | from .tensorboard_utils import load_tb_summaries_as_df 3 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/torch/logger/tensorboard_logger.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import numbers, numpy as np 3 | import matplotlib.pyplot as plt 4 | 5 | 6 | class TensorboardLogger: 7 | def __init__(self, log_dir=None): 8 | from torch.utils.tensorboard import SummaryWriter 9 | 10 | self.writer = SummaryWriter(osp.join(log_dir, "tf_logs")) 11 | 12 | def get_lr_tags(self, runner): 13 | tags = {} 14 | lrs = runner.current_lr() 15 | if isinstance(lrs, dict): 16 | for name, value in lrs.items(): 17 | tags[f"learning_rate/{name}"] = value[0] 18 | else: 19 | tags["learning_rate"] = lrs[0] 20 | return tags 21 | 22 | def get_momentum_tags(self, runner): 23 | tags = {} 24 | momentums = runner.current_momentum() 25 | if isinstance(momentums, dict): 26 | for name, value in momentums.items(): 27 | tags[f"momentum/{name}"] = value[0] 28 | else: 29 | tags["momentum"] = momentums[0] 30 | return tags 31 | 32 | @staticmethod 33 | def is_scalar(val, include_np=True, include_torch=True): 34 | if isinstance(val, numbers.Number): 35 | return True 36 | elif include_np and isinstance(val, np.ndarray) and val.ndim == 0: 37 | return True 38 | else: 39 | import torch 40 | 41 | if include_torch and isinstance(val, torch.Tensor) and len(val) == 1: 42 | return True 43 | else: 44 | return False 45 | 46 | def get_loggable_tags(self, output, allow_scalar=True, allow_text=False, tags_to_skip=("time", "data_time"), add_mode=True, tag_name="train"): 47 | tags = {} 48 | for tag, val in output.items(): 49 | if tag in tags_to_skip: 50 | continue 51 | if self.is_scalar(val) and not allow_scalar: 52 | continue 53 | if isinstance(val, str) and not allow_text: 54 | continue 55 | if add_mode and "/" not in tag: 56 | tag = f"{tag_name}/{tag}" 57 | tags[tag] = val 58 | return tags 59 | 60 | def log(self, tags, n_iter, tag_name="train"): 61 | tags = self.get_loggable_tags(tags, tag_name=tag_name) 62 | for tag, val in tags.items(): 63 | if isinstance(val, str): 64 | self.writer.add_text(tag, val, n_iter) 65 | elif np.isscalar(val) or val.size == 1: 66 | self.writer.add_scalar(tag, val, n_iter) 67 | else: 68 | if val.ndim == 2: 69 | cmap = plt.get_cmap('jet') 70 | val = cmap(val)[..., :3] 71 | assert val.ndim == 3, f"Image should have two dimension! You provide: {tag, val.shape}!" 72 | self.writer.add_image(tag, val, n_iter, dataformats='HWC') 73 | 74 | def close(self): 75 | self.writer.close() 76 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/torch/logger/tensorboard_utils.py: -------------------------------------------------------------------------------- 1 | import pandas, os.path as osp, re, pathlib, os 2 | 3 | 4 | def load_tb_summaries_as_df(tb_files, exp_names, data_title="test"): 5 | from tensorboard.backend.event_processing import event_accumulator 6 | 7 | if isinstance(tb_files, (list, tuple)): 8 | return [load_tb_summaries_as_df(tb_file, exp_name, data_title) for tb_file, exp_name in zip(tb_files, exp_names)] 9 | else: 10 | event = event_accumulator.EventAccumulator(tb_files) 11 | event.Reload() 12 | # df = pandas.DataFrame() 13 | exp_begin_time = event.FirstEventTimestamp() 14 | res = {"exp_name": exp_names} 15 | running_time = -1 16 | step = None 17 | 18 | for key in event.scalars.Keys(): 19 | if key.startswith(data_title): 20 | content = event.scalars.Items(key) 21 | # print(len(content)) 22 | hours = (content[-1].wall_time - exp_begin_time) / 3600 23 | if hours > running_time: 24 | running_time = hours 25 | if step is None: 26 | step = content[-1].step 27 | res["steps"] = [_.step for _ in content] 28 | if step != content[-1].step: 29 | print(f"Step is not consistent: {key}, old {step}, new {content.step}") 30 | new_key = re.sub(r"\W+", "", key[len(data_title) :]) 31 | res[new_key] = [_.value for _ in content] 32 | res["running_time"] = running_time 33 | # df = df.append(res, ignore_index=True) 34 | return res 35 | 36 | 37 | def parse_tb_for_rl_exp(path, output_name, data_title="test"): 38 | """ 39 | Folder structure: Env_name/tf_logs/*tfevents* 40 | """ 41 | path = pathlib.Path(path) 42 | tb_files = [] 43 | exp_names = [] 44 | os.makedirs(osp.dirname(output_name), exist_ok=True) 45 | for env_name in path.glob("*"): 46 | if env_name.is_dir(): 47 | print(f"Find env {env_name.name}") 48 | exp_names.append(env_name.name) 49 | events = list(env_name.glob("*/**/*.tfevents.*")) 50 | if len(events) > 1: 51 | print(f"There are mote than one tensorbaord logs in the folder: {str(env_name)}") 52 | tb_files.append(str(events[0])) 53 | df = load_tb_summaries_as_df(tb_files, exp_names, data_title) 54 | print(f"Save log to {osp.abspath(output_name)}") 55 | df.to_csv(output_name, float_format="%.4f") 56 | 57 | 58 | def parse_tb_for_rl_alg(path, output_dir="csv_logs", data_title="test"): 59 | """ 60 | Folder structure: Alogrithm_name/Env_name/tf_logs/*tfevents* 61 | """ 62 | path = pathlib.Path(path) 63 | for alg_name in path.glob("*"): 64 | if alg_name.is_dir(): 65 | print(f"Find alg {alg_name.name}") 66 | parse_tb_for_rl_exp(str(alg_name), osp.join(output_dir, f"{alg_name.name}.csv"), data_title) 67 | 68 | 69 | if __name__ == "__main__": 70 | path = "/home/lz/data/Projects/MBRL/final_results/" 71 | parse_tb_for_rl_alg(path) 72 | 73 | # df = loaf_tb_summary_as_df(path, 'Ant') 74 | # print(df) 75 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/torch/logger/wandb_logger.py: -------------------------------------------------------------------------------- 1 | import wandb 2 | 3 | 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/torch/misc.py: -------------------------------------------------------------------------------- 1 | from functools import wraps 2 | import numpy as np 3 | import torch 4 | from maniskill2_learn.utils.math import split_num 5 | from maniskill2_learn.utils.meta import get_logger 6 | from maniskill2_learn.utils.data import DictArray, to_np, GDict, to_torch 7 | 8 | 9 | def disable_gradients(network): 10 | for param in network.parameters(): 11 | param.requires_grad = False 12 | 13 | 14 | def worker_init_fn(worker_id): 15 | """The function is designed for pytorch multi-process dataloader. 16 | Note that we use the pytorch random generator to generate a base_seed. Please try to be consistent. 17 | References: 18 | https://pytorch.org/docs/stable/notes/faq.html#dataloader-workers-random-seed 19 | """ 20 | base_seed = torch.IntTensor(1).random_().item() 21 | np.random.seed(base_seed + worker_id) 22 | 23 | 24 | def no_grad(f): 25 | wraps(f) 26 | 27 | def wrapper(*args, **kwargs): 28 | with torch.no_grad(): 29 | return f(*args, **kwargs) 30 | 31 | return wrapper 32 | 33 | 34 | def get_seq_info(done_mask): 35 | # It will sort the length of the sequence to improve the performance 36 | 37 | # input: done_mask [L] 38 | # return: index [#seq, max_seq_len]; sorted_idx [#seq]; is_valid [#seq, max_seq_len] 39 | done_mask = to_np(done_mask) 40 | 41 | indices = np.where(done_mask)[0] 42 | one = np.ones(1, dtype=indices.dtype) 43 | indices = np.concatenate([one * -1, indices]) 44 | len_seq = indices[1:] - indices[:-1] 45 | 46 | sorted_idx = np.argsort(-len_seq, kind="stable") # From max to min 47 | max_len = len_seq[sorted_idx[0]] 48 | index = np.zeros([len(sorted_idx), max_len], dtype=np.int64) 49 | is_valid = np.zeros([len(sorted_idx), max_len], dtype=np.bool_) 50 | 51 | for i, idx in enumerate(sorted_idx): 52 | index[i, : len_seq[idx]] = np.arange(len_seq[idx]) + indices[idx] + 1 53 | is_valid[i, : len_seq[idx]] = True 54 | return index, sorted_idx, is_valid 55 | 56 | 57 | def run_with_mini_batch( 58 | function, 59 | *args, 60 | batch_size=None, 61 | wrapper=True, 62 | device=None, 63 | ret_device=None, 64 | episode_dones=None, 65 | **kwargs, 66 | ): 67 | """ 68 | Run a pytorch function with mini-batch when the batch size of dat is very large. 69 | :param function: the function 70 | :param data: the input data which should be in dict array structure 71 | :param batch_size: the num of samples in the whole batch 72 | :return: all the outputs. 73 | """ 74 | capacity = None 75 | # print('In mini batch', batch_size, DictArray(kwargs).shape) 76 | 77 | def process_kwargs(x): 78 | if x is None or len(x) == 0: 79 | return None 80 | # print(type(x)) 81 | nonlocal capacity, device, ret_device 82 | x = DictArray(x) 83 | # print(x.shape, x.type) 84 | # exit(0) 85 | capacity = x.capacity 86 | if device is None: 87 | device = x.one_device 88 | if ret_device is None: 89 | ret_device = x.one_device 90 | return x 91 | 92 | args, kwargs = list(args), dict(kwargs) 93 | # print(GDict(args).type) 94 | # print(GDict(kwargs).type) 95 | 96 | args = process_kwargs(args) 97 | kwargs = process_kwargs(kwargs) 98 | 99 | assert capacity is not None, "Input is None" 100 | if batch_size is None: 101 | batch_size = capacity 102 | 103 | ret = [] 104 | # print(capacity, batch_size) 105 | for i in range(0, capacity, batch_size): 106 | num_i = min(capacity - i, batch_size) 107 | args_i = args.slice(slice(i, i + num_i)).to_torch(device=device, wrapper=False) if args is not None else [] 108 | kwargs_i = kwargs.slice(slice(i, i + num_i)).to_torch(device=device, wrapper=False) if kwargs is not None else {} 109 | ret.append(GDict(function(*args_i, **kwargs_i)).to_torch(device=ret_device, wrapper=False)) 110 | ret = DictArray.concat(ret, axis=0, wrapper=wrapper) 111 | 112 | return ret 113 | 114 | 115 | def mini_batch(wrapper_=True): 116 | def actual_mini_batch(f): 117 | wraps(f) 118 | 119 | def wrapper(*args, batch_size=None, wrapper=None, device=None, ret_device=None, **kwargs): 120 | if wrapper is None: 121 | wrapper = wrapper_ 122 | # print(batch_size, dict(kwargs)) 123 | 124 | return run_with_mini_batch(f, *args, **kwargs, batch_size=batch_size, wrapper=wrapper, device=device, ret_device=ret_device) 125 | 126 | return wrapper 127 | 128 | return actual_mini_batch 129 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/torch/optimizer_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import copy, inspect, torch 3 | from maniskill2_learn.utils.meta import Registry, build_from_cfg 4 | from maniskill2_learn.utils.data import regex_match 5 | 6 | 7 | OPTIMIZERS = Registry("optimizer") 8 | 9 | def register_torch_optimizers(): 10 | torch_optimizers = [] 11 | for module_name in dir(torch.optim): 12 | if module_name.startswith("__"): 13 | continue 14 | _optim = getattr(torch.optim, module_name) 15 | if inspect.isclass(_optim) and issubclass(_optim, torch.optim.Optimizer): 16 | OPTIMIZERS.register_module()(_optim) 17 | torch_optimizers.append(module_name) 18 | return torch_optimizers 19 | 20 | 21 | TORCH_OPTIMIZERS = register_torch_optimizers() 22 | 23 | 24 | def get_mean_lr(optimizer, mean=True): 25 | ret = [] 26 | for param_group in optimizer.param_groups: 27 | ret.append(param_group["lr"]) 28 | return np.mean(ret) if mean else ret 29 | 30 | 31 | def build_optimizer(model, cfg): 32 | cfg = copy.deepcopy(cfg) 33 | constructor_type = cfg.pop("constructor", "default") # keyword "type" is saved for specify optimizer. 34 | if constructor_type == "default": 35 | param_cfg = cfg.pop("param_cfg", None) 36 | 37 | param_i_template = copy.deepcopy(cfg) 38 | param_i_template.pop("type", None) 39 | 40 | params = [] 41 | existing_params = [] 42 | if hasattr(model, "named_parameters"): 43 | # It is a mode not a Parameter. 44 | for name, param in model.named_parameters(): 45 | if id(param) in existing_params or not param.requires_grad: 46 | continue 47 | param_i = copy.deepcopy(param_i_template) 48 | existing_params.append(id(param)) 49 | if param_cfg is not None: 50 | for pattern, param_config in param_cfg.items(): 51 | if regex_match(name, pattern): 52 | param_i = param_config 53 | break 54 | if param_i is None: 55 | continue 56 | param_i = {'params': param} 57 | params.append(param_i) 58 | else: 59 | params = [model, ] 60 | cfg["params"] = params 61 | optimizer = build_from_cfg(cfg, OPTIMIZERS) 62 | else: 63 | raise NotImplementedError 64 | return optimizer 65 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/visualization/__init__.py: -------------------------------------------------------------------------------- 1 | from .o3d_utils import visualize_3d, visualize_pcd 2 | -------------------------------------------------------------------------------- /maniskill2_learn/utils/visualization/o3d_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np, open3d as o3d 2 | from ..lib3d import np2pcd, to_o3d 3 | 4 | 5 | def visualize_3d(objects, show_frame=True, frame_size=1.0, frame_origin=(0, 0, 0)): 6 | if not isinstance(objects, (list, tuple)): 7 | objects = [objects] 8 | objects = [to_o3d(obj) for obj in objects] 9 | if show_frame: 10 | objects.append(o3d.geometry.TriangleMesh.create_coordinate_frame(size=frame_size, origin=frame_origin)) 11 | return o3d.visualization.draw_geometries(objects) 12 | 13 | 14 | def visualize_pcd(points, colors=None, normals=None, bbox=None, show_frame=False, frame_size=1.0, frame_origin=(0, 0, 0)): 15 | """Visualize a point cloud.""" 16 | pc = np2pcd(points, colors, normals) 17 | geometries = [pc] 18 | if show_frame: 19 | coord_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=frame_size, origin=frame_origin) 20 | geometries.append(coord_frame) 21 | if bbox is None: 22 | bbox = [] 23 | elif not isinstance(bbox, (tuple, list)): 24 | bbox = [bbox] 25 | o3d.visualization.draw_geometries(geometries + bbox) 26 | -------------------------------------------------------------------------------- /maniskill2_learn/version.py: -------------------------------------------------------------------------------- 1 | __version__ = "1.8.0b0" 2 | 3 | 4 | def parse_version_info(version_str): 5 | version_info = [] 6 | for x in version_str.split("."): 7 | if x.isdigit(): 8 | version_info.append(int(x)) 9 | elif x.find("rc") != -1: 10 | patch_version = x.split("rc") 11 | version_info.append(int(patch_version[0])) 12 | version_info.append(f"rc{patch_version[1]}") 13 | return tuple(version_info) 14 | 15 | 16 | version_info = parse_version_info(__version__) 17 | __all__ = ["__version__", "version_info", "parse_version_info"] 18 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | addict 2 | cython 3 | numpy 4 | Pillow 5 | pyyaml 6 | regex;sys_platform=='win32' 7 | yapf 8 | ffmpeg-python 9 | pandas 10 | tables 11 | h5py 12 | coverage 13 | lmdb 14 | pytest 15 | PyTurboJPEG 16 | open3d 17 | trimesh 18 | gymnasium 19 | shapely 20 | transforms3d 21 | sorcery 22 | psutil 23 | opencv-python 24 | tensorboard 25 | tabulate 26 | rtree 27 | GitPython 28 | pynvml 29 | kubernetes 30 | pynput 31 | sqlalchemy 32 | docker 33 | crc32c 34 | pypi-simple 35 | numpy-quaternion 36 | termcolor 37 | pymeshlab 38 | plyfile 39 | einops 40 | ninja 41 | pytorch3d 42 | deprecation 43 | -------------------------------------------------------------------------------- /scripts/example_demo_conversion/general_rigid_body_multi_object_envs.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # For general envs with multiple objects, where each object has its own demos, 4 | # and goal information is not provided in the observation 5 | # i.e. TurnFaucet-v0; for ManiSkill1 envs (OpenCabinetDoor/Drawer, PushChair, MoveBucket), please use maniskill1.sh instead of this script) 6 | ENV="TurnFaucet-v0" 7 | 8 | # unzip the demonstration for TurnFaucet, if not already 9 | if [[ $ENV =~ "TurnFaucet-v0" ]]; then 10 | cd ../ManiSkill2/demos/v0/rigid_body/$ENV/ 11 | if [[ -f "20220815.zip" ]]; then 12 | unzip 20220815.zip 13 | rm 20220815.zip 14 | fi 15 | cd - # ManiSkill2-Learn 16 | python tools/merge_trajectory.py \ 17 | -i ../ManiSkill2/demos/v0/v0/rigid_body/$ENV/ \ 18 | -o ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.h5 \ 19 | -p 5*.h5 # this can be replaced with other patterns 20 | else 21 | # Assume current working directory is ManiSkill2-Learn/ 22 | # Assume conda environment contains all dependencies of ManiSkill2 and ManiSkill2-Learn 23 | # Inside ManiSkill2's directory, run merge_trajectory to output merged h5 and json files 24 | python tools/merge_trajectory.py \ 25 | -i ../ManiSkill2/demos/v0/rigid_body/$ENV/ \ 26 | -o ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.h5 \ 27 | -p trajectory.h5 # this can be replaced with other patterns 28 | fi 29 | 30 | # Inside the ManiSkill2's directory, run replay_trajectory.py. See wiki page 31 | # of ManiSkill2 for more information. 32 | cd ../ManiSkill2 33 | python mani_skill2/trajectory/replay_trajectory.py --num-procs 32 \ 34 | --traj-path demos/v0/rigid_body/$ENV/trajectory_merged.h5 \ 35 | --save-traj \ 36 | --target-control-mode pd_ee_delta_pose \ 37 | --obs-mode none 38 | 39 | # Inside ManiSkill2-Learn's directory, run convert_state.py to generate visual observations 40 | # for the demonstrations. 41 | cd ../ManiSkill2-Learn 42 | # Generate pointcloud demo 43 | python tools/convert_state.py \ 44 | --env-name=$ENV \ 45 | --num-procs=12 \ 46 | --traj-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose.h5 \ 47 | --json-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose.json \ 48 | --output-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_pointcloud.h5 \ 49 | --control-mode=pd_ee_delta_pose \ 50 | --max-num-traj=-1 \ 51 | --obs-mode=pointcloud \ 52 | --reward-mode=dense \ 53 | --obs-frame=ee \ 54 | --n-points=1200 55 | 56 | # Generate rgbd demo 57 | python tools/convert_state.py \ 58 | --env-name=$ENV \ 59 | --num-procs=12 \ 60 | --traj-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose.h5 \ 61 | --json-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose.json \ 62 | --output-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_rgbd.h5 \ 63 | --control-mode=pd_ee_delta_pose \ 64 | --max-num-traj=-1 \ 65 | --obs-mode=rgbd \ 66 | --reward-mode=dense 67 | 68 | # Shuffle pointcloud demos 69 | python tools/shuffle_demo.py \ 70 | --source-file ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_pointcloud.h5 \ 71 | --target-file ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_pointcloud_shuffled.h5 72 | 73 | # Shuffle rgbd demos 74 | python tools/shuffle_demo.py \ 75 | --source-file ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_rgbd.h5 \ 76 | --target-file ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_rgbd_shuffled.h5 77 | -------------------------------------------------------------------------------- /scripts/example_demo_conversion/general_rigid_body_single_object_envs.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # For all rigid-body environments which do not have demo for each individual asset 4 | # and which do provide goal position in observation 5 | # i.e. StackCube-v0, PegInsertionSide-v0, PlugCharger-v0, PandaAvoidObstacles-v0, AssemblingKits-v0 6 | ENV="StackCube-v0" 7 | 8 | # Assume current working directory is ManiSkill2-Learn/ 9 | # Assume conda environment contains all dependencies of ManiSkill2 and ManiSkill2-Learn 10 | # Inside the ManiSkill2's directory, run replay_trajectory.py. See wiki page 11 | # of ManiSkill2 for more information. 12 | cd ../ManiSkill2 13 | python mani_skill2/trajectory/replay_trajectory.py --num-procs 32 \ 14 | --traj-path demos/v0/rigid_body/$ENV/trajectory.h5 \ 15 | --save-traj \ 16 | --target-control-mode pd_ee_delta_pose \ 17 | --obs-mode none 18 | 19 | # Inside ManiSkill2-Learn's directory, run convert_state.py to generate visual observations 20 | # for the demonstrations. 21 | cd ../ManiSkill2-Learn 22 | 23 | # Generate pointcloud demo 24 | python tools/convert_state.py \ 25 | --env-name=$ENV \ 26 | --num-procs=12 \ 27 | --traj-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose.h5 \ 28 | --json-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose.json \ 29 | --output-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose_pointcloud.h5 \ 30 | --control-mode=pd_ee_delta_pose \ 31 | --max-num-traj=-1 \ 32 | --obs-mode=pointcloud \ 33 | --reward-mode=dense \ 34 | --obs-frame=ee \ 35 | --n-points=1200 36 | 37 | # Generate rgbd demo 38 | python tools/convert_state.py \ 39 | --env-name=$ENV \ 40 | --num-procs=12 \ 41 | --traj-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose.h5 \ 42 | --json-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose.json \ 43 | --output-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose_rgbd.h5 \ 44 | --control-mode=pd_ee_delta_pose \ 45 | --max-num-traj=-1 \ 46 | --obs-mode=rgbd \ 47 | --reward-mode=dense -------------------------------------------------------------------------------- /scripts/example_demo_conversion/general_soft_body_envs.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # For all soft-body environments 4 | # i.e. Fill-v0, Write-v0, Pinch-v0, Hang-v0, Pour-v0, Excavate-v0 5 | ENV="Excavate-v0" 6 | 7 | # Create a dummy environment to initialize cache if it has not been done in the history, otherwise 8 | # we can't do multi-processing below 9 | CUDA_VISIBLE_DEVICES=1 python -c "import mani_skill2.envs, gym; env=gym.make('$ENV'); env.reset();" 10 | 11 | # Assume current working directory is ManiSkill2-Learn/ 12 | # Assume conda environment contains all dependencies of ManiSkill2 and ManiSkill2-Learn 13 | # Inside the ManiSkill2's directory, run replay_trajectory.py. See wiki page 14 | # of ManiSkill2 for more information. 15 | cd ../ManiSkill2 16 | CUDA_VISIBLE_DEVICES=1 python mani_skill2/trajectory/replay_trajectory.py --num-procs 4 \ 17 | --traj-path demos/v0/soft_body/$ENV/trajectory.h5 \ 18 | --save-traj \ 19 | --target-control-mode pd_ee_delta_pose \ 20 | --obs-mode none 21 | 22 | # Inside ManiSkill2-Learn's directory, run convert_state.py to generate visual observations 23 | # for the demonstrations. 24 | cd ../ManiSkill2-Learn 25 | 26 | # Generate pointcloud demo 27 | CUDA_VISIBLE_DEVICES=1 python tools/convert_state.py \ 28 | --env-name=$ENV \ 29 | --num-procs=4 \ 30 | --traj-name=../ManiSkill2/demos/v0/soft_body/$ENV/trajectory.none.pd_ee_delta_pose.h5 \ 31 | --json-name=../ManiSkill2/demos/v0/soft_body/$ENV/trajectory.none.pd_ee_delta_pose.json \ 32 | --output-name=../ManiSkill2/demos/v0/soft_body/$ENV/trajectory.none.pd_ee_delta_pose_pointcloud.h5 \ 33 | --control-mode=pd_ee_delta_pose \ 34 | --max-num-traj=-1 \ 35 | --obs-mode=pointcloud \ 36 | --reward-mode=dense \ 37 | --obs-frame=ee \ 38 | --n-points=1200 39 | 40 | # Generate rgbd demo 41 | CUDA_VISIBLE_DEVICES=1 python tools/convert_state.py \ 42 | --env-name=$ENV \ 43 | --num-procs=4 \ 44 | --traj-name=../ManiSkill2/demos/v0/soft_body/$ENV/trajectory.none.pd_ee_delta_pose.h5 \ 45 | --json-name=../ManiSkill2/demos/v0/soft_body/$ENV/trajectory.none.pd_ee_delta_pose.json \ 46 | --output-name=../ManiSkill2/demos/v0/soft_body/$ENV/trajectory.none.pd_ee_delta_pose_rgbd.h5 \ 47 | --control-mode=pd_ee_delta_pose \ 48 | --max-num-traj=-1 \ 49 | --obs-mode=rgbd \ 50 | --reward-mode=dense -------------------------------------------------------------------------------- /scripts/example_demo_conversion/maniskill1.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # For environments migrated from ManiSkill1 4 | # i.e. OpenCabinetDoor-v1, OpenCabinetDrawer-v1, PushChair-v1, MoveBucket-v1 5 | ENV="OpenCabinetDrawer-v1" 6 | 7 | # Assume current working directory is ManiSkill2-Learn/ 8 | cd ../ManiSkill2 9 | 10 | find demos/v0/rigid_body/$ENV -name "*trajectory.h5" | while read line; do 11 | 12 | CUDA_VISIBLE_DEVICES=1 python mani_skill2/trajectory/replay_trajectory.py --num-procs 16 \ 13 | --traj-path $line \ 14 | --use-env-states \ 15 | --save-traj \ 16 | --obs-mode none 17 | 18 | # Inside ManiSkill2-Learn's directory, run convert_state.py to generate visual observations 19 | # for the demonstrations. 20 | cd ../ManiSkill2-Learn 21 | 22 | Generate pointcloud demo 23 | python tools/convert_state.py \ 24 | --env-name=$ENV \ 25 | --num-procs=12 \ 26 | --traj-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.base_pd_joint_vel_arm_pd_ee_delta_pose.h5 \ 27 | --json-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.base_pd_joint_vel_arm_pd_ee_delta_pose.json \ 28 | --output-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.base_pd_joint_vel_arm_pd_ee_delta_pose_pointcloud.h5 \ 29 | --control-mode=base_pd_joint_vel_arm_pd_joint_vel \ 30 | --max-num-traj=-1 \ 31 | --obs-mode=pointcloud \ 32 | --reward-mode=dense \ 33 | --obs-frame=ee \ 34 | --n-points=1200 35 | 36 | # Generate rgbd demo 37 | CUDA_VISIBLE_DEVICES=1 python tools/convert_state.py \ 38 | --env-name=$ENV \ 39 | --num-procs=12 \ 40 | --traj-name=../ManiSkill2/"$(dirname "$line")"/trajectory.none.base_pd_joint_vel_arm_pd_joint_vel.h5 \ 41 | --json-name=../ManiSkill2/"$(dirname "$line")"/trajectory.none.base_pd_joint_vel_arm_pd_joint_vel.json \ 42 | --output-name=../ManiSkill2/"$(dirname "$line")"/trajectory.none.base_pd_joint_vel_arm_pd_joint_vel_rgbd.h5 \ 43 | --control-mode=base_pd_joint_vel_arm_pd_joint_vel \ 44 | --max-num-traj=-1 \ 45 | --obs-mode=rgbd \ 46 | --reward-mode=dense 47 | 48 | cd ../ManiSkill2 49 | 50 | done 51 | -------------------------------------------------------------------------------- /scripts/example_demo_conversion/pick_place_multi_object_envs.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # For Pick & Place envs with multiple objects, where each object has its own demos, 4 | # and goal information is provided in the observation 5 | # i.e. PickSingleYCB-v0, PickSingleEGAD-v0, PickClutterYCB-v0 6 | ENV="PickSingleYCB-v0" 7 | 8 | if [[ $ENV =~ "PickSingleYCB-v0" ]]; then 9 | cd ../ManiSkill2/demos/v0/rigid_body/$ENV/ 10 | if [[ -f "220815.zip" ]]; then 11 | unzip 220815.zip 12 | rm 220815.zip 13 | fi 14 | cd - # ManiSkill2-Learn 15 | python tools/merge_trajectory.py \ 16 | -i ../ManiSkill2/demos/v0/rigid_body/$ENV/ \ 17 | -o ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.h5 \ 18 | -p 0*.h5 # this can be replaced with other patterns 19 | else 20 | # Assume current working directory is ManiSkill2-Learn/ 21 | # Assume conda environment contains all dependencies of ManiSkill2 and ManiSkill2-Learn 22 | # Inside ManiSkill2's directory, run merge_trajectory to output merged h5 and json files 23 | python tools/merge_trajectory.py \ 24 | -i ../ManiSkill2/demos/v0/rigid_body/$ENV/ \ 25 | -o ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.h5 \ 26 | -p trajectory.h5 # this can be replaced with other patterns 27 | fi 28 | 29 | # Inside the ManiSkill2's directory, run replay_trajectory.py. See wiki page 30 | # of ManiSkill2 for more information. 31 | cd ../ManiSkill2 32 | python mani_skill2/trajectory/replay_trajectory.py --num-procs 32 \ 33 | --traj-path demos/v0/rigid_body/$ENV/trajectory_merged.h5 \ 34 | --save-traj \ 35 | --target-control-mode pd_ee_delta_pose \ 36 | --obs-mode none 37 | 38 | # Inside ManiSkill2-Learn's directory, run convert_state.py to generate visual observations 39 | # for the demonstrations. 40 | cd ../ManiSkill2-Learn 41 | # Generate pointcloud demo 42 | python tools/convert_state.py \ 43 | --env-name=$ENV \ 44 | --num-procs=12 \ 45 | --traj-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose.h5 \ 46 | --json-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose.json \ 47 | --output-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_pointcloud.h5 \ 48 | --control-mode=pd_ee_delta_pose \ 49 | --max-num-traj=-1 \ 50 | --obs-mode=pointcloud \ 51 | --reward-mode=dense \ 52 | --obs-frame=ee \ 53 | --n-points=1200 \ 54 | --n-goal-points=50 55 | 56 | # Generate rgbd demo 57 | python tools/convert_state.py \ 58 | --env-name=$ENV \ 59 | --num-procs=12 \ 60 | --traj-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose.h5 \ 61 | --json-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose.json \ 62 | --output-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_rgbd.h5 \ 63 | --control-mode=pd_ee_delta_pose \ 64 | --max-num-traj=-1 \ 65 | --obs-mode=rgbd \ 66 | --reward-mode=dense 67 | 68 | # Shuffle pointcloud demos 69 | python tools/shuffle_demo.py \ 70 | --source-file ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_pointcloud.h5 \ 71 | --target-file ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_pointcloud_shuffled.h5 72 | 73 | # Shuffle rgbd demos 74 | python tools/shuffle_demo.py \ 75 | --source-file ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_rgbd.h5 \ 76 | --target-file ../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory_merged.none.pd_ee_delta_pose_rgbd_shuffled.h5 77 | -------------------------------------------------------------------------------- /scripts/example_demo_conversion/pick_place_single_object_envs.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # For Pick & Place envs with single object. These envs provide goal position in observation; 4 | # Currently only PickCube-v0 ibelongs to such environment 5 | ENV="PickCube-v0" 6 | 7 | # Assume current working directory is ManiSkill2-Learn/ 8 | # Assume conda environment contains all dependencies of ManiSkill2 and ManiSkill2-Learn 9 | # Inside the ManiSkill2's directory, run replay_trajectory.py. See wiki page 10 | # of ManiSkill2 for more information. 11 | cd ../ManiSkill2 12 | python mani_skill2/trajectory/replay_trajectory.py --num-procs 32 \ 13 | --traj-path demos/v0/rigid_body/$ENV/trajectory.h5 \ 14 | --save-traj \ 15 | --target-control-mode pd_ee_delta_pose \ 16 | --obs-mode none 17 | 18 | # Inside ManiSkill2-Learn's directory, run convert_state.py to generate visual observations 19 | # for the demonstrations. 20 | cd ../ManiSkill2-Learn 21 | 22 | # Generate pointcloud demo; num_procs=64 takes about 30GB memory and 32 cpus 23 | python tools/convert_state.py \ 24 | --env-name=$ENV \ 25 | --num-procs=12 \ 26 | --traj-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose.h5 \ 27 | --json-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose.json \ 28 | --output-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose_pointcloud.h5 \ 29 | --control-mode=pd_ee_delta_pose \ 30 | --max-num-traj=-1 \ 31 | --obs-mode=pointcloud \ 32 | --reward-mode=dense \ 33 | --obs-frame=ee \ 34 | --n-points=1200 \ 35 | --n-goal-points=50 36 | 37 | # Generate rgbd demo 38 | python tools/convert_state.py \ 39 | --env-name=$ENV \ 40 | --num-procs=12 \ 41 | --traj-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose.h5 \ 42 | --json-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose.json \ 43 | --output-name=../ManiSkill2/demos/v0/rigid_body/$ENV/trajectory.none.pd_ee_delta_pose_rgbd.h5 \ 44 | --control-mode=pd_ee_delta_pose \ 45 | --max-num-traj=-1 \ 46 | --obs-mode=rgbd \ 47 | --reward-mode=dense -------------------------------------------------------------------------------- /scripts/example_training/pretrained_model/bc_soft_body_pointcloud.sh: -------------------------------------------------------------------------------- 1 | # Replace --work-dir, env_cfg.env_name, and replay_cfg.buffer_filenames when you run other environments 2 | 3 | python maniskill2_learn/apis/run_rl.py configs/brl/bc/pointnet_soft_body.py \ 4 | --work-dir ./logs/bc_excavate_pointcloud --gpu-ids 0 \ 5 | --cfg-options "env_cfg.env_name=Excavate-v0" "env_cfg.obs_mode=pointcloud" "env_cfg.n_points=1200" \ 6 | "env_cfg.control_mode=pd_joint_delta_pos" \ 7 | "replay_cfg.buffer_filenames=../ManiSkill2/demos/v0/soft_body/Excavate-v0/trajectory.none.pd_ee_delta_pose_pointcloud.h5" \ 8 | "eval_cfg.num=100" "eval_cfg.save_traj=False" "eval_cfg.save_video=True" \ 9 | "train_cfg.n_eval=50000" "train_cfg.total_steps=50000" "train_cfg.n_checkpoint=50000" "train_cfg.n_updates=500" -------------------------------------------------------------------------------- /scripts/example_training/pretrained_model/bc_soft_body_rgbd.sh: -------------------------------------------------------------------------------- 1 | # Replace --work-dir, env_cfg.env_name, and replay_cfg.buffer_filenames when you run other environments 2 | 3 | python maniskill2_learn/apis/run_rl.py configs/brl/bc/rgbd_soft_body.py \ 4 | --work-dir ./logs/bc_excavate_rgbd --gpu-ids 0 \ 5 | --cfg-options "env_cfg.env_name=Excavate-v0" "env_cfg.obs_mode=rgbd" "env_cfg.n_points=1200" \ 6 | "env_cfg.control_mode=pd_joint_delta_pos" \ 7 | "replay_cfg.buffer_filenames=../ManiSkill2/demos/v0/soft_body/Excavate-v0/trajectory.none.pd_ee_delta_pose_rgbd.h5" \ 8 | "eval_cfg.num=100" "eval_cfg.save_traj=False" "eval_cfg.save_video=True" \ 9 | "train_cfg.n_eval=50000" "train_cfg.total_steps=50000" "train_cfg.n_checkpoint=50000" "train_cfg.n_updates=500" -------------------------------------------------------------------------------- /scripts/example_training/pretrained_model/dapg_pickcube_pointcloud.sh: -------------------------------------------------------------------------------- 1 | # Assuming 2 gpus each with 12GB memory; 2 | # if you have a GPU with more memory (e.g. 24GB), you can set --gpu-ids and --sim-gpu-ids to be the same; 3 | # if you only have one GPU with small memory, then you can set a smaller rollout_cfg.num_procs (e.g. =5) 4 | 5 | python maniskill2_learn/apis/run_rl.py configs/mfrl/dapg/maniskill2_pn.py \ 6 | --work-dir ./logs/dapg_pickcube_pointcloud --gpu-ids 0 --sim-gpu-ids 1 \ 7 | --cfg-options "env_cfg.env_name=PickCube-v0" "env_cfg.obs_mode=pointcloud" "env_cfg.n_points=1200" \ 8 | "rollout_cfg.num_procs=16" "env_cfg.reward_mode=dense" \ 9 | "env_cfg.control_mode=pd_ee_delta_pose" "env_cfg.obs_frame=ee" "env_cfg.n_goal_points=50" \ 10 | "agent_cfg.demo_replay_cfg.buffer_filenames=../ManiSkill2/demos/v0/rigid_body/PickCube-v0/trajectory.none.pd_ee_delta_pose_pointcloud.h5" \ 11 | "eval_cfg.num=100" "eval_cfg.save_traj=False" "eval_cfg.save_video=True" \ 12 | "train_cfg.total_steps=25000000" "train_cfg.n_checkpoint=5000000" -------------------------------------------------------------------------------- /scripts/example_training/pretrained_model/dapg_pickcube_rgbd.sh: -------------------------------------------------------------------------------- 1 | # Assuming 2 gpus each with 12GB memory; 2 | # if you have a GPU with more memory (e.g. 24GB), you can set --gpu-ids and --sim-gpu-ids to be the same; 3 | # if you only have one GPU with small memory, then you can set a smaller rollout_cfg.num_procs (e.g. =5) 4 | 5 | python maniskill2_learn/apis/run_rl.py configs/mfrl/dapg/maniskill2_rgbd.py \ 6 | --work-dir ./logs/dapg_pickcube_rgbd --gpu-ids 0 --sim-gpu-ids 1 \ 7 | --cfg-options "env_cfg.env_name=PickCube-v0" "env_cfg.obs_mode=rgbd" \ 8 | "env_cfg.control_mode=pd_ee_delta_pose" \ 9 | "rollout_cfg.num_procs=16" "env_cfg.reward_mode=dense" \ 10 | "agent_cfg.demo_replay_cfg.buffer_filenames=../ManiSkill2/demos/v0/rigid_body/PickCube-v0/trajectory.none.pd_ee_delta_pose_rgbd.h5" \ 11 | "eval_cfg.num=100" "eval_cfg.save_traj=False" "eval_cfg.save_video=True" \ 12 | "train_cfg.total_steps=25000000" "train_cfg.n_checkpoint=5000000" -------------------------------------------------------------------------------- /scripts/example_training/pretrained_model/dapg_picksingleycb_pointcloud.sh: -------------------------------------------------------------------------------- 1 | # Assuming 2 gpus each with 12GB memory; 2 | # if you have a GPU with more memory (e.g. 24GB), you can set --gpu-ids and --sim-gpu-ids to be the same; 3 | # if you only have one GPU with small memory, then you can set a smaller rollout_cfg.num_procs (e.g. =5) 4 | 5 | # Since the demo file is very large, we use dynamic loading to save memory 6 | python maniskill2_learn/apis/run_rl.py configs/mfrl/dapg/maniskill2_pn.py \ 7 | --work-dir ./logs/dapg_picksingleycb_pointcloud --gpu-ids 0 --sim-gpu-ids 1 \ 8 | --cfg-options "env_cfg.env_name=PickSingleYCB-v0" "env_cfg.obs_mode=pointcloud" "env_cfg.n_points=1200" \ 9 | "rollout_cfg.num_procs=16" "env_cfg.reward_mode=dense" \ 10 | "env_cfg.control_mode=pd_ee_delta_pose" "env_cfg.obs_frame=ee" "env_cfg.n_goal_points=50" \ 11 | "agent_cfg.demo_replay_cfg.capacity=20000" "agent_cfg.demo_replay_cfg.cache_size=20000" \ 12 | "agent_cfg.demo_replay_cfg.dynamic_loading=True" "agent_cfg.demo_replay_cfg.num_samples=-1" \ 13 | "agent_cfg.demo_replay_cfg.buffer_filenames=../ManiSkill2/demos/v0/rigid_body/PickSingleYCB-v0/trajectory_merged.none.pd_ee_delta_pose_pointcloud.h5" \ 14 | "eval_cfg.num=100" "eval_cfg.save_traj=False" "eval_cfg.save_video=True" \ 15 | "train_cfg.total_steps=25000000" "train_cfg.n_checkpoint=5000000" -------------------------------------------------------------------------------- /scripts/example_training/pretrained_model/dapg_picksingleycb_rgbd.sh: -------------------------------------------------------------------------------- 1 | # Assuming 2 gpus each with 12GB memory; 2 | # if you have a GPU with more memory (e.g. 24GB), you can set --gpu-ids and --sim-gpu-ids to be the same; 3 | # if you only have one GPU with small memory, then you can set a smaller rollout_cfg.num_procs (e.g. =5) 4 | 5 | # Since the demo file is very large, we use dynamic loading to save memory 6 | python maniskill2_learn/apis/run_rl.py configs/mfrl/dapg/maniskill2_rgbd.py \ 7 | --work-dir ./logs/dapg_picksingleycb_rgbd --gpu-ids 0 --sim-gpu-ids 1 \ 8 | --cfg-options "env_cfg.env_name=PickSingleYCB-v0" "env_cfg.obs_mode=rgbd" \ 9 | "rollout_cfg.num_procs=16" "env_cfg.reward_mode=dense" \ 10 | "env_cfg.control_mode=pd_ee_delta_pose" \ 11 | "agent_cfg.demo_replay_cfg.capacity=20000" "agent_cfg.demo_replay_cfg.cache_size=20000" \ 12 | "agent_cfg.demo_replay_cfg.dynamic_loading=True" "agent_cfg.demo_replay_cfg.num_samples=-1" \ 13 | "agent_cfg.demo_replay_cfg.buffer_filenames=../ManiSkill2/demos/v0/rigid_body/PickSingleYCB-v0/trajectory_merged.none.pd_ee_delta_pose_rgbd.h5" \ 14 | "eval_cfg.num=100" "eval_cfg.save_traj=False" "eval_cfg.save_video=True" \ 15 | "train_cfg.total_steps=25000000" "train_cfg.n_checkpoint=5000000" -------------------------------------------------------------------------------- /scripts/example_training/pretrained_model/dapg_stackcube_pointcloud.sh: -------------------------------------------------------------------------------- 1 | # Assuming 2 gpus each with 12GB memory; 2 | # if you have a GPU with more memory (e.g. 24GB), you can set --gpu-ids and --sim-gpu-ids to be the same; 3 | # if you only have one GPU with small memory, then you can set a smaller rollout_cfg.num_procs (e.g. =5) 4 | 5 | python maniskill2_learn/apis/run_rl.py configs/mfrl/dapg/maniskill2_pn.py \ 6 | --work-dir ./logs/dapg_stackcube_pointcloud --gpu-ids 0 --sim-gpu-ids 1 \ 7 | --cfg-options "env_cfg.env_name=StackCube-v0" "env_cfg.obs_mode=pointcloud" "env_cfg.n_points=1200" \ 8 | "rollout_cfg.num_procs=16" "env_cfg.reward_mode=dense" \ 9 | "env_cfg.control_mode=pd_ee_delta_pose" "env_cfg.obs_frame=ee" \ 10 | "agent_cfg.demo_replay_cfg.buffer_filenames=../ManiSkill2/demos/v0/rigid_body/StackCube-v0/trajectory.none.pd_ee_delta_pose_pointcloud.h5" \ 11 | "eval_cfg.num=100" "eval_cfg.save_traj=False" "eval_cfg.save_video=True" \ 12 | "train_cfg.total_steps=25000000" "train_cfg.n_checkpoint=5000000" -------------------------------------------------------------------------------- /scripts/example_training/pretrained_model/dapg_stackcube_rgbd.sh: -------------------------------------------------------------------------------- 1 | # Assuming 2 gpus each with 12GB memory; 2 | # if you have a GPU with more memory (e.g. 24GB), you can set --gpu-ids and --sim-gpu-ids to be the same; 3 | # if you only have one GPU with small memory, then you can set a smaller rollout_cfg.num_procs (e.g. =5) 4 | 5 | python maniskill2_learn/apis/run_rl.py configs/mfrl/dapg/maniskill2_rgbd.py \ 6 | --work-dir ./logs/dapg_stackcube_rgbd --gpu-ids 0 --sim-gpu-ids 1 \ 7 | --cfg-options "env_cfg.env_name=StackCube-v0" "env_cfg.obs_mode=rgbd" \ 8 | "env_cfg.control_mode=pd_ee_delta_pose" \ 9 | "rollout_cfg.num_procs=16" "env_cfg.reward_mode=dense" \ 10 | "agent_cfg.demo_replay_cfg.buffer_filenames=../ManiSkill2/demos/v0/rigid_body/StackCube-v0/trajectory.none.pd_ee_delta_pose_rgbd.h5" \ 11 | "eval_cfg.num=100" "eval_cfg.save_traj=False" "eval_cfg.save_video=True" \ 12 | "train_cfg.total_steps=25000000" "train_cfg.n_checkpoint=5000000" -------------------------------------------------------------------------------- /scripts/example_training/scratch_pointcloud_template/run_bc.sh: -------------------------------------------------------------------------------- 1 | python maniskill2_learn/apis/run_rl.py configs/brl/bc/pointnet.py \ 2 | --work-dir YOUR_LOGGING_DIRECTORY --gpu-ids 0 \ 3 | --cfg-options "env_cfg.env_name=PickCube-v0" "env_cfg.obs_mode=pointcloud" "env_cfg.n_points=1200" \ 4 | "env_cfg.control_mode=pd_joint_delta_pos" \ 5 | "replay_cfg.buffer_filenames=PATH_TO_POINT_CLOUD_DEMO" \ 6 | "eval_cfg.num=100" "eval_cfg.save_traj=False" "eval_cfg.save_video=True" \ 7 | "train_cfg.n_eval=50000" "train_cfg.total_steps=50000" "train_cfg.n_checkpoint=50000" "train_cfg.n_updates=500" -------------------------------------------------------------------------------- /scripts/example_training/scratch_pointcloud_template/run_dapg.sh: -------------------------------------------------------------------------------- 1 | python maniskill2_learn/apis/run_rl.py configs/mfrl/ppo/maniskill2_pn_dapg.py \ 2 | --work-dir YOUR_LOGGING_DIRECTORY --gpu-ids 0 \ 3 | --cfg-options "env_cfg.env_name=PegInsertionSide-v0" "env_cfg.obs_mode=pointcloud" "env_cfg.n_points=1200" \ 4 | "env_cfg.reward_mode=dense" "env_cfg.control_mode=pd_joint_delta_pos" \ 5 | "agent_cfg.demo_replay_cfg.buffer_filenames=PATH_TO_POINT_CLOUD_DEMO.h5" 6 | "eval_cfg.num=100" "eval_cfg.save_traj=False" "eval_cfg.save_video=True" 7 | # To manually evaluate the model, add --evaluation and --resume-from YOUR_LOGGING_DIRECTORY/models/SOME_CHECKPOINT.ckpt 8 | # to the above commands. 9 | 10 | # Using multiple GPUs will increase training speed; 11 | # Note that train_cfg.n_steps will also be multiplied by the number of gpus you use, so you may want to divide it by the number of gpus 12 | -------------------------------------------------------------------------------- /scripts/example_training/scratch_pointcloud_template/run_gail.sh: -------------------------------------------------------------------------------- 1 | # Point Cloud-based GAIL 2 | 3 | # ** Before you run it, check that demonstrations are converted with "--with-next" and rewards behave as intended. ** 4 | 5 | python maniskill2_learn/apis/run_rl.py configs/mfrl/gail/maniskill2_pn.py \ 6 | --gpu-ids 0 --work-dir YOUR_LOGGING_DIRECTORY --print-steps 16 \ 7 | --cfg-options "env_cfg.env_name=PlugCharger-v0" "env_cfg.control_mode=pd_joint_delta_pos" \ 8 | "env_cfg.reward_mode=dense" "replay_cfg.buffer_filenames=[PATH_TO_DEMO.h5]" \ 9 | "expert_replay_cfg.buffer_filenames=[PATH_TO_DEMO.h5]" 10 | 11 | # Using multiple GPUs will increase training speed; 12 | # Note that the effective batch size is multiplied by the number of gpus; large batch can be crucial for stabilizing GAIL training 13 | 14 | 15 | 16 | 17 | -------------------------------------------------------------------------------- /scripts/example_training/scratch_pointcloud_template/run_ppo.sh: -------------------------------------------------------------------------------- 1 | python maniskill2_learn/apis/run_rl.py configs/mfrl/ppo/maniskill2_pn.py \ 2 | --work-dir YOUR_LOGGING_DIRECTORY --gpu-ids 0 \ 3 | --cfg-options "env_cfg.env_name=PickCube-v0" "env_cfg.obs_mode=pointcloud" "env_cfg.n_points=1200" \ 4 | "env_cfg.reward_mode=dense" "env_cfg.control_mode=pd_joint_delta_pos" \ 5 | "eval_cfg.num=100" "eval_cfg.save_traj=False" "eval_cfg.save_video=True" 6 | 7 | # The above command does automatic evaluation after training. Alternatively, you can manually evaluate a model checkpoint 8 | # by appending --evaluation and --resume-from YOUR_LOGGING_DIRECTORY/models/SOME_CHECKPOINT.ckpt to the above commands. 9 | 10 | 11 | # Using multiple GPUs will increase training speed; 12 | # Note that train_cfg.n_steps will also be multiplied by the number of gpus you use, so you may want to divide it by the number of gpus -------------------------------------------------------------------------------- /scripts/example_training/scratch_pointcloud_template/run_sac.sh: -------------------------------------------------------------------------------- 1 | # Point Cloud-based SAC 2 | python maniskill2_learn/apis/run_rl.py configs/mfrl/sac/maniskill2_pn.py \ 3 | --gpu-ids 0 --work-dir YOUR_LOGGING_DIRECTORY --print-steps 16 \ 4 | --cfg-options "env_cfg.env_name=PegInsertionSide-v0" "env_cfg.control_mode=pd_joint_delta_pos" \ 5 | "env_cfg.reward_mode=dense" 6 | 7 | # Using multiple GPUs will increase training speed; 8 | # Note that the effective batch size is multiplied by the number of gpus; large batch can be crucial for stabilizing SAC training 9 | 10 | 11 | # State-based SAC for debugging purposes 12 | """ 13 | python maniskill2_learn/apis/run_rl.py configs/mfrl/sac/maniskill2_state.py \ 14 | --work-dir YOUR_LOGGING_DIRECTORY --gpu-ids 0 \ 15 | --cfg-options 'env_cfg.env_name=PegInsertionSide-v0' 'env_cfg.control_mode=pd_joint_delta_pos' \ 16 | 'env_cfg.reward_mode=dense' 17 | """ 18 | 19 | 20 | 21 | -------------------------------------------------------------------------------- /tools/merge_h5.py: -------------------------------------------------------------------------------- 1 | 2 | import argparse 3 | import os, numpy as np 4 | import os.path as osp 5 | import h5py 6 | import glob 7 | 8 | os.environ["D4RL_SUPPRESS_IMPORT_ERROR"] = "1" 9 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 10 | os.environ["MKL_NUM_THREADS"] = "1" 11 | os.environ["NUMEXPR_NUM_THREADS"] = "1" 12 | os.environ["OMP_NUM_THREADS"] = "1" 13 | 14 | from maniskill2_learn.utils.file import merge_h5_trajectory 15 | 16 | """ 17 | Example: 18 | python tools/merge_h5.py --source-dir /TODO/demos/PickSingleYCB-v0/ --pattern "trajectory_pcd" \ 19 | --output-file /TODO/demos/PickSingleYCB-v0/trajectory_pcd_all.h5 20 | """ 21 | 22 | parser = argparse.ArgumentParser(description="Merge h5 files that match {pattern}.h5 under a directory into a single file") 23 | parser.add_argument("--source-dir", type=str, default="") 24 | parser.add_argument("--pattern", type=str, default="") 25 | parser.add_argument("--output-file", type=str, default="") 26 | args = parser.parse_args() 27 | assert args.source_dir != "" and args.pattern != "" and args.output_file != "" 28 | 29 | files = glob.glob(f'{args.source_dir}/**/{args.pattern}.h5', recursive=True) 30 | print("Input files", files) 31 | try: 32 | os.remove(args.output_file) 33 | except: 34 | pass 35 | merge_h5_trajectory(files, args.output_file) 36 | -------------------------------------------------------------------------------- /tools/merge_trajectory.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | 4 | import h5py 5 | 6 | from mani_skill2.utils.io_utils import dump_json, load_json 7 | 8 | 9 | def merge_h5(output_path: str, traj_paths, recompute_id=True): 10 | print("Merge to", output_path) 11 | 12 | merged_h5_file = h5py.File(output_path, "w") 13 | merged_json_path = output_path.replace(".h5", ".json") 14 | merged_json_data = {"env_info": {}, "episodes": []} 15 | _env_info = None 16 | cnt = 0 17 | 18 | for traj_path in traj_paths: 19 | traj_path = str(traj_path) 20 | print("Merging", traj_path) 21 | 22 | h5_file = h5py.File(traj_path, "r") 23 | json_path = traj_path.replace(".h5", ".json") 24 | json_data = load_json(json_path) 25 | 26 | # Check env info 27 | env_info = json_data["env_info"] 28 | if _env_info is None: 29 | _env_info = env_info 30 | merged_json_data["env_info"] = _env_info 31 | else: 32 | assert str(env_info) == str(_env_info), f"Env info is not consistent in {json_path}" 33 | 34 | # Merge 35 | for ep in json_data["episodes"]: 36 | episode_id = ep["episode_id"] 37 | traj_id = f"traj_{episode_id}" 38 | 39 | # Copy h5 data 40 | if recompute_id: 41 | new_traj_id = f"traj_{cnt}" 42 | else: 43 | new_traj_id = traj_id 44 | 45 | assert new_traj_id not in merged_h5_file, f"{new_traj_id} is already in the merged_h5_file!" 46 | h5_file.copy(traj_id, merged_h5_file, new_traj_id) 47 | 48 | # Copy json data 49 | if recompute_id: 50 | ep["episode_id"] = cnt 51 | merged_json_data["episodes"].append(ep) 52 | 53 | cnt += 1 54 | 55 | h5_file.close() 56 | 57 | # Ignore commit info 58 | merged_h5_file.close() 59 | dump_json(merged_json_path, merged_json_data, indent=2) 60 | 61 | 62 | def main(): 63 | parser = argparse.ArgumentParser() 64 | parser.add_argument("-i", "--input-dirs", nargs="+") 65 | parser.add_argument("-o", "--output-path", type=str) 66 | parser.add_argument("-p", "--pattern", type=str, default="trajectory.h5") 67 | args = parser.parse_args() 68 | 69 | traj_paths = [] 70 | for input_dir in args.input_dirs: 71 | input_dir = Path(input_dir) 72 | traj_paths.extend(sorted(input_dir.rglob(args.pattern))) 73 | 74 | output_dir = Path(args.output_path).parent 75 | output_dir.mkdir(exist_ok=True, parents=True) 76 | 77 | merge_h5(args.output_path, traj_paths) 78 | 79 | 80 | if __name__ == "__main__": 81 | main() 82 | -------------------------------------------------------------------------------- /tools/permute_rgbd_agent_weights.py: -------------------------------------------------------------------------------- 1 | import torch, numpy as np 2 | import argparse 3 | 4 | parser = argparse.ArgumentParser(description="Permute RGBD agent weights due to visual key order changes") 5 | parser.add_argument("--model", type=str, required=True, help="Path to the model file") 6 | 7 | args = parser.parse_args() 8 | 9 | model = torch.load(args.model) 10 | for k in model['state_dict'].keys(): 11 | if k in ['actor.backbone.visual_nn.stem.weight', 'critic.values.0.backbone.visual_nn.stem.weight']: 12 | rgb_weights = model['state_dict'][k][:, :6, :, :] 13 | model['state_dict'][k][:, :6, :, :] = torch.cat([rgb_weights[:, 3:6, :, :], rgb_weights[:, :3, :, :]], dim=1) 14 | depth_weights = model['state_dict'][k][:, 6:8, :, :] 15 | model['state_dict'][k][:, 6:8, :, :] = torch.cat([depth_weights[:, 1:2, :, :], depth_weights[:, :1, :, :]], dim=1) 16 | torch.save(model, args.model.replace('.pth', '_permuted_visual.pth')) -------------------------------------------------------------------------------- /tools/shuffle_demo.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | import random 3 | import argparse 4 | import tqdm 5 | import json 6 | import copy 7 | from multiprocessing import Pool 8 | 9 | 10 | def copy_group(from_file, to_file, key_name, new_key_name): 11 | if new_key_name not in to_file.keys(): 12 | to_file.require_group(new_key_name) 13 | for key in from_file[key_name].keys(): 14 | new_data_key = key 15 | if "dist" in key or "str" in key: 16 | new_data_key = "_".join(key.split("_")[2:]) 17 | if isinstance(from_file[key_name][key], h5py.Group): 18 | copy_group(from_file[key_name], to_file[new_key_name], key, new_data_key) 19 | else: 20 | to_file[new_key_name].create_dataset( 21 | new_data_key, data=from_file[key_name][key] 22 | ) 23 | 24 | 25 | parser = argparse.ArgumentParser() 26 | parser.add_argument("--source-file", type=str) 27 | parser.add_argument("--target-file", type=str) 28 | args = parser.parse_args() 29 | 30 | mapping = {} 31 | source_file = h5py.File(args.source_file, "r") 32 | target_file = h5py.File(args.target_file, "w") 33 | order_list = list(range(len(list(source_file.keys())))) 34 | random.shuffle(order_list) 35 | for index, ori_index in tqdm.tqdm(enumerate(order_list), total=len(order_list)): 36 | old_key = "traj_" + str(ori_index) 37 | new_key = "traj_" + str(index) 38 | mapping[old_key] = new_key 39 | copy_group(source_file, target_file, old_key, new_key) 40 | print(mapping) 41 | source_file.close() 42 | target_file.close() 43 | --------------------------------------------------------------------------------