Add files using upload-large-folder tool

.gitattributes

@@ -34,3 +34,56 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 Piper_ros_private-ros-noetic/asserts/pictures/piper_rviz.jpg filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/utils.cpp.o filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera_node.dir/src/main.cpp.o filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/base_link.STL filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/x64/PS1080Console filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/arm/OpenNI2/Drivers/libOniFile.so filter=lfs diff=lfs merge=lfs -text
+camera_ws/devel/lib/astra_camera/astra_camera_node filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/x64/NiViewer filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/x64/OpenNI2/Drivers/libOniFile.so filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/arm/PS1080Console filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/uvc_camera_driver.cpp.o filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/d2c_viewer.cpp.o filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/ros_service.cpp.o filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/ob_camera_info.cpp.o filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_description/meshes/base_link.STL filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/cfg/1.png filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/realsense-ros/realsense2_camera/CMakeFiles/realsense2_camera.dir/src/base_realsense_node.cpp.o filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/arm64/PS1080Console filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/x64/OpenNI2/Drivers/liborbbec.so filter=lfs diff=lfs merge=lfs -text
+collect_data/docs/episode_0_qpos.png filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/realsense-ros/realsense2_camera/CMakeFiles/realsense2_camera.dir/src/realsense_node_factory.cpp.o filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/point_cloud_proc/point_cloud_xyz.cpp.o filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link4.STL filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_description/meshes/link4.STL filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/realsense-ros/realsense2_description/meshes/plug.stl filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link2.STL filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/ob_camera_node_factory.cpp.o filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/realsense-ros/realsense2_description/meshes/d455.stl filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/realsense-ros/realsense2_camera/CMakeFiles/realsense2_camera.dir/src/t265_realsense_node.cpp.o filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/ob_camera_node.cpp.o filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/x64/libOpenNI2.so filter=lfs diff=lfs merge=lfs -text
+camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/point_cloud_proc/point_cloud_xyzrgb.cpp.o filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_description/meshes/link1.STL filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_description/meshes/link5.STL filter=lfs diff=lfs merge=lfs -text
+camera_ws/devel/lib/libastra_camera.so filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/arm/libOpenNI2.so filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/arm64/NiViewer filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/arm64/libOpenNI2.so filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link1.STL filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_description/meshes/link2.STL filter=lfs diff=lfs merge=lfs -text
+collect_data/docs/1.png filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link5.STL filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/arm/OpenNI2/Drivers/liborbbec.so filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/arm64/OpenNI2/Drivers/liborbbec.so filter=lfs diff=lfs merge=lfs -text
+camera_ws/devel/lib/librealsense2_camera.so filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_description/meshes/link3.STL filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link3.STL filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_description/meshes/link6.STL filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/ros_astra_camera/include/openni2_redist/arm64/OpenNI2/Drivers/libOniFile.so filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/realsense-ros/realsense2_description/meshes/d435.dae filter=lfs diff=lfs merge=lfs -text
+Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link6.STL filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/realsense-ros/realsense2_description/meshes/d415.stl filter=lfs diff=lfs merge=lfs -text
+camera_ws/src/realsense-ros/realsense2_description/meshes/l515.dae filter=lfs diff=lfs merge=lfs -text
+collect_data/docs/1.gif filter=lfs diff=lfs merge=lfs -text

Piper_ros_private-ros-noetic/src/piper_description/meshes/base_link.STL

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

Piper_ros_private-ros-noetic/src/piper_description/meshes/link1.STL

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

Piper_ros_private-ros-noetic/src/piper_description/meshes/link2.STL

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

Piper_ros_private-ros-noetic/src/piper_description/meshes/link3.STL

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

Piper_ros_private-ros-noetic/src/piper_description/meshes/link4.STL

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

Piper_ros_private-ros-noetic/src/piper_description/meshes/link5.STL

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

Piper_ros_private-ros-noetic/src/piper_description/meshes/link6.STL

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

Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/base_link.STL

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

Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link1.STL

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

Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link2.STL

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

Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link3.STL

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

Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link4.STL

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

Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link5.STL

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

Piper_ros_private-ros-noetic/src/piper_mujoco/mujoco_description/meshes_mujoco/link6.STL

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

aloha-devel/act/train.sh

@@ -0,0 +1,20 @@
+num_epochs=600
+batch_size=48
+num_episodes=80
+ROOT=/inspire/hdd/ws-f4d69b29-e0a5-44e6-bd92-acf4de9990f0/public-project/chengdongzhou-240108390137/vla_projects/cobot_magic
+train_dir=$ROOT/tranin_dir
+pretrain_ckpt=$ROOT/policy_best.ckpt
+ws_path=$(pwd)
+dataset_dir=/inspire/ssd/ws-f4d69b29-e0a5-44e6-bd92-acf4de9990f0/public-project/public/aloha_group/blue_new
+# echo "$pretrain_ckpt"
+# echo "$train_dir"
+# echo $(pwd)
+
+cd $ws_path
+
+python act/train.py --dataset $dataset_dir --ckpt_dir $train_dir/pretrain --batch_size $batch_size  --num_epochs $num_epochs --num_episodes $num_episodes --pretrain_ckpt $pretrain_ckpt
+python act/train.py --dataset /media/lin/T7/data0314/ --ckpt_dir $train_dir/no_pretrain --batch_size $batch_size --num_epochs $num_epochs --num_episodes $num_episodes
+
+
+
+python act/train.py --dataset $dataset_dir --ckpt_dir $train_dir/no_pretrain --batch_size $batch_size  --num_epochs $num_epochs --num_episodes $num_episodes

aloha-devel/robomimic/scripts/config_gen/bc_xfmr_gen.py

@@ -0,0 +1,169 @@
+from robomimic.scripts.config_gen.helper import *
+
+def make_generator_helper(args):
+    algo_name_short = "bc_xfmr"
+
+    generator = get_generator(
+        algo_name="bc",
+        config_file=os.path.join(base_path, 'robomimic/exps/templates/bc_transformer.json'),
+        args=args,
+        algo_name_short=algo_name_short,
+        pt=True,
+    )
+    if args.ckpt_mode is None:
+        args.ckpt_mode = "off"
+
+    if args.env == "r2d2":
+        generator.add_param(
+            key="train.data",
+            name="ds",
+            group=2,
+            values=[
+                [{"path": p} for p in scan_datasets("~/Downloads/example_pen_in_cup", postfix="trajectory_im128.h5")],
+            ],
+            value_names=[
+                "pen-in-cup",
+            ],
+        )
+        generator.add_param(
+            key="observation.modalities.obs.rgb",
+            name="cams",
+            group=130,
+            values=[
+                # ["camera/image/hand_camera_left_image"],
+                ["camera/image/hand_camera_left_image", "camera/image/varied_camera_1_left_image", "camera/image/varied_camera_2_left_image"],
+            ],
+            value_names=[
+                # "wrist",
+                "3cams",
+            ]
+        )
+    elif args.env == "kitchen":
+        generator.add_param(
+            key="train.data",
+            name="ds",
+            group=2,
+            values=[
+                [
+                    {
+                        "path": "/home/aaronl/tmp/v2_demos/KitchenPnPCounterToCab_im84.hdf5",
+                        "filter_key": "100_demos",
+                        "lang": "pick and place the object from the counter to the cabinet",
+                    },
+                    {
+                        "path": "/home/aaronl/tmp/v2_demos/KitchenPnPCabToCounter_im84.hdf5",
+                        "filter_key": "100_demos",
+                        "lang": "pick and place the object from the cabinet to the counter",
+                    },
+                ],
+            ],
+            value_names=[
+                "pnp-multi-task"
+            ],
+        )
+        generator.add_param(
+            key="algo.language_conditioned",
+            name="langcond",
+            group=145892,
+            values=[
+                True,
+                False,
+            ],
+        )
+    elif args.env == "square":
+        generator.add_param(
+            key="train.data",
+            name="ds",
+            group=2,
+            values=[
+                [
+                    {"path": "~/datasets/square/ph/square_ph_abs_tmp.hdf5"}, # replace with your own path
+                ],
+            ],
+            value_names=[
+                "square",
+            ],
+        )
+    else:
+        raise ValueError
+
+    # change default settings: predict 10 steps into future
+    generator.add_param(
+        key="algo.transformer.pred_future_acs",
+        name="predfuture",
+        group=1,
+        values=[
+            True,
+            # False,
+        ],
+        hidename=True,
+    )
+    generator.add_param(
+        key="algo.transformer.supervise_all_steps",
+        name="supallsteps",
+        group=1,
+        values=[
+            True,
+            # False,
+        ],
+        hidename=True,
+    )
+    generator.add_param(
+        key="algo.transformer.causal",
+        name="causal",
+        group=1,
+        values=[
+            False,
+            # True,
+        ],
+        hidename=True,
+    )
+    generator.add_param(
+        key="train.seq_length",
+        name="",
+        group=-1,
+        values=[10],
+        hidename=True,
+    )
+
+    generator.add_param(
+        key="algo.gmm.min_std",
+        name="mindstd",
+        group=271314,
+        values=[
+            0.03,
+            #0.0001,
+        ],
+        hidename=True,
+    )
+    generator.add_param(
+        key="train.max_grad_norm",
+        name="maxgradnorm",
+        group=18371,
+        values=[
+            # None,
+            100.0,
+        ],
+        hidename=True,
+    )
+    
+    generator.add_param(
+        key="train.output_dir",
+        name="",
+        group=-1,
+        values=[
+            "~/expdata/{env}/{mod}/{algo_name_short}".format(
+                env=args.env,
+                mod=args.mod,
+                algo_name_short=algo_name_short,
+            )
+        ],
+    )
+
+    return generator
+
+if __name__ == "__main__":
+    parser = get_argparser()
+
+    args = parser.parse_args()
+    make_generator(args, make_generator_helper)

aloha-devel/robomimic/scripts/config_gen/diffusion_gen.py

@@ -0,0 +1,263 @@
+from robomimic.scripts.config_gen.helper import *
+
+def make_generator_helper(args):
+    algo_name_short = "diffusion_policy"
+
+    generator = get_generator(
+        algo_name="diffusion_policy",
+        config_file=os.path.join(base_path, 'robomimic/exps/templates/diffusion_policy.json'),
+        args=args,
+        algo_name_short=algo_name_short,
+        pt=True,
+    )
+    if args.ckpt_mode is None:
+        args.ckpt_mode = "off"
+
+    generator.add_param(
+        key="train.num_data_workers",
+        name="",
+        group=-1,
+        values=[8],
+    )
+
+    generator.add_param(
+        key="train.num_epochs",
+        name="",
+        group=-1,
+        values=[1000],
+    )
+
+    # use ddim by default
+    generator.add_param(
+        key="algo.ddim.enabled",
+        name="ddim",
+        group=1001,
+        values=[
+            True,
+            # False,
+        ],
+        hidename=True,
+    )
+    generator.add_param(
+        key="algo.ddpm.enabled",
+        name="ddpm",
+        group=1001,
+        values=[
+            False,
+            # True,
+        ],
+        hidename=True,
+    )
+
+    if args.env == "r2d2":
+        generator.add_param(
+            key="train.data",
+            name="ds",
+            group=2,
+            values=[
+                [{"path": p, "lang": "put the pen in the cup"} for p in scan_datasets("~/Downloads/example_pen_in_cup", postfix="trajectory_im128.h5")],
+            ],
+            value_names=[
+                "pen-in-cup",
+            ],
+        )
+        generator.add_param(
+            key="train.action_keys",
+            name="ac_keys",
+            group=-1,
+            values=[
+                [
+                    "action/abs_pos",
+                    "action/abs_rot_6d",
+                    "action/gripper_position",
+                ],
+            ],
+            value_names=[
+                "abs",
+            ],
+            hidename=True,
+        )
+        generator.add_param(
+            key="observation.modalities.obs.rgb",
+            name="cams",
+            group=130,
+            values=[
+                # ["camera/image/hand_camera_left_image"],
+                # ["camera/image/hand_camera_left_image", "camera/image/hand_camera_right_image"],
+                ["camera/image/hand_camera_left_image", "camera/image/varied_camera_1_left_image", "camera/image/varied_camera_2_left_image"],
+                # [
+                #     "camera/image/hand_camera_left_image", "camera/image/hand_camera_right_image",
+                #     "camera/image/varied_camera_1_left_image", "camera/image/varied_camera_1_right_image",
+                #     "camera/image/varied_camera_2_left_image", "camera/image/varied_camera_2_right_image",
+                # ],
+            ],
+            value_names=[
+                # "wrist",
+                # "wrist-stereo",
+                "3cams",
+                # "3cams-stereo",
+            ]
+        )
+
+        generator.add_param(
+            key="observation.modalities.obs.low_dim",
+            name="ldkeys",
+            group=2498,
+            values=[
+                ["robot_state/cartesian_position", "robot_state/gripper_position"],
+                # [
+                #     "robot_state/cartesian_position", "robot_state/gripper_position",
+                #     "camera/extrinsics/hand_camera_left", "camera/extrinsics/hand_camera_left_gripper_offset",
+                #     "camera/extrinsics/hand_camera_right", "camera/extrinsics/hand_camera_right_gripper_offset",
+                #     "camera/extrinsics/varied_camera_1_left", "camera/extrinsics/varied_camera_1_right",
+                #     "camera/extrinsics/varied_camera_2_left", "camera/extrinsics/varied_camera_2_right",
+                # ]
+            ],
+            value_names=[
+                "proprio",
+                # "proprio-extrinsics",
+            ]
+        )
+
+        generator.add_param(
+            key="observation.encoder.rgb.core_kwargs.backbone_class",
+            name="backbone",
+            group=1234,
+            values=[
+                "ResNet18Conv",
+                # "ResNet50Conv",
+            ],
+        )
+        generator.add_param(
+            key="observation.encoder.rgb.core_kwargs.feature_dimension",
+            name="visdim",
+            group=1234,
+            values=[
+                64,
+                # 512,
+            ],
+        )
+
+        generator.add_param(
+            key="algo.language_conditioned",
+            name="langcond",
+            group=145892,
+            values=[
+                True,
+                False,
+            ],
+        )
+
+    elif args.env == "kitchen":
+        generator.add_param(
+            key="train.data",
+            name="ds",
+            group=2,
+            values=[
+                # [{"path": "~/datasets/kitchen/prior/human_demos/pnp_table_to_cab/bowls/20230816_im84.hdf5", "filter_key": "100_demos"}],
+                [{"path": "~/datasets/kitchen/prior/human_demos/pnp_table_to_cab/all/20230806_im84.hdf5", "filter_key": "100_demos"}],
+                # [{"path": "~/datasets/kitchen/prior/mimicgen/pnp_table_to_cab/viraj_mg_2023-08-10-20-31-14/demo_im84.hdf5", "filter_key": "100_demos"}],
+                # [{"path": "~/datasets/kitchen/prior/mimicgen/pnp_table_to_cab/viraj_mg_2023-08-10-20-31-14/demo_im84.hdf5", "filter_key": "1000_demos"}],
+            ],
+            value_names=[
+                # "bowls-human-100",
+                "human-100",
+                # "mg-100",
+                # "mg-1000",
+            ],
+        )
+        
+        # update env config to use absolute action control
+        generator.add_param(
+            key="experiment.env_meta_update_dict",
+            name="",
+            group=-1,
+            values=[
+                {"env_kwargs": {"controller_configs": {"control_delta": False}}}
+            ],
+        )
+
+        generator.add_param(
+            key="train.action_keys",
+            name="ac_keys",
+            group=-1,
+            values=[
+                [
+                    "action_dict/abs_pos",
+                    "action_dict/abs_rot_6d",
+                    "action_dict/gripper",
+                    "action_dict/base_mode",
+                    # "actions",
+                ],
+            ],
+            value_names=[
+                "abs",
+            ],
+            hidename=True,
+        )
+    elif args.env == "square":
+        generator.add_param(
+            key="train.data",
+            name="ds",
+            group=2,
+            values=[
+                [
+                    {"path": "~/datasets/square/ph/square_ph_abs_tmp.hdf5"}, # replace with your own path
+                ],
+            ],
+            value_names=[
+                "square",
+            ],
+        )
+
+        # update env config to use absolute action control
+        generator.add_param(
+            key="experiment.env_meta_update_dict",
+            name="",
+            group=-1,
+            values=[
+                {"env_kwargs": {"controller_configs": {"control_delta": False}}}
+            ],
+        )
+        
+        generator.add_param(
+            key="train.action_keys",
+            name="ac_keys",
+            group=-1,
+            values=[
+                [
+                    "action_dict/abs_pos",
+                    "action_dict/abs_rot_6d",
+                    "action_dict/gripper",
+                    # "actions",
+                ],
+            ],
+            value_names=[
+                "abs",
+            ],
+        )
+
+
+    else:
+        raise ValueError
+    
+    generator.add_param(
+        key="train.output_dir",
+        name="",
+        group=-1,
+        values=[
+            "~/expdata/{env}/{mod}/{algo_name_short}".format(
+                env=args.env,
+                mod=args.mod,
+                algo_name_short=algo_name_short,
+            )
+        ],
+    )
+
+    return generator
+
+if __name__ == "__main__":
+    parser = get_argparser()
+
+    args = parser.parse_args()
+    make_generator(args, make_generator_helper)

aloha-devel/robomimic/scripts/conversion/convert_roboturk_pilot.py

@@ -0,0 +1,192 @@
+"""
+Helper script to convert the RoboTurk Pilot datasets (https://roboturk.stanford.edu/dataset_sim.html)
+into a format compatible with this repository. It will also create some useful filter keys
+in the file (e.g. training, validation, and fastest n trajectories). Prior work
+(https://arxiv.org/abs/1911.05321) has found this useful (for example, training on the 
+fastest 225 demonstrations for bins-Can).
+
+Direct download link for dataset: http://cvgl.stanford.edu/projects/roboturk/RoboTurkPilot.zip
+
+Args:
+    folder (str): path to a folder containing a demo.hdf5 and a models directory containing
+        mujoco xml files. For example, RoboTurkPilot/bins-Can.
+
+    n (int): creates a filter key corresponding to the n fastest trajectories. Defaults to 225.
+
+Example usage:
+
+    python convert_roboturk_pilot.py --folder /path/to/RoboTurkPilot/bins-Can --n 225
+"""
+
+import os
+import h5py
+import json
+import argparse
+import numpy as np
+from tqdm import tqdm
+
+import robomimic
+import robomimic.envs.env_base as EB
+from robomimic.utils.file_utils import create_hdf5_filter_key
+from robomimic.scripts.split_train_val import split_train_val_from_hdf5
+
+
+def convert_rt_pilot_hdf5(ref_folder):
+    """
+    Uses the reference demo hdf5 to write a new converted hdf5 compatible with
+    the repository.
+
+    Args:
+        ref_folder (str): path to a folder containing a demo.hdf5 and a models directory containing
+            mujoco xml files.
+    """
+    hdf5_path = os.path.join(ref_folder, "demo.hdf5")
+    new_path = os.path.join(ref_folder, "demo_new.hdf5")
+
+    f = h5py.File(hdf5_path, "r")
+    f_new = h5py.File(new_path, "w")
+    f_new_grp = f_new.create_group("data")
+
+    # sorted list of demonstrations by demo number
+    demos = list(f["data"].keys())
+    inds = np.argsort([int(elem[5:]) for elem in demos])
+    demos = [demos[i] for i in inds]
+
+    # write each demo
+    num_samples_arr = []
+    for demo_id in tqdm(range(len(demos))):
+        ep = demos[demo_id]
+
+        # create group for this demonstration
+        ep_data_grp = f_new_grp.create_group(ep)
+
+        # copy states over
+        states = f["data/{}/states".format(ep)][()]
+        ep_data_grp.create_dataset("states", data=np.array(states))
+
+        # concat jvels and gripper actions to form full actions
+        jvels = f["data/{}/joint_velocities".format(ep)][()]
+        gripper_acts = f["data/{}/gripper_actuations".format(ep)][()]
+        actions = np.concatenate([jvels, gripper_acts], axis=1)
+
+        # IMPORTANT: clip actions to -1, 1, since this is expected by the codebase
+        actions = np.clip(actions, -1., 1.)
+        ep_data_grp.create_dataset("actions", data=actions)
+
+        # store model xml directly in the new hdf5 file
+        model_path = os.path.join(ref_folder, "models", f["data/{}".format(ep)].attrs["model_file"])
+        f_model = open(model_path, "r")
+        model_xml = f_model.read()
+        f_model.close()
+        ep_data_grp.attrs["model_file"] = model_xml
+
+        # store num samples for this ep
+        num_samples = actions.shape[0]
+        ep_data_grp.attrs["num_samples"] = num_samples # number of transitions in this episode
+        num_samples_arr.append(num_samples)
+
+    # write dataset attributes (metadata)
+    f_new_grp.attrs["total"] = np.sum(num_samples_arr)
+
+    # construct and save env metadata
+    env_meta = dict()
+    env_meta["type"] = EB.EnvType.ROBOSUITE_TYPE
+    env_meta["env_name"] = (f["data"].attrs["env"] + "Teleop")
+    # hardcode robosuite v0.3 args
+    robosuite_args = {
+        "has_renderer": False,
+        "has_offscreen_renderer": False,
+        "ignore_done": True,
+        "use_object_obs": True,
+        "use_camera_obs": False,
+        "camera_depth": False,
+        "camera_height": 84,
+        "camera_width": 84,
+        "camera_name": "agentview",
+        "gripper_visualization": False,
+        "reward_shaping": False,
+        "control_freq": 100,
+    }
+    env_meta["env_kwargs"] = robosuite_args
+    f_new_grp.attrs["env_args"] = json.dumps(env_meta, indent=4) # environment info
+
+    print("\n====== Added env meta ======")
+    print(f_new_grp.attrs["env_args"])
+
+    f.close()
+    f_new.close()
+
+    # back up the old dataset, and replace with new dataset
+    os.rename(hdf5_path, os.path.join(ref_folder, "demo_bak.hdf5"))
+    os.rename(new_path, hdf5_path)
+
+
+def split_fastest_from_hdf5(hdf5_path, n):
+    """
+    Creates filter key for fastest N trajectories, named
+    "fastest_{}".format(n).
+
+    Args:
+        hdf5_path (str): path to the hdf5 file
+
+        n (int): fastest n demos to create filter key for
+    """
+
+    # retrieve fastest n demos
+    f = h5py.File(hdf5_path, "r")
+    demos = sorted(list(f["data"].keys()))
+    traj_lengths = []
+    for ep in demos:
+        traj_lengths.append(f["data/{}/actions".format(ep)].shape[0])
+    inds = np.argsort(traj_lengths)[:n]
+    filtered_demos = [demos[i] for i in inds]
+    f.close()
+
+    # create filter key
+    name = "fastest_{}".format(n)
+    lengths = create_hdf5_filter_key(hdf5_path=hdf5_path, demo_keys=filtered_demos, key_name=name)
+
+    print("Total number of samples in fastest {} demos: {}".format(n, np.sum(lengths)))
+    print("Average number of samples in fastest {} demos: {}".format(n, np.mean(lengths)))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--folder",
+        type=str,
+        help="path to a folder containing a demo.hdf5 and a models directory containing \
+            mujoco xml files. For example, RoboTurkPilot/bins-Can.",
+    )
+    parser.add_argument(
+        "--n",
+        type=int,
+        default=225,
+        help="creates a filter key corresponding to the n fastest trajectories. Defaults to 225.",
+    )
+    args = parser.parse_args()
+
+    # convert hdf5
+    convert_rt_pilot_hdf5(ref_folder=args.folder)
+
+    # create 90-10 train-validation split in the dataset
+    print("\nCreating 90-10 train-validation split...\n")
+    hdf5_path = os.path.join(args.folder, "demo.hdf5")
+    split_train_val_from_hdf5(hdf5_path=hdf5_path, val_ratio=0.1)
+
+    print("\nCreating filter key for fastest {} trajectories...".format(args.n))
+    split_fastest_from_hdf5(hdf5_path=hdf5_path, n=args.n)
+
+    print("\nCreating 90-10 train-validation split for fastest {} trajectories...".format(args.n))
+    split_train_val_from_hdf5(hdf5_path=hdf5_path, val_ratio=0.1, filter_key="fastest_{}".format(args.n))
+
+    print(
+        "\nWARNING: new dataset has replaced old one in demo.hdf5 file. "
+        "The old dataset file has been moved to demo_bak.hdf5"
+    )
+
+    print(
+        "\nNOTE: the new dataset also contains a fastest_{} filter key, for an easy way "
+        "to train on the fastest trajectories. Just set config.train.hdf5_filter to train on this "
+        "subset. A common choice is 225 when training on the bins-Can dataset.\n".format(args.n)
+    )

aloha-devel/robomimic/scripts/conversion/convert_to_robosuite_v141.py

@@ -0,0 +1,156 @@
+import h5py
+import json
+import argparse
+import os
+from shutil import copyfile
+import robosuite
+import xml.etree.ElementTree as ET
+
+import robomimic.utils.obs_utils as ObsUtils
+import robomimic.utils.env_utils as EnvUtils
+import robomimic.utils.file_utils as FileUtils
+
+from robosuite.utils.mjcf_utils import find_elements
+
+def replace_elem(parent, old_elem, new_elem):
+    """
+    code adapted from https://stackoverflow.com/a/20931505
+    """
+    parent_index = list(parent).index(old_elem)
+    parent.remove(old_elem)
+    parent.insert(parent_index, new_elem)
+
+def convert_xml(old_xml_str, env_name, env):
+    """
+    Postprocess xml string generated by robosuite to be compatible with robosuite v1.3
+    This script should not the xml string if it was already generated using robosuite v1.3
+    Args:
+        xml_str (str): xml string to process (from robosuite v1.2)
+    """
+    
+    if env_name in ["PickPlaceCan", "NutAssemblySquare", "ToolHang"]:
+        xml_str = env.env.sim.model.get_xml()
+    elif env_name == "Lift":
+        xml_str = env.env.sim.model.get_xml()
+        # replace the cube_g0 and cube_g0_vis with elements in old_xml_str
+        old_et = ET.ElementTree(ET.fromstring(old_xml_str)).getroot()
+        new_et = ET.ElementTree(ET.fromstring(xml_str)).getroot()
+
+        cube_new = find_elements(
+            root=new_et,
+            tags="body",
+            attribs={"name": "cube_main"},
+            return_first=True
+        )
+
+        cube_old = find_elements(
+            root=old_et,
+            tags="body",
+            attribs={"name": "cube_main"},
+            return_first=True
+        )
+
+        worldbody_new = find_elements(
+            root=new_et,
+            tags="worldbody",
+            return_first=True
+        )
+
+        replace_elem(worldbody_new, cube_new, cube_old)
+
+        xml_str = ET.tostring(new_et, encoding="utf8").decode("utf8")
+    elif env_name == "TwoArmTransport":
+        xml_str = env.env.sim.model.get_xml()
+        # replace the cube_g0 and cube_g0_vis with elements in old_xml_str
+        old_et = ET.ElementTree(ET.fromstring(old_xml_str)).getroot()
+        new_et = ET.ElementTree(ET.fromstring(xml_str)).getroot()
+
+        worldbody_new = find_elements(
+            root=new_et,
+            tags="worldbody",
+            return_first=True
+        )
+        for bname in [
+            "payload_root",
+            
+            ### ignore all these other following assets (makes playback worse for some reason...)
+            # "trash_main",
+            # "transport_start_bin_root", "transport_target_bin_root",
+            # "transport_trash_bin_root", "transport_start_bin_lid_root"
+        ]:
+            body_new = find_elements(
+                root=new_et,
+                tags="body",
+                attribs={"name": bname},
+                return_first=True
+            )
+
+            body_old = find_elements(
+                root=old_et,
+                tags="body",
+                attribs={"name": bname},
+                return_first=True
+            )
+
+        replace_elem(worldbody_new, body_new, body_old)
+
+        xml_str = ET.tostring(new_et, encoding="utf8").decode("utf8")
+
+    return xml_str
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        help="path to input hdf5 dataset",
+    )
+    parser.add_argument(
+        "--output_dataset",
+        type=str,
+        help="path to output hdf5 dataset",
+    )
+    args = parser.parse_args()
+    
+    args.dataset = os.path.expanduser(args.dataset)
+    args.output_dataset = os.path.expanduser(args.output_dataset)
+    
+    assert args.output_dataset != args.dataset
+    assert robosuite.__version__ == '1.4.1'
+    
+    copyfile(args.dataset, args.output_dataset)
+    
+    f = h5py.File(args.output_dataset, "r+")
+
+    env_args = json.loads(f["data"].attrs["env_args"])
+    env_name = env_args["env_name"]
+
+    env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=args.dataset)
+    env_type = EnvUtils.get_env_type(env_meta=env_meta)
+
+    # need to make sure ObsUtils knows which observations are images, but it doesn't matter 
+    # for playback since observations are unused. Pass a dummy spec here.
+    dummy_spec = dict(
+        obs=dict(
+                low_dim=["robot0_eef_pos"],
+                rgb=[],
+            ),
+    )
+    ObsUtils.initialize_obs_utils_with_obs_specs(obs_modality_specs=dummy_spec)
+
+    env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=args.dataset)
+    env = EnvUtils.create_env_from_metadata(env_meta=env_meta, render=False, render_offscreen=True)
+    env.reset()
+
+    for demo_key in list(f["data"].keys()):
+        ep_data_grp = f["data/{}".format(demo_key)]
+        model_file = ep_data_grp.attrs["model_file"]
+        
+        coverted_model_file = convert_xml(model_file, env_name, env)
+        ep_data_grp.attrs["model_file"] = coverted_model_file
+        
+    env_args = json.loads(f["data"].attrs["env_args"])
+    env_args["env_version"] = robosuite.__version__ 
+    f["data"].attrs["env_args"] = json.dumps(env_args, indent=4)
+
+    f.close()

aloha-devel/robomimic/scripts/conversion/extract_action_dict.py

@@ -0,0 +1,81 @@
+import argparse
+import pathlib
+import sys
+import tqdm
+import h5py
+import numpy as np
+import torch
+import os
+
+import robomimic.utils.torch_utils as TorchUtils
+
+def extract_action_dict(dataset):
+    # find files
+    f = h5py.File(os.path.expanduser(dataset), mode="r+")
+
+    SPECS = [
+        dict(
+            key="actions",
+            is_absolute=False,
+        ),
+        dict(
+            key="actions_abs",
+            is_absolute=True,
+        )
+    ]
+
+    # execute
+    for spec in SPECS:
+        input_action_key = spec["key"]
+        is_absolute = spec["is_absolute"]
+
+        if is_absolute:
+            prefix = "abs_"
+        else:
+            prefix = "rel_"
+
+        for demo in f["data"].values():
+            in_action = demo[str(input_action_key)][:]
+            in_pos = in_action[:,:3].astype(np.float32)
+            in_rot = in_action[:,3:6].astype(np.float32)
+            in_grip = in_action[:,6:7].astype(np.float32)
+            
+            rot_6d = TorchUtils.axis_angle_to_rot_6d(
+                axis_angle=torch.from_numpy(in_rot)
+            )
+            rot_6d = rot_6d.numpy().astype(np.float32) # convert to numpy
+            
+            this_action_dict = {
+                prefix + "pos": in_pos,
+                prefix + "rot_axis_angle": in_rot,
+                prefix + "rot_6d": rot_6d,
+                "gripper": in_grip
+            }
+
+            # special case: 8 dim actions mean there is a mobile base mode in the action space
+            if in_action.shape[1] == 8:
+                this_action_dict["base_mode"] = in_action[:,7:8].astype(np.float32)
+
+            action_dict_group = demo.require_group("action_dict")
+            for key, data in this_action_dict.items():
+                if key in action_dict_group:
+                    del action_dict_group[key]
+                action_dict_group.create_dataset(key, data=data)
+
+    f.close()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        required=True
+    )
+    
+    args = parser.parse_args()
+    extract_action_dict(args.dataset)
+    
+if __name__ == "__main__":
+    main()

aloha-devel/robomimic/scripts/conversion/robosuite_add_absolute_actions.py

@@ -0,0 +1,306 @@
+import multiprocessing
+import os
+import pathlib
+import h5py
+from tqdm import tqdm
+import collections
+import pickle
+import argparse
+import numpy as np
+import copy
+
+import h5py
+import robomimic.utils.obs_utils as ObsUtils
+import robomimic.utils.file_utils as FileUtils
+import robomimic.utils.env_utils as EnvUtils
+from scipy.spatial.transform import Rotation
+
+from robomimic.config import config_factory
+
+"""
+copied/adapted from https://github.com/columbia-ai-robotics/diffusion_policy/blob/main/diffusion_policy/common/robomimic_util.py
+"""
+class RobomimicAbsoluteActionConverter:
+    def __init__(self, dataset_path, algo_name='bc'):
+        # default BC config
+        config = config_factory(algo_name=algo_name)
+
+        # read config to set up metadata for observation modalities (e.g. detecting rgb observations)
+        # must ran before create dataset
+        ObsUtils.initialize_obs_utils_with_config(config)
+
+        env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path)
+        abs_env_meta = copy.deepcopy(env_meta)
+        abs_env_meta['env_kwargs']['controller_configs']['control_delta'] = False
+
+        env = EnvUtils.create_env_from_metadata(
+            env_meta=env_meta,
+            render=False, 
+            render_offscreen=False,
+            use_image_obs=False, 
+        )
+        assert len(env.env.robots) in (1, 2)
+
+        abs_env = EnvUtils.create_env_from_metadata(
+            env_meta=abs_env_meta,
+            render=False, 
+            render_offscreen=False,
+            use_image_obs=False, 
+        )
+        assert not abs_env.env.robots[0].controller.use_delta
+
+        self.env = env
+        self.abs_env = abs_env
+        self.file = h5py.File(dataset_path, 'r')
+    
+    def get_demo_keys(self):
+        return list(self.file['data'].keys())
+
+    def convert_actions(self, 
+            states: np.ndarray, 
+            actions: np.ndarray,
+            initial_state: dict) -> np.ndarray:
+        """
+        Given state and delta action sequence
+        generate equivalent goal position and orientation for each step
+        keep the original gripper action intact.
+        """
+        env = self.env
+        d_a = len(env.env.robots[0].action_limits[0])
+
+        # in case of multi robot
+        # reshape (N,14) to (N,2,7)
+        # or (N,7) to (N,1,7)
+        stacked_actions = actions.reshape(*actions.shape[:-1], -1, d_a)
+
+        # generate abs actions
+        action_goal_pos = np.zeros(
+            stacked_actions.shape[:-1]+(3,), 
+            dtype=stacked_actions.dtype)
+        action_goal_ori = np.zeros(
+            stacked_actions.shape[:-1]+(3,), 
+            dtype=stacked_actions.dtype)
+        action_remainder = stacked_actions[...,6:]
+        for i in range(len(states)):
+            if i == 0:
+                _ = env.reset_to(initial_state)
+            else:
+                _ = env.reset_to({'states': states[i]})
+
+            # taken from robot_env.py L#454
+            for idx, robot in enumerate(env.env.robots):
+                # run controller goal generator
+                robot.control(stacked_actions[i,idx], policy_step=True)
+            
+                # read pos and ori from robots
+                controller = robot.controller
+                action_goal_pos[i,idx] = controller.goal_pos
+                action_goal_ori[i,idx] = Rotation.from_matrix(
+                    controller.goal_ori).as_rotvec()
+
+        stacked_abs_actions = np.concatenate([
+            action_goal_pos,
+            action_goal_ori,
+            action_remainder
+        ], axis=-1)
+        abs_actions = stacked_abs_actions.reshape(actions.shape)
+        return abs_actions
+
+    def convert_demo(self, demo_key):
+        file = self.file
+        demo = file["data/{}".format(demo_key)]
+        # input
+        states = demo['states'][:]
+        actions = demo['actions'][:]
+        initial_state = dict(states=states[0])
+        initial_state["model"] = demo.attrs["model_file"]
+        initial_state["ep_meta"] = demo.attrs.get("ep_meta", None)
+
+        # generate abs actions
+        abs_actions = self.convert_actions(states, actions, initial_state=initial_state)
+        return abs_actions
+
+    def convert_and_eval_demo(self, demo_key):
+        raise NotImplementedError
+        env = self.env
+        abs_env = self.abs_env
+        file = self.file
+        # first step have high error for some reason, not representative
+        eval_skip_steps = 1
+
+        demo = file["data/{}".format(demo_key)]
+        # input
+        states = demo['states'][:]
+        actions = demo['actions'][:]
+
+        # generate abs actions
+        abs_actions = self.convert_actions(states, actions)
+
+        # verify
+        robot0_eef_pos = demo['obs']['robot0_eef_pos'][:]
+        robot0_eef_quat = demo['obs']['robot0_eef_quat'][:]
+
+        delta_error_info = self.evaluate_rollout_error(
+            env, states, actions, robot0_eef_pos, robot0_eef_quat, 
+            metric_skip_steps=eval_skip_steps)
+        abs_error_info = self.evaluate_rollout_error(
+            abs_env, states, abs_actions, robot0_eef_pos, robot0_eef_quat,
+            metric_skip_steps=eval_skip_steps)
+
+        info = {
+            'delta_max_error': delta_error_info,
+            'abs_max_error': abs_error_info
+        }
+        return abs_actions, info
+
+    @staticmethod
+    def evaluate_rollout_error(env, 
+            states, actions, 
+            robot0_eef_pos, 
+            robot0_eef_quat, 
+            metric_skip_steps=1):
+        # first step have high error for some reason, not representative
+
+        # evaluate abs actions
+        rollout_next_states = list()
+        rollout_next_eef_pos = list()
+        rollout_next_eef_quat = list()
+        obs = env.reset_to({'states': states[0]})
+        for i in range(len(states)):
+            obs = env.reset_to({'states': states[i]})
+            obs, reward, done, info = env.step(actions[i])
+            obs = env.get_observation()
+            rollout_next_states.append(env.get_state()['states'])
+            rollout_next_eef_pos.append(obs['robot0_eef_pos'])
+            rollout_next_eef_quat.append(obs['robot0_eef_quat'])
+        rollout_next_states = np.array(rollout_next_states)
+        rollout_next_eef_pos = np.array(rollout_next_eef_pos)
+        rollout_next_eef_quat = np.array(rollout_next_eef_quat)
+
+        next_state_diff = states[1:] - rollout_next_states[:-1]
+        max_next_state_diff = np.max(np.abs(next_state_diff[metric_skip_steps:]))
+
+        next_eef_pos_diff = robot0_eef_pos[1:] - rollout_next_eef_pos[:-1]
+        next_eef_pos_dist = np.linalg.norm(next_eef_pos_diff, axis=-1)
+        max_next_eef_pos_dist = next_eef_pos_dist[metric_skip_steps:].max()
+
+        next_eef_rot_diff = Rotation.from_quat(robot0_eef_quat[1:]) \
+            * Rotation.from_quat(rollout_next_eef_quat[:-1]).inv()
+        next_eef_rot_dist = next_eef_rot_diff.magnitude()
+        max_next_eef_rot_dist = next_eef_rot_dist[metric_skip_steps:].max()
+
+        info = {
+            'state': max_next_state_diff,
+            'pos': max_next_eef_pos_dist,
+            'rot': max_next_eef_rot_dist
+        }
+        return info
+
+"""
+copied/adapted from https://github.com/columbia-ai-robotics/diffusion_policy/blob/main/diffusion_policy/scripts/robomimic_dataset_conversion.py
+"""
+def worker(x):
+    path, demo_key, do_eval = x
+    converter = RobomimicAbsoluteActionConverter(path)
+    if do_eval:
+        abs_actions, info = converter.convert_and_eval_demo(demo_key)
+    else:
+        abs_actions = converter.convert_demo(demo_key)
+        info = dict()
+    return abs_actions, info
+
+
+def add_absolute_actions_to_dataset(dataset, eval_dir, num_workers):
+    # process inputs
+    dataset = pathlib.Path(dataset).expanduser()
+    assert dataset.is_file()
+
+    do_eval = False
+    if eval_dir is not None:
+        eval_dir = pathlib.Path(eval_dir).expanduser()
+        assert eval_dir.parent.exists()
+        do_eval = True
+    
+    converter = RobomimicAbsoluteActionConverter(dataset)
+    demo_keys = converter.get_demo_keys()
+    del converter
+    
+    # run
+    with multiprocessing.Pool(num_workers) as pool:
+        results = pool.map(worker, [(dataset, demo_key, do_eval) for demo_key in demo_keys])
+
+    # modify action
+    with h5py.File(dataset, 'r+') as out_file:
+        for i in tqdm(range(len(results)), desc="Writing to output"):
+            abs_actions, info = results[i]
+            demo = out_file["data/{}".format(demo_keys[i])]
+            if "actions_abs" not in demo:
+                demo.create_dataset("actions_abs", data=np.array(abs_actions))
+            else:
+                demo['actions_abs'][:] = abs_actions
+    
+    # save eval
+    if do_eval:
+        eval_dir.mkdir(parents=False, exist_ok=True)
+
+        print("Writing error_stats.pkl")
+        infos = [info for _, info in results]
+        pickle.dump(infos, eval_dir.joinpath('error_stats.pkl').open('wb'))
+
+        print("Generating visualization")
+        metrics = ['pos', 'rot']
+        metrics_dicts = dict()
+        for m in metrics:
+            metrics_dicts[m] = collections.defaultdict(list)
+
+        for i in range(len(infos)):
+            info = infos[i]
+            for k, v in info.items():
+                for m in metrics:
+                    metrics_dicts[m][k].append(v[m])
+
+        from matplotlib import pyplot as plt
+        plt.switch_backend('PDF')
+
+        fig, ax = plt.subplots(1, len(metrics))
+        for i in range(len(metrics)):
+            axis = ax[i]
+            data = metrics_dicts[metrics[i]]
+            for key, value in data.items():
+                axis.plot(value, label=key)
+            axis.legend()
+            axis.set_title(metrics[i])
+        fig.set_size_inches(10,4)
+        fig.savefig(str(eval_dir.joinpath('error_stats.pdf')))
+        fig.savefig(str(eval_dir.joinpath('error_stats.png')))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        required=True
+    )
+
+    parser.add_argument(
+        "--eval_dir",
+        type=str,
+        help="directory to output evaluation metrics",
+    )
+
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=10,
+    )
+    
+    args = parser.parse_args()
+    
+    
+    add_absolute_actions_to_dataset(
+        dataset=args.dataset,
+        eval_dir=args.eval_dir,
+        num_workers=args.num_workers,
+    )

aloha-devel/robomimic/scripts/dataset_states_to_obs.py

@@ -0,0 +1,375 @@
+"""
+Script to extract observations from low-dimensional simulation states in a robosuite dataset.
+
+Args:
+    dataset (str): path to input hdf5 dataset
+
+    output_name (str): name of output hdf5 dataset
+
+    n (int): if provided, stop after n trajectories are processed
+
+    shaped (bool): if flag is set, use dense rewards
+
+    camera_names (str or [str]): camera name(s) to use for image observations. 
+        Leave out to not use image observations.
+
+    camera_height (int): height of image observation.
+
+    camera_width (int): width of image observation
+
+    done_mode (int): how to write done signal. If 0, done is 1 whenever s' is a success state.
+        If 1, done is 1 at the end of each trajectory. If 2, both.
+
+    copy_rewards (bool): if provided, copy rewards from source file instead of inferring them
+
+    copy_dones (bool): if provided, copy dones from source file instead of inferring them
+
+Example usage:
+    
+    # extract low-dimensional observations
+    python dataset_states_to_obs.py --dataset /path/to/demo.hdf5 --output_name low_dim.hdf5 --done_mode 2
+    
+    # extract 84x84 image observations
+    python dataset_states_to_obs.py --dataset /path/to/demo.hdf5 --output_name image.hdf5 \
+        --done_mode 2 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84
+
+    # (space saving option) extract 84x84 image observations with compression and without 
+    # extracting next obs (not needed for pure imitation learning algos)
+    python dataset_states_to_obs.py --dataset /path/to/demo.hdf5 --output_name image.hdf5 \
+        --done_mode 2 --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84 \
+        --compress --exclude-next-obs
+
+    # use dense rewards, and only annotate the end of trajectories with done signal
+    python dataset_states_to_obs.py --dataset /path/to/demo.hdf5 --output_name image_dense_done_1.hdf5 \
+        --done_mode 1 --dense --camera_names agentview robot0_eye_in_hand --camera_height 84 --camera_width 84
+"""
+import os
+import json
+import h5py
+import argparse
+import numpy as np
+from copy import deepcopy
+
+import robomimic.utils.tensor_utils as TensorUtils
+import robomimic.utils.file_utils as FileUtils
+import robomimic.utils.env_utils as EnvUtils
+from robomimic.envs.env_base import EnvBase
+
+
+def extract_trajectory(
+    env, 
+    initial_state, 
+    states, 
+    actions,
+    actions_abs,
+    done_mode,
+):
+    """
+    Helper function to extract observations, rewards, and dones along a trajectory using
+    the simulator environment.
+
+    Args:
+        env (instance of EnvBase): environment
+        initial_state (dict): initial simulation state to load
+        states (np.array): array of simulation states to load to extract information
+        actions (np.array): array of actions
+        done_mode (int): how to write done signal. If 0, done is 1 whenever s' is a 
+            success state. If 1, done is 1 at the end of each trajectory. 
+            If 2, do both.
+    """
+    assert isinstance(env, EnvBase)
+    assert states.shape[0] == actions.shape[0]
+
+    # load the initial state
+    ## this reset call doesn't seem necessary.
+    ## seems ok to remove but haven't fully tested it.
+    ## removing for now
+    # env.reset()
+    obs = env.reset_to(initial_state)
+
+    traj = dict(
+        obs=[], 
+        next_obs=[], 
+        rewards=[], 
+        dones=[], 
+        actions=np.array(actions), 
+        states=np.array(states), 
+        initial_state_dict=initial_state,
+    )
+    if actions_abs is not None:
+        traj["actions_abs"] = np.array(actions_abs)
+    
+    traj_len = states.shape[0]
+    # iteration variable @t is over "next obs" indices
+    for t in range(1, traj_len + 1):
+
+        # get next observation
+        if t == traj_len:
+            # play final action to get next observation for last timestep
+            next_obs, _, _, _ = env.step(actions[t - 1])
+        else:
+            # reset to simulator state to get observation
+            next_obs = env.reset_to({"states" : states[t]})
+
+        # infer reward signal
+        # note: our tasks use reward r(s'), reward AFTER transition, so this is
+        #       the reward for the current timestep
+        r = env.get_reward()
+
+        # infer done signal
+        done = False
+        if (done_mode == 1) or (done_mode == 2):
+            # done = 1 at end of trajectory
+            done = done or (t == traj_len)
+        if (done_mode == 0) or (done_mode == 2):
+            # done = 1 when s' is task success state
+            done = done or env.is_success()["task"]
+        done = int(done)
+
+        # collect transition
+        traj["obs"].append(obs)
+        traj["next_obs"].append(next_obs)
+        traj["rewards"].append(r)
+        traj["dones"].append(done)
+
+        # update for next iter
+        obs = deepcopy(next_obs)
+
+    # convert list of dict to dict of list for obs dictionaries (for convenient writes to hdf5 dataset)
+    traj["obs"] = TensorUtils.list_of_flat_dict_to_dict_of_list(traj["obs"])
+    traj["next_obs"] = TensorUtils.list_of_flat_dict_to_dict_of_list(traj["next_obs"])
+
+    # list to numpy array
+    for k in traj:
+        if k == "initial_state_dict":
+            continue
+        if isinstance(traj[k], dict):
+            for kp in traj[k]:
+                traj[k][kp] = np.array(traj[k][kp])
+        else:
+            traj[k] = np.array(traj[k])
+
+    return traj
+
+
+def dataset_states_to_obs(args):
+    # create environment to use for data processing
+    env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=args.dataset)
+    env = EnvUtils.create_env_for_data_processing(
+        env_meta=env_meta,
+        camera_names=args.camera_names, 
+        camera_height=args.camera_height, 
+        camera_width=args.camera_width, 
+        reward_shaping=args.shaped,
+    )
+
+    print("==== Using environment with the following metadata ====")
+    print(json.dumps(env.serialize(), indent=4))
+    print("")
+
+    # some operations for playback are robosuite-specific, so determine if this environment is a robosuite env
+    is_robosuite_env = EnvUtils.is_robosuite_env(env_meta)
+
+    # list of all demonstration episodes (sorted in increasing number order)
+    f = h5py.File(args.dataset, "r")
+    demos = list(f["data"].keys())
+    inds = np.argsort([int(elem[5:]) for elem in demos])
+    demos = [demos[i] for i in inds]
+
+    # maybe reduce the number of demonstrations to playback
+    if args.n is not None:
+        demos = demos[:args.n]
+
+    # output file in same directory as input file
+    output_name = args.output_name
+    if output_name is None:
+        if len(args.camera_names) == 0:
+            output_name = os.path.basename(args.dataset)[:-5] + "_ld.hdf5"
+        else:
+            output_name = os.path.basename(args.dataset)[:-5] + "_im{}.hdf5".format(args.camera_width)
+
+    output_path = os.path.join(os.path.dirname(args.dataset), output_name)
+    f_out = h5py.File(output_path, "w")
+    data_grp = f_out.create_group("data")
+    print("input file: {}".format(args.dataset))
+    print("output file: {}".format(output_path))
+
+    total_samples = 0
+    for ind in range(len(demos)):
+        ep = demos[ind]
+
+        # prepare initial state to reload from
+        states = f["data/{}/states".format(ep)][()]
+        initial_state = dict(states=states[0])
+        if is_robosuite_env:
+            initial_state["model"] = f["data/{}".format(ep)].attrs["model_file"]
+            initial_state["ep_meta"] = f["data/{}".format(ep)].attrs.get("ep_meta", None)
+
+        # extract obs, rewards, dones
+        actions = f["data/{}/actions".format(ep)][()]
+        if "data/{}/actions_abs".format(ep) in f:
+            actions_abs = f["data/{}/actions_abs".format(ep)][()]
+        else:
+            actions_abs = None
+        traj = extract_trajectory(
+            env=env, 
+            initial_state=initial_state, 
+            states=states, 
+            actions=actions,
+            actions_abs=actions_abs,
+            done_mode=args.done_mode,
+        )
+
+        # maybe copy reward or done signal from source file
+        if args.copy_rewards:
+            traj["rewards"] = f["data/{}/rewards".format(ep)][()]
+        if args.copy_dones:
+            traj["dones"] = f["data/{}/dones".format(ep)][()]
+
+        # store transitions
+
+        # IMPORTANT: keep name of group the same as source file, to make sure that filter keys are
+        #            consistent as well
+        ep_data_grp = data_grp.create_group(ep)
+        ep_data_grp.create_dataset("actions", data=np.array(traj["actions"]))
+        ep_data_grp.create_dataset("states", data=np.array(traj["states"]))
+        ep_data_grp.create_dataset("rewards", data=np.array(traj["rewards"]))
+        ep_data_grp.create_dataset("dones", data=np.array(traj["dones"]))
+        if "actions_abs" in traj:
+            ep_data_grp.create_dataset("actions_abs", data=np.array(traj["actions_abs"]))
+        for k in traj["obs"]:
+            if args.compress:
+                ep_data_grp.create_dataset("obs/{}".format(k), data=np.array(traj["obs"][k]), compression="gzip")
+            else:
+                ep_data_grp.create_dataset("obs/{}".format(k), data=np.array(traj["obs"][k]))
+            if not args.exclude_next_obs:
+                if args.compress:
+                    ep_data_grp.create_dataset("next_obs/{}".format(k), data=np.array(traj["next_obs"][k]), compression="gzip")
+                else:
+                    ep_data_grp.create_dataset("next_obs/{}".format(k), data=np.array(traj["next_obs"][k]))
+
+        # copy action dict (if applicable)
+        if "data/{}/action_dict".format(ep) in f:
+            action_dict = f["data/{}/action_dict".format(ep)]
+            for k in action_dict:
+                ep_data_grp.create_dataset("action_dict/{}".format(k), data=np.array(action_dict[k][()]))
+
+        # episode metadata
+        if is_robosuite_env:
+            ep_data_grp.attrs["model_file"] = traj["initial_state_dict"]["model"] # model xml for this episode
+        if "ep_meta" in f["data/{}".format(ep)].attrs:
+            ep_data_grp.attrs["ep_meta"] = f["data/{}".format(ep)].attrs["ep_meta"]
+        ep_data_grp.attrs["num_samples"] = traj["actions"].shape[0] # number of transitions in this episode
+        total_samples += traj["actions"].shape[0]
+        print("ep {}: wrote {} transitions to group {}".format(ind, ep_data_grp.attrs["num_samples"], ep))
+
+
+    # copy over all filter keys that exist in the original hdf5
+    if "mask" in f:
+        f.copy("mask", f_out)
+
+    # global metadata
+    data_grp.attrs["total"] = total_samples
+    data_grp.attrs["env_args"] = json.dumps(env.serialize(), indent=4) # environment info
+    print("Wrote {} trajectories to {}".format(len(demos), output_path))
+
+    f.close()
+    f_out.close()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        required=True,
+        help="path to input hdf5 dataset",
+    )
+    # name of hdf5 to write - it will be in the same directory as @dataset
+    parser.add_argument(
+        "--output_name",
+        type=str,
+        help="name of output hdf5 dataset",
+    )
+
+    # specify number of demos to process - useful for debugging conversion with a handful
+    # of trajectories
+    parser.add_argument(
+        "--n",
+        type=int,
+        default=None,
+        help="(optional) stop after n trajectories are processed",
+    )
+
+    # flag for reward shaping
+    parser.add_argument(
+        "--shaped", 
+        action='store_true',
+        help="(optional) use shaped rewards",
+    )
+
+    # camera names to use for observations
+    parser.add_argument(
+        "--camera_names",
+        type=str,
+        nargs='+',
+        default=[],
+        help="(optional) camera name(s) to use for image observations. Leave out to not use image observations.",
+    )
+
+    parser.add_argument(
+        "--camera_height",
+        type=int,
+        default=84,
+        help="(optional) height of image observations",
+    )
+
+    parser.add_argument(
+        "--camera_width",
+        type=int,
+        default=84,
+        help="(optional) width of image observations",
+    )
+
+    # specifies how the "done" signal is written. If "0", then the "done" signal is 1 wherever 
+    # the transition (s, a, s') has s' in a task completion state. If "1", the "done" signal 
+    # is one at the end of every trajectory. If "2", the "done" signal is 1 at task completion
+    # states for successful trajectories and 1 at the end of all trajectories.
+    parser.add_argument(
+        "--done_mode",
+        type=int,
+        default=0,
+        help="how to write done signal. If 0, done is 1 whenever s' is a success state.\
+            If 1, done is 1 at the end of each trajectory. If 2, both.",
+    )
+
+    # flag for copying rewards from source file instead of re-writing them
+    parser.add_argument(
+        "--copy_rewards", 
+        action='store_true',
+        help="(optional) copy rewards from source file instead of inferring them",
+    )
+
+    # flag for copying dones from source file instead of re-writing them
+    parser.add_argument(
+        "--copy_dones", 
+        action='store_true',
+        help="(optional) copy dones from source file instead of inferring them",
+    )
+
+    # flag to exclude next obs in dataset
+    parser.add_argument(
+        "--exclude-next-obs", 
+        action='store_true',
+        help="(optional) exclude next obs in dataset",
+    )
+
+    # flag to compress observations with gzip option in hdf5
+    parser.add_argument(
+        "--compress", 
+        action='store_true',
+        help="(optional) compress observations with gzip option in hdf5",
+    )
+
+    args = parser.parse_args()
+    dataset_states_to_obs(args)

aloha-devel/robomimic/scripts/download_momart_datasets.py

@@ -0,0 +1,161 @@
+"""
+Script to download datasets used in MoMaRT paper (https://arxiv.org/abs/2112.05251). By default, all
+datasets will be stored at robomimic/datasets, unless the @download_dir
+argument is supplied. We recommend using the default, as most examples that
+use these datasets assume that they can be found there.
+
+The @tasks and @dataset_types arguments can all be supplied
+to choose which datasets to download. 
+
+Args:
+    download_dir (str): Base download directory. Created if it doesn't exist. 
+        Defaults to datasets folder in repository - only pass in if you would
+        like to override the location.
+
+    tasks (list): Tasks to download datasets for. Defaults to table_setup_from_dishwasher task. Pass 'all' to
+        download all tasks - 5 total:
+            - table_setup_from_dishwasher
+            - table_setup_from_dresser
+            - table_cleanup_to_dishwasher
+            - table_cleanup_to_sink
+            - unload_dishwasher
+    
+    dataset_types (list): Dataset types to download datasets for (expert, suboptimal, generalize, sample).
+        Defaults to expert. Pass 'all' to download datasets for all available dataset
+        types per task, or directly specify the list of dataset types.
+        NOTE: Because these datasets are huge, we will always print out a warning
+        that a user must respond yes to to acknowledge the data size (can be up to >100G for all tasks of a single type)
+
+Example usage:
+
+    # default behavior - just download expert table_setup_from_dishwasher dataset
+    python download_momart_datasets.py
+
+    # download expert datasets for all tasks
+    # (do a dry run first to see which datasets would be downloaded)
+    python download_momart_datasets.py --tasks all --dataset_types expert --dry_run
+    python download_momart_datasets.py --tasks all --dataset_types expert low_dim
+
+    # download all expert and suboptimal datasets for the table_setup_from_dishwasher and table_cleanup_to_dishwasher tasks
+    python download_datasets.py --tasks table_setup_from_dishwasher table_cleanup_to_dishwasher --dataset_types expert suboptimal
+
+    # download the sample datasets
+    python download_datasets.py --tasks all --dataset_types sample
+
+    # download all datasets
+    python download_datasets.py --tasks all --dataset_types all
+"""
+import os
+import argparse
+
+import robomimic
+import robomimic.utils.file_utils as FileUtils
+from robomimic import MOMART_DATASET_REGISTRY
+
+ALL_TASKS = [
+    "table_setup_from_dishwasher",
+    "table_setup_from_dresser",
+    "table_cleanup_to_dishwasher",
+    "table_cleanup_to_sink",
+    "unload_dishwasher",
+]
+ALL_DATASET_TYPES = [
+    "expert",
+    "suboptimal",
+    "generalize",
+    "sample",
+]
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    # directory to download datasets to
+    parser.add_argument(
+        "--download_dir",
+        type=str,
+        default=None,
+        help="Base download directory. Created if it doesn't exist. Defaults to datasets folder in repository.",
+    )
+
+    # tasks to download datasets for
+    parser.add_argument(
+        "--tasks",
+        type=str,
+        nargs='+',
+        default=["table_setup_from_dishwasher"],
+        help="Tasks to download datasets for. Defaults to table_setup_from_dishwasher task. Pass 'all' to download all"
+             f"5 tasks, or directly specify the list of tasks. Options are any of: {ALL_TASKS}",
+    )
+
+    # dataset types to download datasets for
+    parser.add_argument(
+        "--dataset_types",
+        type=str,
+        nargs='+',
+        default=["expert"],
+        help="Dataset types to download datasets for (e.g. expert, suboptimal). Defaults to expert. Pass 'all' to "
+             "download datasets for all available dataset types per task, or directly specify the list of dataset "
+             f"types. Options are any of: {ALL_DATASET_TYPES}",
+    )
+
+    # dry run - don't actually download datasets, but print which datasets would be downloaded
+    parser.add_argument(
+        "--dry_run",
+        action='store_true',
+        help="set this flag to do a dry run to only print which datasets would be downloaded"
+    )
+
+    args = parser.parse_args()
+
+    # set default base directory for downloads
+    default_base_dir = args.download_dir
+    if default_base_dir is None:
+        default_base_dir = os.path.join(robomimic.__path__[0], "../datasets")
+
+    # load args
+    download_tasks = args.tasks
+    if "all" in download_tasks:
+        assert len(download_tasks) == 1, "all should be only tasks argument but got: {}".format(args.tasks)
+        download_tasks = ALL_TASKS
+
+    download_dataset_types = args.dataset_types
+    if "all" in download_dataset_types:
+        assert len(download_dataset_types) == 1, "all should be only dataset_types argument but got: {}".format(args.dataset_types)
+        download_dataset_types = ALL_DATASET_TYPES
+
+    # Run sanity check first to warn user if they're about to download a huge amount of data
+    total_size = 0
+    for task in MOMART_DATASET_REGISTRY:
+        if task in download_tasks:
+            for dataset_type in MOMART_DATASET_REGISTRY[task]:
+                if dataset_type in download_dataset_types:
+                    total_size += MOMART_DATASET_REGISTRY[task][dataset_type]["size"]
+
+    # Verify user acknowledgement if we're not doing a dry run
+    if not args.dry_run:
+        user_response = input(f"Warning: requested datasets will take a total of {total_size}GB. Proceed? y/n\n")
+        assert user_response.lower() in {"yes", "y"}, f"Did not receive confirmation. Aborting download."
+
+    # download requested datasets
+    for task in MOMART_DATASET_REGISTRY:
+        if task in download_tasks:
+            for dataset_type in MOMART_DATASET_REGISTRY[task]:
+                if dataset_type in download_dataset_types:
+                    dataset_info = MOMART_DATASET_REGISTRY[task][dataset_type]
+                    download_dir = os.path.abspath(os.path.join(default_base_dir, task, dataset_type))
+                    print(f"\nDownloading dataset:\n"
+                          f"    task: {task}\n"
+                          f"    dataset type: {dataset_type}\n"
+                          f"    dataset size: {dataset_info['size']}GB\n"
+                          f"    download path: {download_dir}")
+                    if args.dry_run:
+                        print("\ndry run: skip download")
+                    else:
+                        # Make sure path exists and create if it doesn't
+                        os.makedirs(download_dir, exist_ok=True)
+                        FileUtils.download_url(
+                            url=dataset_info["url"],
+                            download_dir=download_dir,
+                        )
+                    print("")

aloha-devel/robomimic/scripts/filter_dataset_size.py

@@ -0,0 +1,81 @@
+import argparse
+import h5py
+import numpy as np
+
+from robomimic.utils.file_utils import create_hdf5_filter_key
+
+
+def filter_dataset_size(hdf5_path, num_demos, input_filter_key=None, output_filter_key=None):
+    # retrieve number of demos
+    f = h5py.File(hdf5_path, "r")
+    if input_filter_key is not None:
+        print("using filter key: {}".format(input_filter_key))
+        demos = sorted([elem.decode("utf-8") for elem in np.array(f["mask/{}".format(input_filter_key)])])
+    else:
+        demos = sorted(list(f["data"].keys()))
+    f.close()
+
+    # get random split
+    total_num_demos = len(demos)
+    mask = np.zeros(total_num_demos)
+    mask[:num_demos] = 1.
+    np.random.shuffle(mask)
+    mask = mask.astype(int)
+    subset_inds = mask.nonzero()[0]
+    subset_keys = [demos[i] for i in subset_inds]
+
+    # pass mask to generate split
+    if output_filter_key is not None:
+        name = output_filter_key
+    else:
+        name = "{}_demos".format(num_demos)
+
+    if input_filter_key is not None:
+        name = "{}_{}".format(input_filter_key, name)
+
+    subset_lengths = create_hdf5_filter_key(hdf5_path=hdf5_path, demo_keys=subset_keys, key_name=name)
+
+    print("Total number of subset samples: {}".format(np.sum(subset_lengths)))
+    print("Average number of subset samples {}".format(np.mean(subset_lengths)))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        required=True,
+        help="path to hdf5 dataset",
+    )
+    parser.add_argument(
+        "--input_filter_key",
+        type=str,
+        default=None,
+        help="if provided, split the subset of trajectories in the file that correspond to\
+            this filter key into a training and validation set of trajectories, instead of\
+            splitting the full set of trajectories",
+    )
+    parser.add_argument(
+        "--num_demos",
+        type=int,
+        nargs='+',
+        required=True,
+    )
+    parser.add_argument(
+        "--output_filter_key",
+        type=str,
+        required=False,
+        help="(optional) use custom name for output filter key name"
+    )
+    args = parser.parse_args()
+
+    # seed to make sure results are consistent
+    np.random.seed(0)
+
+    for n in args.num_demos:
+        filter_dataset_size(
+            args.dataset,
+            input_filter_key=args.input_filter_key,
+            num_demos=n,
+            output_filter_key=args.output_filter_key,
+        )

aloha-devel/robomimic/scripts/plot_model_predictions.py

@@ -0,0 +1,213 @@
+import json
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.gridspec as gridspec
+from copy import deepcopy
+import random
+from sklearn.metrics import mean_squared_error
+import re
+import robomimic.utils.file_utils as FileUtils
+import robomimic.utils.torch_utils as TorchUtils
+import robomimic.utils.tensor_utils as TensorUtils
+import robomimic.utils.train_utils as TrainUtils
+from robomimic.config import config_factory
+import robomimic.utils.obs_utils as ObsUtils
+import torch
+from torch.utils.data import DataLoader
+
+"""
+TODO: track rotation magnitude seperately (https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.transform.Rotation.magnitude.html)
+"""
+
+# the configs of the models to be plotted
+model_config_mapping = {
+    # "bottle_less_obs": {
+    #     "model":"/home/zehan/expdata/r2d2/im/bc_xfmr/google_bc_baseline/bottle_less_obs/20230815225106/models/model_epoch_60.pth",
+    #     'folder':"/home/zehan/expdata/r2d2/im/bc_xfmr/google_bc_baseline/bottle_less_obs/20230815225106/test_inference_figures/",
+    #     # "action_names": ['x', 'y', 'z', 'roll', 'pitch', 'yaw', "gripper_action" , 'terminate'],
+    #     "action_names": None,
+    #     "trajectory_name_regex": r'(\d+_trajectory_im\d+)'
+    # },
+    "r2d2_wire": {
+        # "model": "/home/soroushn/expdata/r2d2/im/diffusion_policy/debug/ds_pen-in-cup_cams_3cams/20230830160945/models/model_epoch_2.pth",
+        "model": "/home/soroushn/expdata/r2d2/im/bc_xfmr/debug/ds_pen-in-cup_cams_3cams_predfuture_True_ac_keys_rel/20230830161631/models/model_epoch_2.pth",
+        "folder": "/home/soroushn/tmp/model_predictions",
+        # "action_names": ['x', 'y', 'z', 'r', 'p', 'y', "gripper_pos"], # use custom names
+        "action_names": None, # use default names, see line 71
+        "trajectory_name_regex": r'(\w+_\w+_\d{2}_\d{2}:\d{2}:\d{2}_\d{4})' # the name of the figure files need to be custom defined (the part of the names of the trajectories that uniquely identifies them)
+    }
+}
+
+NUM_SAMPLES = 2
+
+# loop through each model
+for model_name in model_config_mapping:
+    ckpt_path = model_config_mapping[model_name]['model']
+    saving_folder = model_config_mapping[model_name]['folder']
+    # can custom-define or using default action_names
+    action_names = model_config_mapping[model_name]['action_names']
+    trajectory_name_regex = model_config_mapping[model_name]['trajectory_name_regex']
+    accuracy_thresholds = np.logspace(-3,-5, num=3).tolist()
+
+
+
+    device = TorchUtils.get_torch_device(try_to_use_cuda=True)
+
+    ckpt_dict = FileUtils.maybe_dict_from_checkpoint(ckpt_path=ckpt_path)
+    config = json.loads(ckpt_dict["config"])
+    config["train"]["shuffled_obs_key_groups"] = None
+    ckpt_dict["config"] = json.dumps(config)
+    policy, _ = FileUtils.policy_from_checkpoint(ckpt_dict=ckpt_dict, device=device, verbose=True)
+    shape_meta = ckpt_dict['shape_metadata']
+    ext_cfg = json.loads(ckpt_dict["config"])
+    config = config_factory(ext_cfg["algo_name"])
+    with config.values_unlocked():
+        config.update(ext_cfg)
+    
+    
+    frame_stack = config.train.frame_stack 
+    device = TorchUtils.get_torch_device(try_to_use_cuda=config.train.cuda)
+
+    trainset, validset = TrainUtils.load_data_for_training(config, obs_keys=shape_meta["all_obs_keys"])
+    # trainset.datasets is a list
+    # the trajectories to plot is randomly sampled from the training and validation sets
+    training_sampled_data = random.sample(trainset.datasets, NUM_SAMPLES)
+    # validation_sampled_data = random.sample(validset.datasets, NUM_SAMPLES)
+
+    inference_datasets_mapping = {"training": training_sampled_data} #, "validation": validation_sampled_data}
+
+    
+    if action_names == None:
+        # TODO
+        action_keys = config.train.action_keys # Need to adjust. For Robomimic datasets, there is no `action_keys`, it is config.train.dataset_keys
+        modified_action_keys = [element.replace('action/', '') for element in action_keys]
+        action_names = []
+        for i, action_key in enumerate(action_keys):
+            if isinstance(training_sampled_data[0].__getitem__(0)[action_key][frame_stack-1], np.ndarray):
+                action_names.extend([f'{modified_action_keys[i]}_{j+1}' for j in range(len(training_sampled_data[0].__getitem__(0)[action_key][frame_stack-1]))])
+            else:
+                action_names.append(modified_action_keys[i])
+
+    # loop through training and validation sets
+    for inference_key in inference_datasets_mapping:
+        mse_training_per_traj = []
+        data_name = []
+        actual_actions_all_traj = [] # (NxT, D)
+        predicted_actions_all_traj = [] # (NxT, D)
+
+        # loop through each trajectory
+        for d in inference_datasets_mapping[inference_key]:
+            hdf5_path = d.hdf5_path
+            mse_for_one_traj = []
+            traj_length = len(d)
+            action_dim = len(action_names)
+            actual_actions = [[] for _ in range(action_dim)] # (T, D)
+            predicted_actions = [[] for _ in range(action_dim)] # (T, D)           
+
+            image_keys = [item for item in d.__getitem__(0)['obs'].keys() if "image" in item]
+            images = {key: [] for key in image_keys}
+
+            dataloader = DataLoader(
+                dataset=d,
+                sampler=None,
+                batch_size=1,
+                shuffle=False,
+                num_workers=1,
+                drop_last=True,
+            )
+
+            model = policy.policy
+
+            model.reset()
+
+            # loop through each timestep
+            for batch in iter(dataloader):
+                batch = model.process_batch_for_training(batch)
+
+                for image_key in image_keys:
+                    im = batch["obs"][image_key][0][-1]
+                    im = TensorUtils.to_numpy(im).astype(np.uint32)
+                    images[image_key].append(im)
+
+                batch = model.postprocess_batch_for_training(batch, obs_normalization_stats=None) # ignore obs_normalization for now
+                # model_output = model.nets["policy"](batch["obs"])
+
+                model_output = model.get_action(batch["obs"])
+                
+                actual_action = TensorUtils.to_numpy(
+                    batch["actions"][0][0]
+                )
+                predicted_action = TensorUtils.to_numpy(
+                    model_output[0]
+                )
+
+                actual_actions_all_traj.append(actual_action)
+                predicted_actions_all_traj.append(predicted_action)
+            
+                for dim in range(action_dim):
+                    actual_actions[dim].append(actual_action[dim])
+                    predicted_actions[dim].append(predicted_action[dim])
+
+            # Plot
+            fig, axs = plt.subplots(len(images) + action_dim, 1, figsize=(30, (len(images) + action_dim) * 3))
+            for i, image_key in enumerate(image_keys):
+                interval = int(traj_length/15) # plot `5` images
+                images[image_key] = images[image_key][::interval]
+                combined_images = np.concatenate(images[image_key], axis=1)
+                axs[i].imshow(combined_images)
+                if i == 0:
+                    axs[i].set_title(hdf5_path + '\n' + image_key, fontsize=30)
+                else:
+                    axs[i].set_title(image_key, fontsize=30)
+                axs[i].axis("off")
+            for dim in range(action_dim):
+                mse = mean_squared_error(actual_actions[dim], predicted_actions[dim])
+                mse_for_one_traj.append(mse)
+                axs[len(images)+dim].plot(range(traj_length), actual_actions[dim], label='Actual Action', color='blue')
+                axs[len(images)+dim].plot(range(traj_length), predicted_actions[dim], label='Predicted Action', color='red')
+                # axs[len(images)+dim].set_xlabel('Timestep')
+                # axs[len(images)+dim].set_ylabel('Action Dimension {}'.format(dim + 1))
+                axs[len(images)+dim].set_title(action_names[dim], fontsize=30)
+                axs[len(images)+dim].xaxis.set_tick_params(labelsize=24)
+                axs[len(images)+dim].yaxis.set_tick_params(labelsize=24)
+                axs[len(images)+dim].legend(fontsize=20)
+            plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05, wspace=0.3, hspace=0.6)
+
+            # Save inference figures
+            save_path = saving_folder + inference_key+"/" #remember to add / at the end
+            # data_content = re.search(trajectory_name_regex, hdf5_path).group(1)
+            data_content = "test"
+            filename = "comparison_figure_"+data_content +".png"    
+            if not os.path.exists(save_path):
+                os.makedirs(save_path)
+            print(save_path + filename)
+            # Save the figure with the specified path and filename
+            plt.savefig(save_path + filename) 
+            mse_training_per_traj.append(mse_for_one_traj)
+            data_name.append(hdf5_path)
+
+        # log MSE information
+        accuracy_thresholds = np.logspace(-3,-5, num=3).tolist()
+        mse = torch.nn.functional.mse_loss(torch.tensor(predicted_actions_all_traj), torch.tensor(actual_actions_all_traj), reduction='none') # (NxT, D)
+        step_log = {}
+        step_log[f'{inference_key}_action_mse_error'] = mse.mean().item() # average MSE across all timesteps averaged across all action dimensions (D,)
+        
+        # compute percentage of timesteps that have MSE less than the accuracy thresholds
+        for accuracy_threshold in accuracy_thresholds:
+            step_log[f'{inference_key}_action_accuracy@{accuracy_threshold}'] = (torch.less(mse,accuracy_threshold).float().mean().item())
+        
+
+        average_mse_per_dimension = np.mean(mse_training_per_traj, axis=0) # (D,)
+        txt_path = saving_folder+inference_key+"/" +"output.txt"
+        list_str = '\n'.join(['{} {}'.format(desc, ' '.join(map(str, sublist))) for desc, sublist in zip(data_name, mse_training_per_traj)])
+        
+        # save MSE information
+        with open(txt_path, "w+") as txt_file:
+            txt_file.write(f"MSE per trajectory:\n{list_str}\n")
+            txt_file.write("\n")
+            txt_file.write(f"Average MSE across trajectories per dimension: {average_mse_per_dimension}\n")
+            txt_file.write("\n")
+            txt_file.write(f"MSE log: {step_log}\n")
+
+

aloha-devel/robomimic/scripts/setup_macros.py

@@ -0,0 +1,32 @@
+"""
+This script sets up a private macros file.
+
+The private macros file (macros_private.py) is not tracked by git,
+allowing user-specific settings that are not tracked by git.
+
+This script checks if macros_private.py exists.
+If applicable, it creates the private macros at robomimic/macros_private.py
+"""
+
+import os
+import robomimic
+import shutil
+
+if __name__ == "__main__":
+    base_path = robomimic.__path__[0]
+    macros_path = os.path.join(base_path, "macros.py")
+    macros_private_path = os.path.join(base_path, "macros_private.py")
+
+    if not os.path.exists(macros_path):
+        print("{} does not exist! Aborting...".format(macros_path))
+
+    if os.path.exists(macros_private_path):
+        ans = input("{} already exists! \noverwrite? (y/n)\n".format(macros_private_path))
+
+        if ans == "y":
+            print("REMOVING")
+        else:
+            exit()
+
+    shutil.copyfile(macros_path, macros_private_path)
+    print("copied {}\nto {}".format(macros_path, macros_private_path))

aloha-devel/robomimic/scripts/split_train_val.py

@@ -0,0 +1,105 @@
+"""
+Script for splitting a dataset hdf5 file into training and validation trajectories.
+
+Args:
+    dataset (str): path to hdf5 dataset
+
+    filter_key (str): if provided, split the subset of trajectories
+        in the file that correspond to this filter key into a training
+        and validation set of trajectories, instead of splitting the
+        full set of trajectories
+
+    ratio (float): validation ratio, in (0, 1). Defaults to 0.1, which is 10%.
+
+Example usage:
+    python split_train_val.py --dataset /path/to/demo.hdf5 --ratio 0.1
+"""
+
+import argparse
+import h5py
+import numpy as np
+
+from robomimic.utils.file_utils import create_hdf5_filter_key
+
+
+def split_train_val_from_hdf5(hdf5_path, val_ratio=0.1, filter_key=None):
+    """
+    Splits data into training set and validation set from HDF5 file.
+
+    Args:
+        hdf5_path (str): path to the hdf5 file
+            to load the transitions from
+
+        val_ratio (float): ratio of validation demonstrations to all demonstrations
+
+        filter_key (str): if provided, split the subset of demonstration keys stored
+            under mask/@filter_key instead of the full set of demonstrations
+    """
+
+    # retrieve number of demos
+    f = h5py.File(hdf5_path, "r")
+    if filter_key is not None:
+        print("using filter key: {}".format(filter_key))
+        demos = sorted([elem.decode("utf-8") for elem in np.array(f["mask/{}".format(filter_key)])])
+    else:
+        demos = sorted(list(f["data"].keys()))
+    num_demos = len(demos)
+    f.close()
+
+    # get random split
+    num_demos = len(demos)
+    num_val = int(val_ratio * num_demos)
+    mask = np.zeros(num_demos)
+    mask[:num_val] = 1.
+    np.random.shuffle(mask)
+    mask = mask.astype(int)
+    train_inds = (1 - mask).nonzero()[0]
+    valid_inds = mask.nonzero()[0]
+    train_keys = [demos[i] for i in train_inds]
+    valid_keys = [demos[i] for i in valid_inds]
+    print("{} validation demonstrations out of {} total demonstrations.".format(num_val, num_demos))
+
+    # pass mask to generate split
+    name_1 = "train"
+    name_2 = "valid"
+    if filter_key is not None:
+        name_1 = "{}_{}".format(filter_key, name_1)
+        name_2 = "{}_{}".format(filter_key, name_2)
+
+    train_lengths = create_hdf5_filter_key(hdf5_path=hdf5_path, demo_keys=train_keys, key_name=name_1)
+    valid_lengths = create_hdf5_filter_key(hdf5_path=hdf5_path, demo_keys=valid_keys, key_name=name_2)
+
+    print("Total number of train samples: {}".format(np.sum(train_lengths)))
+    print("Average number of train samples {}".format(np.mean(train_lengths)))
+
+    print("Total number of valid samples: {}".format(np.sum(valid_lengths)))
+    print("Average number of valid samples {}".format(np.mean(valid_lengths)))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        help="path to hdf5 dataset",
+    )
+    parser.add_argument(
+        "--filter_key",
+        type=str,
+        default=None,
+        help="if provided, split the subset of trajectories in the file that correspond to\
+            this filter key into a training and validation set of trajectories, instead of\
+            splitting the full set of trajectories",
+    )
+    parser.add_argument(
+        "--ratio",
+        type=float,
+        default=0.1,
+        help="validation ratio, in (0, 1)"
+    )
+    args = parser.parse_args()
+
+    # seed to make sure results are consistent
+    np.random.seed(0)
+
+    split_train_val_from_hdf5(args.dataset, val_ratio=args.ratio, filter_key=args.filter_key)

aloha-devel/robomimic/scripts/train.py

@@ -0,0 +1,512 @@
+"""
+The main entry point for training policies.
+
+Args:
+    config (str): path to a config json that will be used to override the default settings.
+        If omitted, default settings are used. This is the preferred way to run experiments.
+
+    algo (str): name of the algorithm to run. Only needs to be provided if @config is not
+        provided.
+
+    name (str): if provided, override the experiment name defined in the config
+
+    dataset (str): if provided, override the dataset path defined in the config
+
+    debug (bool): set this flag to run a quick training run for debugging purposes    
+"""
+
+import argparse
+import json
+import numpy as np
+import time
+import os
+import shutil
+import psutil
+import sys
+import socket
+import traceback
+
+from collections import OrderedDict
+
+import torch
+from torch.utils.data import DataLoader
+
+import robomimic
+import robomimic.utils.train_utils as TrainUtils
+import robomimic.utils.torch_utils as TorchUtils
+import robomimic.utils.obs_utils as ObsUtils
+import robomimic.utils.env_utils as EnvUtils
+import robomimic.utils.file_utils as FileUtils
+from robomimic.config import config_factory
+from robomimic.algo import algo_factory, RolloutPolicy
+from robomimic.utils.log_utils import PrintLogger, DataLogger, flush_warnings
+
+
+def train(config, device):
+    """
+    Train a model using the algorithm.
+    """
+
+    # first set seeds
+    np.random.seed(config.train.seed)
+    torch.manual_seed(config.train.seed)
+
+    # set num workers
+    torch.set_num_threads(1)
+
+    print("\n============= New Training Run with Config =============")
+    print(config)
+    print("")
+    log_dir, ckpt_dir, video_dir, vis_dir = TrainUtils.get_exp_dir(config)
+
+    if config.experiment.logging.terminal_output_to_txt:
+        # log stdout and stderr to a text file
+        logger = PrintLogger(os.path.join(log_dir, 'log.txt'))
+        sys.stdout = logger
+        sys.stderr = logger
+
+    # read config to set up metadata for observation modalities (e.g. detecting rgb observations)
+    ObsUtils.initialize_obs_utils_with_config(config)
+
+    # extract the metadata and shape metadata across all datasets
+    env_meta_list = []
+    shape_meta_list = []
+    for dataset_cfg in config.train.data:
+        dataset_path = os.path.expanduser(dataset_cfg["path"])
+        ds_format = config.train.data_format
+        if not os.path.exists(dataset_path):
+            raise Exception("Dataset at provided path {} not found!".format(dataset_path))
+
+        # load basic metadata from training file
+        print("\n============= Loaded Environment Metadata =============")
+        env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=dataset_path, ds_format=ds_format)
+
+        # populate language instruction for env in env_meta
+        env_meta["lang"] = dataset_cfg.get("lang", "dummy")
+
+        # update env meta if applicable
+        from robomimic.utils.script_utils import deep_update
+        deep_update(env_meta, config.experiment.env_meta_update_dict)
+        env_meta_list.append(env_meta)
+
+        shape_meta = FileUtils.get_shape_metadata_from_dataset(
+            dataset_path=dataset_path,
+            action_keys=config.train.action_keys,
+            all_obs_keys=config.all_obs_keys,
+            ds_format=ds_format,
+            verbose=True
+        )
+        shape_meta_list.append(shape_meta)
+
+    if config.experiment.env is not None:
+        env_meta["env_name"] = config.experiment.env
+        print("=" * 30 + "\n" + "Replacing Env to {}\n".format(env_meta["env_name"]) + "=" * 30)
+
+    # create environment
+    envs = OrderedDict()
+    if config.experiment.rollout.enabled:
+        # create environments for validation runs
+        # env_names = [env_meta["env_name"]]
+
+        # # disable this feature for now
+        # if config.experiment.additional_envs is not None:
+        #     raise NotImplementedError
+        #     for name in config.experiment.additional_envs:
+        #         env_names.append(name)
+
+        for (dataset_i, dataset_cfg) in enumerate(config.train.data):
+            do_eval = dataset_cfg.get("eval", True)
+            if do_eval is not True:
+                continue
+            env_meta = env_meta_list[dataset_i]
+            shape_meta = shape_meta_list[dataset_i]
+            env_name = env_meta["env_name"]
+            
+            def create_env(env_i=0):
+                env_kwargs = dict(
+                    env_meta=env_meta,
+                    env_name=env_name,
+                    render=False,
+                    render_offscreen=config.experiment.render_video,
+                    use_image_obs=shape_meta["use_images"],
+                    # seed=config.train.seed * 1000 + env_i # TODO: add seeding across environments
+                )
+                env = EnvUtils.create_env_from_metadata(**env_kwargs)
+                # handle environment wrappers
+                env = EnvUtils.wrap_env_from_config(env, config=config)  # apply environment warpper, if applicable
+
+                return env
+
+            if config.experiment.rollout.batched:
+                from tianshou.env import SubprocVectorEnv
+                env_fns = [lambda env_i=i: create_env(env_i) for i in range(config.experiment.rollout.num_batch_envs)]
+                env = SubprocVectorEnv(env_fns)
+                env_name = env.get_env_attr(key="name", id=0)[0]
+            else:
+                env = create_env()
+                env_name = env.name
+            
+            envs[env_name] = env
+            print(env)
+
+    print("")
+
+    # setup for a new training run
+    data_logger = DataLogger(
+        log_dir,
+        config,
+        log_tb=config.experiment.logging.log_tb,
+        log_wandb=config.experiment.logging.log_wandb,
+    )
+    model = algo_factory(
+        algo_name=config.algo_name,
+        config=config,
+        obs_key_shapes=shape_meta_list[0]["all_shapes"],
+        ac_dim=shape_meta_list[0]["ac_dim"],
+        device=device,
+    )
+    
+    # save the config as a json file
+    with open(os.path.join(log_dir, '..', 'config.json'), 'w') as outfile:
+        json.dump(config, outfile, indent=4)
+
+    # if checkpoint is specified, load in model weights
+    ckpt_path = config.experiment.ckpt_path
+    if ckpt_path is not None:
+        print("LOADING MODEL WEIGHTS FROM " + ckpt_path)
+        from robomimic.utils.file_utils import maybe_dict_from_checkpoint
+        ckpt_dict = maybe_dict_from_checkpoint(ckpt_path=ckpt_path)
+        model.deserialize(ckpt_dict["model"])
+
+    print("\n============= Model Summary =============")
+    print(model)  # print model summary
+    print("")
+
+    # load training data
+    trainset, validset = TrainUtils.load_data_for_training(
+        config, obs_keys=shape_meta["all_obs_keys"])
+    train_sampler = trainset.get_dataset_sampler()
+    print("\n============= Training Dataset =============")
+    print(trainset)
+    print("")
+    if validset is not None:
+        print("\n============= Validation Dataset =============")
+        print(validset)
+        print("")
+
+    # maybe retreve statistics for normalizing observations
+    obs_normalization_stats = None
+    if config.train.hdf5_normalize_obs:
+        obs_normalization_stats = trainset.get_obs_normalization_stats()
+
+    # maybe retreve statistics for normalizing actions
+    action_normalization_stats = trainset.get_action_normalization_stats()
+
+    # initialize data loaders
+    train_loader = DataLoader(
+        dataset=trainset,
+        sampler=train_sampler,
+        batch_size=config.train.batch_size,
+        shuffle=(train_sampler is None),
+        num_workers=config.train.num_data_workers,
+        drop_last=True
+    )
+
+    if config.experiment.validate:
+        # cap num workers for validation dataset at 1
+        num_workers = min(config.train.num_data_workers, 1)
+        valid_sampler = validset.get_dataset_sampler()
+        valid_loader = DataLoader(
+            dataset=validset,
+            sampler=valid_sampler,
+            batch_size=config.train.batch_size,
+            shuffle=(valid_sampler is None),
+            num_workers=num_workers,
+            drop_last=True
+        )
+    else:
+        valid_loader = None
+
+    # print all warnings before training begins
+    print("*" * 50)
+    print("Warnings generated by robomimic have been duplicated here (from above) for convenience. Please check them carefully.")
+    flush_warnings()
+    print("*" * 50)
+    print("")
+
+    # main training loop
+    best_valid_loss = None
+    best_return = {k: -np.inf for k in envs} if config.experiment.rollout.enabled else None
+    best_success_rate = {k: -1. for k in envs} if config.experiment.rollout.enabled else None
+    last_ckpt_time = time.time()
+
+    # number of learning steps per epoch (defaults to a full dataset pass)
+    train_num_steps = config.experiment.epoch_every_n_steps
+    valid_num_steps = config.experiment.validation_epoch_every_n_steps
+
+    for epoch in range(1, config.train.num_epochs + 1): # epoch numbers start at 1
+        step_log = TrainUtils.run_epoch(
+            model=model,
+            data_loader=train_loader,
+            epoch=epoch,
+            num_steps=train_num_steps,
+            obs_normalization_stats=obs_normalization_stats,
+        )
+        model.on_epoch_end(epoch)
+
+        # setup checkpoint path
+        epoch_ckpt_name = "model_epoch_{}".format(epoch)
+
+        # check for recurring checkpoint saving conditions
+        should_save_ckpt = False
+        if config.experiment.save.enabled:
+            time_check = (config.experiment.save.every_n_seconds is not None) and \
+                (time.time() - last_ckpt_time > config.experiment.save.every_n_seconds)
+            epoch_check = (config.experiment.save.every_n_epochs is not None) and \
+                (epoch > 0) and (epoch % config.experiment.save.every_n_epochs == 0)
+            epoch_list_check = (epoch in config.experiment.save.epochs)
+            should_save_ckpt = (time_check or epoch_check or epoch_list_check)
+        ckpt_reason = None
+        if should_save_ckpt:
+            last_ckpt_time = time.time()
+            ckpt_reason = "time"
+
+        print("Train Epoch {}".format(epoch))
+        print(json.dumps(step_log, sort_keys=True, indent=4))
+        for k, v in step_log.items():
+            if k.startswith("Time_"):
+                data_logger.record("Timing_Stats/Train_{}".format(k[5:]), v, epoch)
+            else:
+                data_logger.record("Train/{}".format(k), v, epoch)
+
+        # Evaluate the model on validation set
+        if config.experiment.validate:
+            with torch.no_grad():
+                step_log = TrainUtils.run_epoch(model=model, data_loader=valid_loader, epoch=epoch, validate=True, num_steps=valid_num_steps)
+            for k, v in step_log.items():
+                if k.startswith("Time_"):
+                    data_logger.record("Timing_Stats/Valid_{}".format(k[5:]), v, epoch)
+                else:
+                    data_logger.record("Valid/{}".format(k), v, epoch)
+
+            print("Validation Epoch {}".format(epoch))
+            print(json.dumps(step_log, sort_keys=True, indent=4))
+
+            # save checkpoint if achieve new best validation loss
+            valid_check = "Loss" in step_log
+            if valid_check and (best_valid_loss is None or (step_log["Loss"] <= best_valid_loss)):
+                best_valid_loss = step_log["Loss"]
+                if config.experiment.save.enabled and config.experiment.save.on_best_validation:
+                    epoch_ckpt_name += "_best_validation_{}".format(best_valid_loss)
+                    should_save_ckpt = True
+                    ckpt_reason = "valid" if ckpt_reason is None else ckpt_reason
+
+        # Evaluate the model by by running rollouts
+
+        # do rollouts at fixed rate or if it's time to save a new ckpt
+        video_paths = None
+        rollout_check = (epoch % config.experiment.rollout.rate == 0) or (should_save_ckpt and ckpt_reason == "time")
+        if config.experiment.rollout.enabled and (epoch > config.experiment.rollout.warmstart) and rollout_check:
+            # wrap model as a RolloutPolicy to prepare for rollouts
+            rollout_model = RolloutPolicy(
+                model,
+                obs_normalization_stats=obs_normalization_stats,
+                action_normalization_stats=action_normalization_stats,
+            )
+
+            num_episodes = config.experiment.rollout.n
+            all_rollout_logs, video_paths = TrainUtils.rollout_with_stats(
+                policy=rollout_model,
+                envs=envs,
+                horizon=config.experiment.rollout.horizon,
+                use_goals=config.use_goals,
+                num_episodes=num_episodes,
+                render=False,
+                video_dir=video_dir if config.experiment.render_video else None,
+                epoch=epoch,
+                video_skip=config.experiment.get("video_skip", 5),
+                terminate_on_success=config.experiment.rollout.terminate_on_success,
+            )
+
+            # summarize results from rollouts to tensorboard and terminal
+            for env_name in all_rollout_logs:
+                rollout_logs = all_rollout_logs[env_name]
+                for k, v in rollout_logs.items():
+                    if k.startswith("Time_"):
+                        data_logger.record("Timing_Stats/Rollout_{}_{}".format(env_name, k[5:]), v, epoch)
+                    else:
+                        data_logger.record("Rollout/{}/{}".format(k, env_name), v, epoch, log_stats=True)
+
+                print("\nEpoch {} Rollouts took {}s (avg) with results:".format(epoch, rollout_logs["time"]))
+                print('Env: {}'.format(env_name))
+                print(json.dumps(rollout_logs, sort_keys=True, indent=4))
+
+            # checkpoint and video saving logic
+            updated_stats = TrainUtils.should_save_from_rollout_logs(
+                all_rollout_logs=all_rollout_logs,
+                best_return=best_return,
+                best_success_rate=best_success_rate,
+                epoch_ckpt_name=epoch_ckpt_name,
+                save_on_best_rollout_return=config.experiment.save.on_best_rollout_return,
+                save_on_best_rollout_success_rate=config.experiment.save.on_best_rollout_success_rate,
+            )
+            best_return = updated_stats["best_return"]
+            best_success_rate = updated_stats["best_success_rate"]
+            epoch_ckpt_name = updated_stats["epoch_ckpt_name"]
+            should_save_ckpt = (config.experiment.save.enabled and updated_stats["should_save_ckpt"]) or should_save_ckpt
+            if updated_stats["ckpt_reason"] is not None:
+                ckpt_reason = updated_stats["ckpt_reason"]
+
+        # check if we need to save model MSE
+        should_save_mse = False
+        if config.experiment.mse.enabled:
+            if config.experiment.mse.every_n_epochs is not None and epoch % config.experiment.mse.every_n_epochs == 0:
+                should_save_mse = True
+            if config.experiment.mse.on_save_ckpt and should_save_ckpt:
+                should_save_mse = True
+        if should_save_mse:
+            print("Computing MSE ...")
+            if config.experiment.mse.visualize:
+                save_vis_dir = os.path.join(vis_dir, epoch_ckpt_name)
+            else:
+                save_vis_dir = None
+            mse_log, vis_log = model.compute_mse_visualize(
+                trainset,
+                validset,
+                num_samples=config.experiment.mse.num_samples,
+                savedir=save_vis_dir,
+            )    
+            for k, v in mse_log.items():
+                data_logger.record("{}".format(k), v, epoch)
+            
+            for k, v in vis_log.items():
+                data_logger.record("{}".format(k), v, epoch, data_type='image')
+
+
+            print("MSE Log Epoch {}".format(epoch))
+            print(json.dumps(mse_log, sort_keys=True, indent=4))
+        
+        # # Only keep saved videos if the ckpt should be saved (but not because of validation score)
+        # should_save_video = (should_save_ckpt and (ckpt_reason != "valid")) or config.experiment.keep_all_videos
+        # if video_paths is not None and not should_save_video:
+        #     for env_name in video_paths:
+        #         os.remove(video_paths[env_name])
+
+        # Save model checkpoints based on conditions (success rate, validation loss, etc)
+        if should_save_ckpt:    
+            TrainUtils.save_model(
+                model=model,
+                config=config,
+                env_meta=env_meta,
+                shape_meta=shape_meta,
+                ckpt_path=os.path.join(ckpt_dir, epoch_ckpt_name + ".pth"),
+                obs_normalization_stats=obs_normalization_stats,
+                action_normalization_stats=action_normalization_stats,
+            )
+
+        # Finally, log memory usage in MB
+        process = psutil.Process(os.getpid())
+        mem_usage = int(process.memory_info().rss / 1000000)
+        data_logger.record("System/RAM Usage (MB)", mem_usage, epoch)
+        print("\nEpoch {} Memory Usage: {} MB\n".format(epoch, mem_usage))
+
+    # terminate logging
+    data_logger.close()
+
+
+def main(args):
+
+    if args.config is not None:
+        ext_cfg = json.load(open(args.config, 'r'))
+        config = config_factory(ext_cfg["algo_name"])
+        # update config with external json - this will throw errors if
+        # the external config has keys not present in the base algo config
+        with config.values_unlocked():
+            config.update(ext_cfg)
+    else:
+        config = config_factory(args.algo)
+
+    if args.dataset is not None:
+        config.train.data = args.dataset
+
+    if args.name is not None:
+        config.experiment.name = args.name
+
+    # get torch device
+    device = TorchUtils.get_torch_device(try_to_use_cuda=config.train.cuda)
+
+    # maybe modify config for debugging purposes
+    if args.debug:
+        # shrink length of training to test whether this run is likely to crash
+        config.unlock()
+        config.lock_keys()
+
+        # train and validate (if enabled) for 3 gradient steps, for 2 epochs
+        config.experiment.epoch_every_n_steps = 3
+        config.experiment.validation_epoch_every_n_steps = 3
+        config.train.num_epochs = 2
+
+        # if rollouts are enabled, try 2 rollouts at end of each epoch, with 10 environment steps
+        config.experiment.rollout.rate = 1
+        config.experiment.rollout.n = 2
+        config.experiment.rollout.horizon = 10
+
+        # send output to a temporary directory
+        config.train.output_dir = "/tmp/tmp_trained_models"
+
+    # lock config to prevent further modifications and ensure missing keys raise errors
+    config.lock()
+
+    # catch error during training and print it
+    res_str = "finished run successfully!"
+    try:
+        train(config, device=device)
+    except Exception as e:
+        res_str = "run failed with error:\n{}\n\n{}".format(e, traceback.format_exc())
+    print(res_str)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    # External config file that overwrites default config
+    parser.add_argument(
+        "--config",
+        type=str,
+        default=None,
+        help="(optional) path to a config json that will be used to override the default settings. \
+            If omitted, default settings are used. This is the preferred way to run experiments.",
+    )
+
+    # Algorithm Name
+    parser.add_argument(
+        "--algo",
+        type=str,
+        help="(optional) name of algorithm to run. Only needs to be provided if --config is not provided",
+    )
+
+    # Experiment Name (for tensorboard, saving models, etc.)
+    parser.add_argument(
+        "--name",
+        type=str,
+        default=None,
+        help="(optional) if provided, override the experiment name defined in the config",
+    )
+
+    # Dataset path, to override the one in the config
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default=None,
+        help="(optional) if provided, override the dataset path defined in the config",
+    )
+
+    # debug mode
+    parser.add_argument(
+        "--debug",
+        action='store_true',
+        help="set this flag to run a quick training run for debugging purposes"
+    )
+
+    args = parser.parse_args()
+    main(args)

camera_ws/build/realsense-ros/realsense2_camera/CMakeFiles/realsense2_camera.dir/src/base_realsense_node.cpp.o

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

camera_ws/build/realsense-ros/realsense2_camera/CMakeFiles/realsense2_camera.dir/src/realsense_node_factory.cpp.o

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

camera_ws/build/realsense-ros/realsense2_camera/CMakeFiles/realsense2_camera.dir/src/t265_realsense_node.cpp.o

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

camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/d2c_viewer.cpp.o

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

camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/ob_camera_info.cpp.o

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

camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/ob_camera_node.cpp.o

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

camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/ob_camera_node_factory.cpp.o

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

camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/point_cloud_proc/point_cloud_xyz.cpp.o

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

camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/point_cloud_proc/point_cloud_xyzrgb.cpp.o

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

camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/ros_service.cpp.o

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

camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/utils.cpp.o

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

camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera.dir/src/uvc_camera_driver.cpp.o

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

camera_ws/build/ros_astra_camera/CMakeFiles/astra_camera_node.dir/src/main.cpp.o

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

camera_ws/devel/lib/astra_camera/astra_camera_node

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

camera_ws/devel/lib/libastra_camera.so

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