diff --git a/.gitignore b/.gitignore index d674f037c..4fce82ba9 100644 --- a/.gitignore +++ b/.gitignore @@ -15,34 +15,21 @@ web-ui/user_code/tmp/webgpu-examples ObservableRobotRemote data/db/code_annotation.db data/alan_kay/images -notebooks/robot_remote_control -notebooks/deno-webgpu -notebooks/2_perception -notebooks/course-content-observable-data-app + vendor/ #dist/ -notebooks/deno-webgpu/ + ./data/alan_kay data/alan_kay scripts/infra/uv_example -notebooks/2_perception/test_lottiesrm -notebooks/2_perception/lottie_2 -notebooks/diffusion_ + bun.lockb .air.toml scripts/infra/supporting_site/chrome-extension/ scripts/infra/Dot-files/ data.jsong -notebooks/flow/ -notebooks/deno-webgpu -notebooks/robot_remote_control -notebooks/deno-webgpu/obs3-breakout/ -notebooks/course-content-observable-data-app -notebooks/3lang-notebook/ -notebooks/3_planning-prediction/ -notebooks/2_perception/pixels2svg/ -notebooks/2_perception/ + alan_kay static/pdfs diff --git a/.replit b/.replit new file mode 100644 index 000000000..f823a32a9 --- /dev/null +++ b/.replit @@ -0,0 +1,9 @@ +entrypoint = "web-ui/my-app/src/index.ts"" + +[nix] +channel = "stable-24_05" + + +[deployment] +run = ["bun", "run", "web-ui/my-app/src/index.ts"] +deploymentTarget = "cloudrun" diff --git a/notebooks_py/1_hardware/robot_move.py b/notebooks/1_hardware/robot_move.py similarity index 100% rename from notebooks_py/1_hardware/robot_move.py rename to notebooks/1_hardware/robot_move.py diff --git a/notebooks/2_perception/readme.md b/notebooks/2_perception/readme.md new file mode 100644 index 000000000..ce4a9ee09 --- /dev/null +++ b/notebooks/2_perception/readme.md @@ -0,0 +1 @@ +asdlfkanslfknasd \ No newline at end of file diff --git a/notebooks/3_planning-prediction/readme.md b/notebooks/3_planning-prediction/readme.md new file mode 100644 index 000000000..ce4a9ee09 --- /dev/null +++ b/notebooks/3_planning-prediction/readme.md @@ -0,0 +1 @@ +asdlfkanslfknasd \ No newline at end of file diff --git a/notebooks_py/3lang-notebook/51af771b6ed99dfe@388.js b/notebooks/3lang-notebook/51af771b6ed99dfe@388.js similarity index 100% rename from notebooks_py/3lang-notebook/51af771b6ed99dfe@388.js rename to notebooks/3lang-notebook/51af771b6ed99dfe@388.js diff --git a/notebooks_py/3lang-notebook/5611c97c55aa59b7@1061.js b/notebooks/3lang-notebook/5611c97c55aa59b7@1061.js similarity index 100% rename from notebooks_py/3lang-notebook/5611c97c55aa59b7@1061.js rename to notebooks/3lang-notebook/5611c97c55aa59b7@1061.js diff --git a/notebooks_py/3lang-notebook/README.md b/notebooks/3lang-notebook/README.md similarity index 100% rename from notebooks_py/3lang-notebook/README.md rename to notebooks/3lang-notebook/README.md diff --git a/notebooks_py/3lang-notebook/__index.html b/notebooks/3lang-notebook/__index.html similarity index 100% rename from notebooks_py/3lang-notebook/__index.html rename to notebooks/3lang-notebook/__index.html diff --git a/notebooks_py/3lang-notebook/e93997d5089d7165@2303.js b/notebooks/3lang-notebook/e93997d5089d7165@2303.js similarity index 100% rename from notebooks_py/3lang-notebook/e93997d5089d7165@2303.js rename to notebooks/3lang-notebook/e93997d5089d7165@2303.js diff --git a/notebooks_py/3lang-notebook/files/c051fbc024553912e31968b35e537d4ad3592201b5f8e7bd13fd9d02e38599c5d541a704d0858c676328babb3e5c9c35dd7c6d67240090d094882a1cad8eece4.gif b/notebooks/3lang-notebook/files/c051fbc024553912e31968b35e537d4ad3592201b5f8e7bd13fd9d02e38599c5d541a704d0858c676328babb3e5c9c35dd7c6d67240090d094882a1cad8eece4.gif similarity index 100% rename from notebooks_py/3lang-notebook/files/c051fbc024553912e31968b35e537d4ad3592201b5f8e7bd13fd9d02e38599c5d541a704d0858c676328babb3e5c9c35dd7c6d67240090d094882a1cad8eece4.gif rename to notebooks/3lang-notebook/files/c051fbc024553912e31968b35e537d4ad3592201b5f8e7bd13fd9d02e38599c5d541a704d0858c676328babb3e5c9c35dd7c6d67240090d094882a1cad8eece4.gif diff --git a/notebooks_py/3lang-notebook/files/cbdae86fbaad8452a4c7a5b6613a97e7b4846aa365213e2ac9dc3f4a631feee61fbbc6fdddd20444b0094adde2019aab464a50b5d15dcaefddef9163ba4f4329.jpeg b/notebooks/3lang-notebook/files/cbdae86fbaad8452a4c7a5b6613a97e7b4846aa365213e2ac9dc3f4a631feee61fbbc6fdddd20444b0094adde2019aab464a50b5d15dcaefddef9163ba4f4329.jpeg similarity index 100% rename from notebooks_py/3lang-notebook/files/cbdae86fbaad8452a4c7a5b6613a97e7b4846aa365213e2ac9dc3f4a631feee61fbbc6fdddd20444b0094adde2019aab464a50b5d15dcaefddef9163ba4f4329.jpeg rename to notebooks/3lang-notebook/files/cbdae86fbaad8452a4c7a5b6613a97e7b4846aa365213e2ac9dc3f4a631feee61fbbc6fdddd20444b0094adde2019aab464a50b5d15dcaefddef9163ba4f4329.jpeg diff --git a/notebooks_py/3lang-notebook/files/cda1fa4c2ef9d968f0c2ace4d37caa95431025144bc8aa45d363d5cfff03dd8d21cac6c842b947d1cecb8ff92bd2f270e4ed4f3f9059bb845192a47462796329.json b/notebooks/3lang-notebook/files/cda1fa4c2ef9d968f0c2ace4d37caa95431025144bc8aa45d363d5cfff03dd8d21cac6c842b947d1cecb8ff92bd2f270e4ed4f3f9059bb845192a47462796329.json similarity index 100% rename from notebooks_py/3lang-notebook/files/cda1fa4c2ef9d968f0c2ace4d37caa95431025144bc8aa45d363d5cfff03dd8d21cac6c842b947d1cecb8ff92bd2f270e4ed4f3f9059bb845192a47462796329.json rename to notebooks/3lang-notebook/files/cda1fa4c2ef9d968f0c2ace4d37caa95431025144bc8aa45d363d5cfff03dd8d21cac6c842b947d1cecb8ff92bd2f270e4ed4f3f9059bb845192a47462796329.json diff --git a/notebooks_py/3lang-notebook/files/ed3b20404ebd695193d70ec3b8486e1e4b0bc9e7634c7e5b86ed9f3ff7ecc03e65dab7e2c182403116aade573b46e0a40460f94ec03c924add8a033d8463d328.jpeg b/notebooks/3lang-notebook/files/ed3b20404ebd695193d70ec3b8486e1e4b0bc9e7634c7e5b86ed9f3ff7ecc03e65dab7e2c182403116aade573b46e0a40460f94ec03c924add8a033d8463d328.jpeg similarity index 100% rename from notebooks_py/3lang-notebook/files/ed3b20404ebd695193d70ec3b8486e1e4b0bc9e7634c7e5b86ed9f3ff7ecc03e65dab7e2c182403116aade573b46e0a40460f94ec03c924add8a033d8463d328.jpeg rename to notebooks/3lang-notebook/files/ed3b20404ebd695193d70ec3b8486e1e4b0bc9e7634c7e5b86ed9f3ff7ecc03e65dab7e2c182403116aade573b46e0a40460f94ec03c924add8a033d8463d328.jpeg diff --git a/notebooks_py/3lang-notebook/files/ff8d2292b6871fd29c4ffd50f030890c885a59e629d8061eccfc245f73cf5e598c019fbd098accf77000fa8f4eeb8c8c709820ce139aabd8fb80370cf764a586.jpeg b/notebooks/3lang-notebook/files/ff8d2292b6871fd29c4ffd50f030890c885a59e629d8061eccfc245f73cf5e598c019fbd098accf77000fa8f4eeb8c8c709820ce139aabd8fb80370cf764a586.jpeg similarity index 100% rename from notebooks_py/3lang-notebook/files/ff8d2292b6871fd29c4ffd50f030890c885a59e629d8061eccfc245f73cf5e598c019fbd098accf77000fa8f4eeb8c8c709820ce139aabd8fb80370cf764a586.jpeg rename to notebooks/3lang-notebook/files/ff8d2292b6871fd29c4ffd50f030890c885a59e629d8061eccfc245f73cf5e598c019fbd098accf77000fa8f4eeb8c8c709820ce139aabd8fb80370cf764a586.jpeg diff --git a/notebooks_py/3lang-notebook/index.html b/notebooks/3lang-notebook/index.html similarity index 100% rename from notebooks_py/3lang-notebook/index.html rename to notebooks/3lang-notebook/index.html diff --git a/notebooks_py/3lang-notebook/index.js b/notebooks/3lang-notebook/index.js similarity index 100% rename from notebooks_py/3lang-notebook/index.js rename to notebooks/3lang-notebook/index.js diff --git a/notebooks_py/3lang-notebook/inspector.css b/notebooks/3lang-notebook/inspector.css similarity index 100% rename from notebooks_py/3lang-notebook/inspector.css rename to notebooks/3lang-notebook/inspector.css diff --git a/notebooks_py/3lang-notebook/package.json b/notebooks/3lang-notebook/package.json similarity index 100% rename from notebooks_py/3lang-notebook/package.json rename to notebooks/3lang-notebook/package.json diff --git a/notebooks_py/3lang-notebook/runtime/.eslintignore b/notebooks/3lang-notebook/runtime/.eslintignore similarity index 100% rename from notebooks_py/3lang-notebook/runtime/.eslintignore rename to notebooks/3lang-notebook/runtime/.eslintignore diff --git a/notebooks_py/3lang-notebook/runtime/.eslintrc.json b/notebooks/3lang-notebook/runtime/.eslintrc.json similarity index 100% rename from notebooks_py/3lang-notebook/runtime/.eslintrc.json rename to notebooks/3lang-notebook/runtime/.eslintrc.json diff --git a/notebooks_py/3lang-notebook/runtime/.gitignore b/notebooks/3lang-notebook/runtime/.gitignore similarity index 100% rename from notebooks_py/3lang-notebook/runtime/.gitignore rename to notebooks/3lang-notebook/runtime/.gitignore diff --git a/notebooks_py/3lang-notebook/runtime/CONTRIBUTING.md b/notebooks/3lang-notebook/runtime/CONTRIBUTING.md similarity index 100% rename from notebooks_py/3lang-notebook/runtime/CONTRIBUTING.md rename to notebooks/3lang-notebook/runtime/CONTRIBUTING.md diff --git a/notebooks_py/3lang-notebook/runtime/LICENSE b/notebooks/3lang-notebook/runtime/LICENSE similarity index 100% rename from notebooks_py/3lang-notebook/runtime/LICENSE rename to notebooks/3lang-notebook/runtime/LICENSE diff --git a/notebooks_py/3lang-notebook/runtime/README.md b/notebooks/3lang-notebook/runtime/README.md similarity index 100% rename from notebooks_py/3lang-notebook/runtime/README.md rename to notebooks/3lang-notebook/runtime/README.md diff --git a/notebooks_py/3lang-notebook/runtime/package.json b/notebooks/3lang-notebook/runtime/package.json similarity index 100% rename from notebooks_py/3lang-notebook/runtime/package.json rename to notebooks/3lang-notebook/runtime/package.json diff --git a/notebooks_py/3lang-notebook/runtime/rollup.config.js b/notebooks/3lang-notebook/runtime/rollup.config.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/rollup.config.js rename to notebooks/3lang-notebook/runtime/rollup.config.js diff --git a/notebooks_py/3lang-notebook/runtime/src/array.js b/notebooks/3lang-notebook/runtime/src/array.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/array.js rename to notebooks/3lang-notebook/runtime/src/array.js diff --git a/notebooks_py/3lang-notebook/runtime/src/constant.js b/notebooks/3lang-notebook/runtime/src/constant.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/constant.js rename to notebooks/3lang-notebook/runtime/src/constant.js diff --git a/notebooks_py/3lang-notebook/runtime/src/errors.js b/notebooks/3lang-notebook/runtime/src/errors.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/errors.js rename to notebooks/3lang-notebook/runtime/src/errors.js diff --git a/notebooks_py/3lang-notebook/runtime/src/generatorish.js b/notebooks/3lang-notebook/runtime/src/generatorish.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/generatorish.js rename to notebooks/3lang-notebook/runtime/src/generatorish.js diff --git a/notebooks_py/3lang-notebook/runtime/src/identity.js b/notebooks/3lang-notebook/runtime/src/identity.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/identity.js rename to notebooks/3lang-notebook/runtime/src/identity.js diff --git a/notebooks_py/3lang-notebook/runtime/src/index.js b/notebooks/3lang-notebook/runtime/src/index.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/index.js rename to notebooks/3lang-notebook/runtime/src/index.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/.eslintignore b/notebooks/3lang-notebook/runtime/src/inspector/.eslintignore similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/.eslintignore rename to notebooks/3lang-notebook/runtime/src/inspector/.eslintignore diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/.eslintrc.json b/notebooks/3lang-notebook/runtime/src/inspector/.eslintrc.json similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/.eslintrc.json rename to notebooks/3lang-notebook/runtime/src/inspector/.eslintrc.json diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/.github/workflows/nodejs.yml b/notebooks/3lang-notebook/runtime/src/inspector/.github/workflows/nodejs.yml similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/.github/workflows/nodejs.yml rename to notebooks/3lang-notebook/runtime/src/inspector/.github/workflows/nodejs.yml diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/.gitignore b/notebooks/3lang-notebook/runtime/src/inspector/.gitignore similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/.gitignore rename to notebooks/3lang-notebook/runtime/src/inspector/.gitignore diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/.prettierrc b/notebooks/3lang-notebook/runtime/src/inspector/.prettierrc similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/.prettierrc rename to notebooks/3lang-notebook/runtime/src/inspector/.prettierrc diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/LICENSE b/notebooks/3lang-notebook/runtime/src/inspector/LICENSE similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/LICENSE rename to notebooks/3lang-notebook/runtime/src/inspector/LICENSE diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/README.md b/notebooks/3lang-notebook/runtime/src/inspector/README.md similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/README.md rename to notebooks/3lang-notebook/runtime/src/inspector/README.md diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/demo.html b/notebooks/3lang-notebook/runtime/src/inspector/demo.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/demo.html rename to notebooks/3lang-notebook/runtime/src/inspector/demo.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/package.json b/notebooks/3lang-notebook/runtime/src/inspector/package.json similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/package.json rename to notebooks/3lang-notebook/runtime/src/inspector/package.json diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/rollup.config.js b/notebooks/3lang-notebook/runtime/src/inspector/rollup.config.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/rollup.config.js rename to notebooks/3lang-notebook/runtime/src/inspector/rollup.config.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/array.js b/notebooks/3lang-notebook/runtime/src/inspector/src/array.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/array.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/array.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/collapsed.js b/notebooks/3lang-notebook/runtime/src/inspector/src/collapsed.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/collapsed.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/collapsed.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/dispatch.js b/notebooks/3lang-notebook/runtime/src/inspector/src/dispatch.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/dispatch.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/dispatch.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/expanded.js b/notebooks/3lang-notebook/runtime/src/inspector/src/expanded.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/expanded.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/expanded.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/formatDate.js b/notebooks/3lang-notebook/runtime/src/inspector/src/formatDate.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/formatDate.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/formatDate.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/formatError.js b/notebooks/3lang-notebook/runtime/src/inspector/src/formatError.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/formatError.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/formatError.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/formatRegExp.js b/notebooks/3lang-notebook/runtime/src/inspector/src/formatRegExp.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/formatRegExp.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/formatRegExp.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/formatString.js b/notebooks/3lang-notebook/runtime/src/inspector/src/formatString.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/formatString.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/formatString.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/formatSymbol.js b/notebooks/3lang-notebook/runtime/src/inspector/src/formatSymbol.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/formatSymbol.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/formatSymbol.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/immutable.js b/notebooks/3lang-notebook/runtime/src/inspector/src/immutable.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/immutable.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/immutable.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/index.js b/notebooks/3lang-notebook/runtime/src/inspector/src/index.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/index.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/index.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/inspect.js b/notebooks/3lang-notebook/runtime/src/inspector/src/inspect.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/inspect.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/inspect.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/inspectFunction.js b/notebooks/3lang-notebook/runtime/src/inspector/src/inspectFunction.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/inspectFunction.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/inspectFunction.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/inspectName.js b/notebooks/3lang-notebook/runtime/src/inspector/src/inspectName.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/inspectName.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/inspectName.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/object.js b/notebooks/3lang-notebook/runtime/src/inspector/src/object.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/object.js rename to notebooks/3lang-notebook/runtime/src/inspector/src/object.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/src/style.css b/notebooks/3lang-notebook/runtime/src/inspector/src/style.css similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/src/style.css rename to notebooks/3lang-notebook/runtime/src/inspector/src/style.css diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/.eslintrc.json b/notebooks/3lang-notebook/runtime/src/inspector/test/.eslintrc.json similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/.eslintrc.json rename to notebooks/3lang-notebook/runtime/src/inspector/test/.eslintrc.json diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/inspector-test.js b/notebooks/3lang-notebook/runtime/src/inspector/test/inspector-test.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/inspector-test.js rename to notebooks/3lang-notebook/runtime/src/inspector/test/inspector-test.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/inspectors.js b/notebooks/3lang-notebook/runtime/src/inspector/test/inspectors.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/inspectors.js rename to notebooks/3lang-notebook/runtime/src/inspector/test/inspectors.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/into-test.js b/notebooks/3lang-notebook/runtime/src/inspector/test/into-test.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/into-test.js rename to notebooks/3lang-notebook/runtime/src/inspector/test/into-test.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/jsdom.js b/notebooks/3lang-notebook/runtime/src/inspector/test/jsdom.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/jsdom.js rename to notebooks/3lang-notebook/runtime/src/inspector/test/jsdom.js diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledArray.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledArray.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledArray.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledArray.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledArrayOpen.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledArrayOpen.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledArrayOpen.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledArrayOpen.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledElement.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledElement.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledElement.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledElement.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableList.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableList.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableList.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableList.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableListOpen.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableListOpen.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableListOpen.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableListOpen.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableMap.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableMap.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableMap.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableMap.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableMapOpen.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableMapOpen.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableMapOpen.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableMapOpen.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableRecord.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableRecord.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableRecord.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableRecord.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableRecordOpen.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableRecordOpen.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableRecordOpen.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableRecordOpen.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableSet.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableSet.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableSet.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableSet.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableSetOpen.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableSetOpen.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableSetOpen.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledImmutableSetOpen.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledIntoMultiple.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledIntoMultiple.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledIntoMultiple.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledIntoMultiple.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledMap.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledMap.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledMap.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledMap.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledMapProto.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledMapProto.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledMapProto.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledMapProto.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledNumber.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledNumber.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledNumber.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledNumber.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledSet.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledSet.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledSet.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledSet.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledSetProto.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledSetProto.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledSetProto.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledSetProto.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledString.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledString.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledString.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledString.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledString10Lines.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledString10Lines.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledString10Lines.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledString10Lines.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledString21Lines.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledString21Lines.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledString21Lines.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledString21Lines.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledString30Lines.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledString30Lines.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledString30Lines.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledString30Lines.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledString30LinesOpen.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledString30LinesOpen.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/fulfilledString30LinesOpen.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/fulfilledString30LinesOpen.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/pending.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/pending.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/pending.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/pending.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/rejectedError.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/rejectedError.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/rejectedError.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/rejectedError.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/test/output/rejectedString.html b/notebooks/3lang-notebook/runtime/src/inspector/test/output/rejectedString.html similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/test/output/rejectedString.html rename to notebooks/3lang-notebook/runtime/src/inspector/test/output/rejectedString.html diff --git a/notebooks_py/3lang-notebook/runtime/src/inspector/yarn.lock b/notebooks/3lang-notebook/runtime/src/inspector/yarn.lock similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/inspector/yarn.lock rename to notebooks/3lang-notebook/runtime/src/inspector/yarn.lock diff --git a/notebooks_py/3lang-notebook/runtime/src/module.js b/notebooks/3lang-notebook/runtime/src/module.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/module.js rename to notebooks/3lang-notebook/runtime/src/module.js diff --git a/notebooks_py/3lang-notebook/runtime/src/noop.js b/notebooks/3lang-notebook/runtime/src/noop.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/noop.js rename to notebooks/3lang-notebook/runtime/src/noop.js diff --git a/notebooks_py/3lang-notebook/runtime/src/rethrow.js b/notebooks/3lang-notebook/runtime/src/rethrow.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/rethrow.js rename to notebooks/3lang-notebook/runtime/src/rethrow.js diff --git a/notebooks_py/3lang-notebook/runtime/src/runtime.js b/notebooks/3lang-notebook/runtime/src/runtime.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/runtime.js rename to notebooks/3lang-notebook/runtime/src/runtime.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/.eslintignore b/notebooks/3lang-notebook/runtime/src/stdlib/.eslintignore similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/.eslintignore rename to notebooks/3lang-notebook/runtime/src/stdlib/.eslintignore diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/.eslintrc.json b/notebooks/3lang-notebook/runtime/src/stdlib/.eslintrc.json similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/.eslintrc.json rename to notebooks/3lang-notebook/runtime/src/stdlib/.eslintrc.json diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/.github/workflows/nodejs.yml b/notebooks/3lang-notebook/runtime/src/stdlib/.github/workflows/nodejs.yml similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/.github/workflows/nodejs.yml rename to notebooks/3lang-notebook/runtime/src/stdlib/.github/workflows/nodejs.yml diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/.gitignore b/notebooks/3lang-notebook/runtime/src/stdlib/.gitignore similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/.gitignore rename to notebooks/3lang-notebook/runtime/src/stdlib/.gitignore diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/.prettierrc b/notebooks/3lang-notebook/runtime/src/stdlib/.prettierrc similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/.prettierrc rename to notebooks/3lang-notebook/runtime/src/stdlib/.prettierrc diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/CONTRIBUTING.md b/notebooks/3lang-notebook/runtime/src/stdlib/CONTRIBUTING.md similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/CONTRIBUTING.md rename to notebooks/3lang-notebook/runtime/src/stdlib/CONTRIBUTING.md diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/LICENSE b/notebooks/3lang-notebook/runtime/src/stdlib/LICENSE similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/LICENSE rename to notebooks/3lang-notebook/runtime/src/stdlib/LICENSE diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/README.md b/notebooks/3lang-notebook/runtime/src/stdlib/README.md similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/README.md rename to notebooks/3lang-notebook/runtime/src/stdlib/README.md diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/package.json b/notebooks/3lang-notebook/runtime/src/stdlib/package.json similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/package.json rename to notebooks/3lang-notebook/runtime/src/stdlib/package.json diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/rollup.config.js b/notebooks/3lang-notebook/runtime/src/stdlib/rollup.config.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/rollup.config.js rename to notebooks/3lang-notebook/runtime/src/stdlib/rollup.config.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/arquero.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/arquero.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/arquero.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/arquero.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/arrow.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/arrow.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/arrow.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/arrow.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dependencies.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dependencies.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dependencies.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dependencies.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dependency.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dependency.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dependency.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dependency.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/canvas.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dom/canvas.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/canvas.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dom/canvas.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/context2d.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dom/context2d.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/context2d.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dom/context2d.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/download.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dom/download.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/download.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dom/download.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/element.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dom/element.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/element.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dom/element.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/index.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dom/index.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/index.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dom/index.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/input.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dom/input.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/input.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dom/input.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/range.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dom/range.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/range.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dom/range.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/select.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dom/select.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/select.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dom/select.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/svg.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dom/svg.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/svg.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dom/svg.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/text.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dom/text.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/text.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dom/text.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/uid.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/dom/uid.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/dom/uid.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/dom/uid.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/duckdb.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/duckdb.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/duckdb.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/duckdb.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/fileAttachment.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/fileAttachment.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/fileAttachment.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/fileAttachment.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/files/buffer.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/files/buffer.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/files/buffer.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/files/buffer.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/files/index.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/files/index.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/files/index.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/files/index.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/files/text.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/files/text.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/files/text.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/files/text.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/files/url.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/files/url.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/files/url.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/files/url.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/disposable.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/generators/disposable.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/disposable.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/generators/disposable.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/filter.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/generators/filter.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/filter.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/generators/filter.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/index.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/generators/index.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/index.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/generators/index.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/input.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/generators/input.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/input.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/generators/input.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/map.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/generators/map.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/map.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/generators/map.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/observe.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/generators/observe.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/observe.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/generators/observe.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/queue.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/generators/queue.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/queue.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/generators/queue.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/range.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/generators/range.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/range.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/generators/range.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/valueAt.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/generators/valueAt.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/valueAt.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/generators/valueAt.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/worker.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/generators/worker.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/generators/worker.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/generators/worker.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/html.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/html.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/html.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/html.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/index.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/index.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/index.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/index.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/leaflet.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/leaflet.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/leaflet.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/leaflet.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/library.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/library.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/library.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/library.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/md.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/md.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/md.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/md.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/mermaid.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/mermaid.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/mermaid.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/mermaid.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/mutable.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/mutable.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/mutable.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/mutable.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/now.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/now.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/now.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/now.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/promises/delay.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/promises/delay.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/promises/delay.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/promises/delay.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/promises/index.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/promises/index.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/promises/index.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/promises/index.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/promises/tick.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/promises/tick.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/promises/tick.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/promises/tick.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/promises/when.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/promises/when.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/promises/when.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/promises/when.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/require.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/require.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/require.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/require.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/resolve.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/resolve.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/resolve.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/resolve.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/sqlite.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/sqlite.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/sqlite.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/sqlite.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/svg.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/svg.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/svg.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/svg.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/table.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/table.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/table.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/table.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/template.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/template.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/template.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/template.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/tex.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/tex.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/tex.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/tex.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/that.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/that.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/that.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/that.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/vegalite.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/vegalite.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/vegalite.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/vegalite.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/width.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/width.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/width.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/width.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/src/xlsx.js b/notebooks/3lang-notebook/runtime/src/stdlib/src/xlsx.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/src/xlsx.js rename to notebooks/3lang-notebook/runtime/src/stdlib/src/xlsx.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/stdlib.sublime-project b/notebooks/3lang-notebook/runtime/src/stdlib/stdlib.sublime-project similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/stdlib.sublime-project rename to notebooks/3lang-notebook/runtime/src/stdlib/stdlib.sublime-project diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/test/.eslintrc.json b/notebooks/3lang-notebook/runtime/src/stdlib/test/.eslintrc.json similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/test/.eslintrc.json rename to notebooks/3lang-notebook/runtime/src/stdlib/test/.eslintrc.json diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/test/DOM/uid-test.js b/notebooks/3lang-notebook/runtime/src/stdlib/test/DOM/uid-test.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/test/DOM/uid-test.js rename to notebooks/3lang-notebook/runtime/src/stdlib/test/DOM/uid-test.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/test/fileAttachments-test.js b/notebooks/3lang-notebook/runtime/src/stdlib/test/fileAttachments-test.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/test/fileAttachments-test.js rename to notebooks/3lang-notebook/runtime/src/stdlib/test/fileAttachments-test.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/test/generators/disposable-test.js b/notebooks/3lang-notebook/runtime/src/stdlib/test/generators/disposable-test.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/test/generators/disposable-test.js rename to notebooks/3lang-notebook/runtime/src/stdlib/test/generators/disposable-test.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/test/generators/generators-test.js b/notebooks/3lang-notebook/runtime/src/stdlib/test/generators/generators-test.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/test/generators/generators-test.js rename to notebooks/3lang-notebook/runtime/src/stdlib/test/generators/generators-test.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/test/index-test.js b/notebooks/3lang-notebook/runtime/src/stdlib/test/index-test.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/test/index-test.js rename to notebooks/3lang-notebook/runtime/src/stdlib/test/index-test.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/test/invalidation.js b/notebooks/3lang-notebook/runtime/src/stdlib/test/invalidation.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/test/invalidation.js rename to notebooks/3lang-notebook/runtime/src/stdlib/test/invalidation.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/test/promises/delay-test.js b/notebooks/3lang-notebook/runtime/src/stdlib/test/promises/delay-test.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/test/promises/delay-test.js rename to notebooks/3lang-notebook/runtime/src/stdlib/test/promises/delay-test.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/test/query-test.js b/notebooks/3lang-notebook/runtime/src/stdlib/test/query-test.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/test/query-test.js rename to notebooks/3lang-notebook/runtime/src/stdlib/test/query-test.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/test/table-test.js b/notebooks/3lang-notebook/runtime/src/stdlib/test/table-test.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/test/table-test.js rename to notebooks/3lang-notebook/runtime/src/stdlib/test/table-test.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/test/xlsx-test.js b/notebooks/3lang-notebook/runtime/src/stdlib/test/xlsx-test.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/test/xlsx-test.js rename to notebooks/3lang-notebook/runtime/src/stdlib/test/xlsx-test.js diff --git a/notebooks_py/3lang-notebook/runtime/src/stdlib/yarn.lock b/notebooks/3lang-notebook/runtime/src/stdlib/yarn.lock similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/stdlib/yarn.lock rename to notebooks/3lang-notebook/runtime/src/stdlib/yarn.lock diff --git a/notebooks_py/3lang-notebook/runtime/src/variable.js b/notebooks/3lang-notebook/runtime/src/variable.js similarity index 100% rename from notebooks_py/3lang-notebook/runtime/src/variable.js rename to notebooks/3lang-notebook/runtime/src/variable.js diff --git a/notebooks_py/3lang-notebook/runtime/yarn.lock b/notebooks/3lang-notebook/runtime/yarn.lock similarity index 100% rename from notebooks_py/3lang-notebook/runtime/yarn.lock rename to notebooks/3lang-notebook/runtime/yarn.lock diff --git a/notebooks_py/3lang-notebook/scratch/README.md b/notebooks/3lang-notebook/scratch/README.md similarity index 100% rename from notebooks_py/3lang-notebook/scratch/README.md rename to notebooks/3lang-notebook/scratch/README.md diff --git a/notebooks_py/3lang-notebook/scratch/index.js b/notebooks/3lang-notebook/scratch/index.js similarity index 100% rename from notebooks_py/3lang-notebook/scratch/index.js rename to notebooks/3lang-notebook/scratch/index.js diff --git a/notebooks_py/3lang-notebook/scratch/inspector.css b/notebooks/3lang-notebook/scratch/inspector.css similarity index 100% rename from notebooks_py/3lang-notebook/scratch/inspector.css rename to notebooks/3lang-notebook/scratch/inspector.css diff --git a/notebooks_py/3lang-notebook/scratch/package.json b/notebooks/3lang-notebook/scratch/package.json similarity index 100% rename from notebooks_py/3lang-notebook/scratch/package.json rename to notebooks/3lang-notebook/scratch/package.json diff --git a/notebooks_py/3lang-notebook/scratch/runtime.js b/notebooks/3lang-notebook/scratch/runtime.js similarity index 100% rename from notebooks_py/3lang-notebook/scratch/runtime.js rename to notebooks/3lang-notebook/scratch/runtime.js diff --git a/notebooks_py/figma_to_animation.py b/notebooks/figma_to_animation.py similarity index 100% rename from notebooks_py/figma_to_animation.py rename to notebooks/figma_to_animation.py diff --git a/notebooks_py/diffusion_/readme.md b/notebooks/readme.md similarity index 80% rename from notebooks_py/diffusion_/readme.md rename to notebooks/readme.md index 0480af413..10fda9ecd 100644 --- a/notebooks_py/diffusion_/readme.md +++ b/notebooks/readme.md @@ -1,4 +1,18 @@ -https://chatgpt.com/share/6706919d-429c-8013-9fba-6cfa870b1a2c +# notebook - parse this to include more repos + +duffison 3d policy + +1. https://github.com/fastai/diffusion-nbs/blob/master/stable_diffusion.ipynb + +2. https://github.com/real-stanford/diffusion_policy + +3. prediction notebook - norvig + +4. + + + + https://chatgpt.com/share/6706919d-429c-8013-9fba-6cfa870b1a2c most important paper of the year - find like all of them - vote diff --git a/notebooks_py/3_planning-prediction/Untitled.ipynb b/notebooks_py/3_planning-prediction/Untitled.ipynb deleted file mode 100644 index 2c9914794..000000000 --- a/notebooks_py/3_planning-prediction/Untitled.ipynb +++ /dev/null @@ -1 +0,0 @@ -{"cells":[{"id":"727807f0-44c0-4335-ad0e-8ae799be88b4","cell_type":"code","execution_count":1,"outputs":[{"output_type":"stream","name":"stdout","text":["123",""]}],"source":["\n","print(123)\n","\n","\n","\n","\n","# all you need is \n","# 1. one zed frame\n","# 2. one bluetooth call or call to irobot_sdk_edu which can be found on github lol \n","# 3. arm wtf lol \n"],"metadata":{}}],"metadata":{"kernelspec":{"name":"python3","display_name":"python3","language":"python"},"language_info":{"name":"python","codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","pygments_lexer":"ipython3"}},"nbformat":4,"nbformat_minor":5} \ No newline at end of file diff --git a/notebooks_py/3_planning-prediction/poetry.lock b/notebooks_py/3_planning-prediction/poetry.lock deleted file mode 100644 index ddb15267a..000000000 --- a/notebooks_py/3_planning-prediction/poetry.lock +++ /dev/null @@ -1,7 +0,0 @@ -# This file is automatically @generated by Poetry 1.5.1 and should not be changed by hand. -package = [] - -[metadata] -lock-version = "2.0" -python-versions = "^3.10" -content-hash = "53f2eabc9c26446fbcc00d348c47878e118afc2054778c3c803a0a8028af27d9" diff --git a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/Dockerfile b/notebooks_py/3_planning-prediction/puppeteer-js-renderer/Dockerfile deleted file mode 100644 index c24cb29b8..000000000 --- a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/Dockerfile +++ /dev/null @@ -1,12 +0,0 @@ -FROM zenika/alpine-chrome:83-with-node-12 - -USER root -ENV NODE_ENV=production -WORKDIR /src - -COPY package*.json ./ -RUN npm install - -COPY . . -EXPOSE 8080 -CMD ["node" , "index.js"] diff --git a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/README.md b/notebooks_py/3_planning-prediction/puppeteer-js-renderer/README.md deleted file mode 100644 index 99446efb0..000000000 --- a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/README.md +++ /dev/null @@ -1,140 +0,0 @@ -# Puppeteer JS Renderer on Fly - - -JavaScript is the bane of a web scraper's life. Scraping is all about extracting data from a web page and JavaScript is there adding content, hiding blocks, and moving the DOM around. Reading the HTML from the server is just not enough. What you ideally want is a way to run all that JavaScript on the page so you can see what's left after that. Then you can get down to some serious scraping. - -There are tools to do this out there but most have their own complications or restrictions that make them difficult to deploy. Puppeteer has none of those problems and with [Fly](https://fly.io), you can deploy puppeteer based applications close to your users. At its core, this project is a puppeteer-based service. [Puppeteer](https://pptr.dev/) is a package that renders pages using a headless Chrome instance, executing the JavaScript within the page. - -## Quick Start - -A typical YouTube page will add the views count to itself only when the JavaScript on that page executes. These are the types of use cases where Puppeteer comes in very handy. Puppeteer can execute the JavaScript on the YouTube page and, from the rendered page, extract the view count and the title of the video. - -Lets get started - -### Install example project - -To install this project locally, execute the following commands: - -``` -git clone https://github.com/fly-examples/puppeteer-js-renderer.git -cd puppeteer-js-renderer -npm install -``` - -### Your first scraper - -To quickly try the project out after installing it, run the command below: - -``` -node yt-views.js -``` - -This will show something like below: - -``` -Pulling views from YouTube, please wait... -Luis Fonsi - Despacito ft. Daddy Yankee has 6,881,463,846 views -``` - -If you want to try any other video just pass the YouTube video URL as a parameter to the script like below: - -``` -node yt-views https://www.youtube.com/watch?v=XqZsoesa55w -``` - -It should give you an output like this: - -``` -Pulling views from YouTube, please wait... -Baby Shark Dance | Sing and Dance! | @Baby Shark Official | PINKFONG Songs for Children has 6,077,338,169 views -``` - -### What just happened here? - -To show the YouTube video views count, a small scraper written with [Axrio](https://www.npmjs.com/package/@geshan/axrio) npm package was executed. Axrio combines the popular [Axios](https://www.npmjs.com/package/axios) library and [Cheerio](https://www.npmjs.com/package/cheerio) to create a mini scraper. Axios is used to make requests and Cheerio acts like DOM navigator parsing the markup and giving us an API for traversing/manipulating the resulting data structure. You can kickstart a small scraper with Axrio too. - -The [yt-views.js](./yt-views) is a basic scraper which performs a GET request for the given YouTube URL with Puppeteer. This will return the final DOM after the page's JavaScript executes. Then it parses the title and views count out of the rendered page's markup and prints it on the console. - -It uses the puppeteer-js-renderer service that is already deployed and running on Fly.io. Have a look at he "[how to use it as a service](#how-to-use-it-as-a-service)" section for more information. To start your own instance on Fly.io jump to the "[how to deploy on fly.io](#how-to-deploy-it-on-flyio)" section. - -## Run locally - -If you have node installed on your machine, you already cloned this repository and ran `npm install`. Now run `npm start` to get this service running locally. - -The next step is to navigate to `http://localhost:8080/api/render?url=https://www.youtube.com/watch?v=kJQP7kiw5Fk` on your browser to view the JavaScript rendered output. - -### Run with Docker - -If you want to run with Docker, execute the following after you clone the repository: - -``` -cd puppeteer-js-renderer -docker-compose up -``` - -Then go to `http://localhost:8080/api/render?url=https://www.youtube.com/watch?v=kJQP7kiw5Fk` in your browser to see the output. - -## How to use it as a service - -If you want to use puppeteer-js-renderer for scraping, you can use the following URL on Fly.io: - -``` -https://js-renderer-fly.fly.dev/api/render?url=https://www.youtube.com/watch?v=kJQP7kiw5Fk -``` - -The YouTube URL is an example above, you can use any URL that needs javascript to be executed to be functional. - -### Styles broken - -Styles and images will look broken but the HTML tags will be there. Happy Web Scraping! - -## How to deploy it on fly.io - -[Fly.io](https://fly.io) is a platform for applications that need to run globally. It runs your code close to users and scales compute in cities where your app is busiest. Write your code, package it into a Docker image, deploy it to Fly's platform, and let that do all the work to keep your app snappy. - -Fly.io has great [documentation](https://fly.io/docs/) to get started. You can find a quick speed run showing how to get your app running closer to your users with this [guide](https://fly.io/docs/speedrun/). - -Please follow these steps to deploy your own puppeteer-js-renderer service on Fly.io: - -### Prerequisites - -1. [Install](https://fly.io/docs/getting-started/installing-flyctl/) the flyctl CLI command. -1. Register on fly with `flyctl auth signup`, if you already have a fly account log in with `flyctl auth login`. - -### Steps - -1. Clone this repo with `git clone git@github.com:fly-examples/puppeteer-js-renderer.git` if you are logged in with SSH support enabled. Otherwise try `git clone https://github.com/fly-examples/puppeteer-js-renderer.git`. -1. Then run `cd puppeteer-js-renderer`. -1. After that execute `flyctl init --dockerfile` and when asked for an app name, hit return to have one generated (unless there's a name you really want). I ran it with js-renderer-fly as the app name for this example. -1. Subsequently, you can select an organization. Usually this will be your first name-last name on the prompt. -1. It should create a fly.toml file in the project root (I have not committed it, it is in .gitignore). -1. Now run `flyctl deploy` to deploy the app. It will build the docker image, push it to the fly docker image registry and deploy it. In the process of deploying it, it will display information about the number of instances and their health. -1. You can then run `flyctl info`. This will give the details of the app including hostname. -1. You can view your app in the browser with `flyctl open`. For me, it opened `https://js-renderer-fly.fly.dev` and displays `{message: "alive"}`. -1. Add `/api/render?url=https://www.youtube.com/watch?v=kJQP7kiw5Fk` on the address bar of your browser after your puppeteer-js-renderer URL. Mine looked like `https://js-renderer-fly.fly.dev/api/render?url=https://www.youtube.com/watch?v=kJQP7kiw5Fk`. This will render the final DOM of that YouTube page after the JavaScript is executed. -1. Enjoy! - -The YouTube video URL is an example and it can, of course, be replaced with other web pages which need JavaScript to work. - -## More fly commands - -You can suspend your service with `flyctl suspend`; this will pause your service until you resume it. If you try `flyctl status` after suspending it will not show any instances running. Suspending the service means all resources allocated to run the service will be deallocated, resulting in the application running nowhere. To get the instances running again execute `flyctl resume`. - -### Fly default resources - -I wanted to see what resources were allocated to the App on Fly by default. The scale commands allowed me to find it out pretty easily like below: - -1. `flyctl scale show` - showed me VM Size: micro-2x -1. `flyctl scale vm` - showed me micro-2x is a 0.25 CPU cores with 512 MB of memory. - -If you want to increase CPU/memory or run more instances in a particular region please refer to the official Fly docs on [scaling](https://fly.io/docs/scaling/). - -### Flying the app on 3 continents - -Our service is running in one data center. For me, it's iad (Ashburn, Virginia) but yours will likely be different based on where you are working from. We can add instances around the world to speed up responses, let's get rolling: - -1. To see the regions available run `flyctl platform regions`, I could see regions all over the world from Oregon to Sydney. -1. Let's add an instance to Australia in Sydney, to do this run `flyctl regions add syd`, yes it is that easy. -1. Now check `flyctl status` and you will see an instance running in Sydney -1. Let’s add one more in Europe in Amsterdam with `flyctl regions add ams`. So now we are mostly covered with the app running in 3 continents. -1. Of course, you can run `flyctl status` again to see your app shining on 3 continents. diff --git a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/docker-compose.yml b/notebooks_py/3_planning-prediction/puppeteer-js-renderer/docker-compose.yml deleted file mode 100644 index c684af803..000000000 --- a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/docker-compose.yml +++ /dev/null @@ -1,12 +0,0 @@ -version: '3.5' -services: - js-renderer: - build: . - volumes: - - .:/src - ports: - - '8080:8080' - command: npm start - environment: - NODE_ENV: dev - PORT: 8080 diff --git a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/imgs/01fly-init.png b/notebooks_py/3_planning-prediction/puppeteer-js-renderer/imgs/01fly-init.png deleted file mode 100644 index a20166e07..000000000 Binary files a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/imgs/01fly-init.png and /dev/null differ diff --git a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/imgs/02fly-deploy.png b/notebooks_py/3_planning-prediction/puppeteer-js-renderer/imgs/02fly-deploy.png deleted file mode 100644 index a783d5607..000000000 Binary files a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/imgs/02fly-deploy.png and /dev/null differ diff --git a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/index.js b/notebooks_py/3_planning-prediction/puppeteer-js-renderer/index.js deleted file mode 100644 index f3051ae0f..000000000 --- a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/index.js +++ /dev/null @@ -1,50 +0,0 @@ -const express = require('express'); -const puppeteer = require('puppeteer-core'); - -const app = express(); - -console.log('hi') - -app.get('/', (req, res) => { - //res.json({ message: 'alive' }); - console.log('wtf') - res.send('hi') - res.end() -}); - -app.get('/api/render', async (req, res) => { - const url = req.query.url; - if (!url) { - return res.send('please provide url'); - } - try { - const browser = await puppeteer.launch( - { - executablePath: process.env.CHROME_BIN, - args: [ - // Required for Docker version of Puppeteer - '--no-sandbox', - '--disable-setuid-sandbox', - // This will write shared memory files into /tmp instead of /dev/shm, - // because Docker’s default for /dev/shm is 64MB - '--disable-dev-shm-usage' - ], - }); - - const page = await browser.newPage(); - await page.goto(url, {waitUntil: 'networkidle0'}); - const pageContent = await page.content(); - console.log(`Response first 200 chars from ${url} : ${pageContent.substring(0, 200)}`); - await browser.close(); - - res.send(pageContent); - } catch (err) { - console.log(`Error while fetching ${url} `, err); - res.send(`Error fetching ${url}`); - } -}); - -const PORT = process.env.PORT || 8080; -app.listen(PORT, () => { - console.log(`Listening on port ${PORT}`); -}); diff --git a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/package.json b/notebooks_py/3_planning-prediction/puppeteer-js-renderer/package.json deleted file mode 100644 index b5780ffb3..000000000 --- a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/package.json +++ /dev/null @@ -1,30 +0,0 @@ -{ - "name": "js-renderer", - "version": "1.0.0", - "description": "A online puppeteer service to render pages with javascript (js). Mainly useful for web scraping (not using splash).", - "main": "index.js", - "scripts": { - "test": "echo \"Error: no test specified\" && exit 1", - "start": "node index.js" - }, - "repository": { - "type": "git", - "url": "git+https://github.com/geshan/js-renderer.git" - }, - "keywords": [ - "javascript", - "nodejs", - "puppeteer" - ], - "author": "Geshan Manandhar", - "license": "ISC", - "bugs": { - "url": "https://github.com/geshan/js-renderer/issues" - }, - "homepage": "https://github.com/geshan/js-renderer#readme", - "dependencies": { - "@geshan/axrio": "^1.0.1", - "express": "^4.17.1", - "puppeteer-core": "^3.3.0" - } -} diff --git a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/yt-views.js b/notebooks_py/3_planning-prediction/puppeteer-js-renderer/yt-views.js deleted file mode 100644 index 56a9e7acf..000000000 --- a/notebooks_py/3_planning-prediction/puppeteer-js-renderer/yt-views.js +++ /dev/null @@ -1,17 +0,0 @@ -const axrio = require('@geshan/axrio'); - -(async function run() { - console.log(`Pulling views from YouTube, please wait...`); - try { - const ytVideoUrl = process.argv[2] ? process.argv[2] : 'https://www.youtube.com/watch?v=kJQP7kiw5Fk'; - const baseName = 'http://localhost:8080' - //`https://js-renderer-fly.fly.dev/api/render?url=${ytVideoUrl}` - const $ = await axrio.getPage(baseName+`/api/render?url=${ytVideoUrl}`, 12000); - const title = $('h1.title>yt-formatted-string').text(); - const views = $('span.view-count').text(); - console.log(`${title} has ${views}`); - } catch(e) { - console.log(`Error while fetching views: `, e); - } - process.exit(0); -})(); diff --git a/notebooks_py/3_planning-prediction/pyproject.toml b/notebooks_py/3_planning-prediction/pyproject.toml deleted file mode 100644 index 444f40ed4..000000000 --- a/notebooks_py/3_planning-prediction/pyproject.toml +++ /dev/null @@ -1,17 +0,0 @@ -[tool.poetry] -name = "sample-project" -version = "0.1.0" -description = "" -authors = ["Author Name "] -# readme = "README.md" -# license = "BSD" -packages = [ - { include = "sample_package" } -] - -[tool.poetry.dependencies] -python = "^3.10" - -[build-system] -requires = ["poetry-core"] -build-backend = "poetry.core.masonry.api" \ No newline at end of file diff --git a/notebooks_py/3_planning-prediction/sample_package/__init__.py b/notebooks_py/3_planning-prediction/sample_package/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/notebooks_py/3_planning-prediction/sample_package/__main__.py b/notebooks_py/3_planning-prediction/sample_package/__main__.py deleted file mode 100644 index 7576b1dcf..000000000 --- a/notebooks_py/3_planning-prediction/sample_package/__main__.py +++ /dev/null @@ -1,2 +0,0 @@ -if __name__ == "__main__": - print("Hello, world!") \ No newline at end of file diff --git a/notebooks_py/3_planning-prediction/scraper-wow/.github/workflows/playwright.yml b/notebooks_py/3_planning-prediction/scraper-wow/.github/workflows/playwright.yml deleted file mode 100644 index 467190be6..000000000 --- a/notebooks_py/3_planning-prediction/scraper-wow/.github/workflows/playwright.yml +++ /dev/null @@ -1,27 +0,0 @@ -name: Playwright Tests -on: - push: - branches: [ main, master ] - pull_request: - branches: [ main, master ] -jobs: - test: - timeout-minutes: 60 - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-node@v4 - with: - node-version: lts/* - - name: Install dependencies - run: npm ci - - name: Install Playwright Browsers - run: npx playwright install --with-deps - - name: Run Playwright tests - run: npx playwright test - - uses: actions/upload-artifact@v4 - if: always() - with: - name: playwright-report - path: playwright-report/ - retention-days: 30 diff --git a/notebooks_py/3_planning-prediction/scraper-wow/.gitignore b/notebooks_py/3_planning-prediction/scraper-wow/.gitignore deleted file mode 100644 index 68c5d18f0..000000000 --- a/notebooks_py/3_planning-prediction/scraper-wow/.gitignore +++ /dev/null @@ -1,5 +0,0 @@ -node_modules/ -/test-results/ -/playwright-report/ -/blob-report/ -/playwright/.cache/ diff --git a/notebooks_py/3_planning-prediction/scraper-wow/package-lock.json b/notebooks_py/3_planning-prediction/scraper-wow/package-lock.json deleted file mode 100644 index 363b6800f..000000000 --- a/notebooks_py/3_planning-prediction/scraper-wow/package-lock.json +++ /dev/null @@ -1,91 +0,0 @@ -{ - "name": "scraper-wow", - "version": "1.0.0", - "lockfileVersion": 3, - "requires": true, - "packages": { - "": { - "name": "scraper-wow", - "version": "1.0.0", - "license": "ISC", - "devDependencies": { - "@playwright/test": "^1.44.0", - "@types/node": "^20.12.12" - } - }, - "node_modules/@playwright/test": { - "version": "1.44.0", - "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.44.0.tgz", - "integrity": "sha512-rNX5lbNidamSUorBhB4XZ9SQTjAqfe5M+p37Z8ic0jPFBMo5iCtQz1kRWkEMg+rYOKSlVycpQmpqjSFq7LXOfg==", - "dev": true, - "dependencies": { - "playwright": "1.44.0" - }, - "bin": { - "playwright": "cli.js" - }, - "engines": { - "node": ">=16" - } - }, - "node_modules/@types/node": { - "version": "20.12.12", - "resolved": "https://registry.npmjs.org/@types/node/-/node-20.12.12.tgz", - "integrity": "sha512-eWLDGF/FOSPtAvEqeRAQ4C8LSA7M1I7i0ky1I8U7kD1J5ITyW3AsRhQrKVoWf5pFKZ2kILsEGJhsI9r93PYnOw==", - "dev": true, - "dependencies": { - "undici-types": "~5.26.4" - } - }, - "node_modules/fsevents": { - "version": "2.3.2", - "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", - "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", - "dev": true, - "hasInstallScript": true, - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": "^8.16.0 || ^10.6.0 || >=11.0.0" - } - }, - "node_modules/playwright": { - "version": "1.44.0", - "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.44.0.tgz", - "integrity": "sha512-F9b3GUCLQ3Nffrfb6dunPOkE5Mh68tR7zN32L4jCk4FjQamgesGay7/dAAe1WaMEGV04DkdJfcJzjoCKygUaRQ==", - "dev": true, - "dependencies": { - "playwright-core": "1.44.0" - }, - "bin": { - "playwright": "cli.js" - }, - "engines": { - "node": ">=16" - }, - "optionalDependencies": { - "fsevents": "2.3.2" - } - }, - "node_modules/playwright-core": { - "version": "1.44.0", - "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.44.0.tgz", - "integrity": "sha512-ZTbkNpFfYcGWohvTTl+xewITm7EOuqIqex0c7dNZ+aXsbrLj0qI8XlGKfPpipjm0Wny/4Lt4CJsWJk1stVS5qQ==", - "dev": true, - "bin": { - "playwright-core": "cli.js" - }, - "engines": { - "node": ">=16" - } - }, - "node_modules/undici-types": { - "version": "5.26.5", - "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-5.26.5.tgz", - "integrity": "sha512-JlCMO+ehdEIKqlFxk6IfVoAUVmgz7cU7zD/h9XZ0qzeosSHmUJVOzSQvvYSYWXkFXC+IfLKSIffhv0sVZup6pA==", - "dev": true - } - } -} diff --git a/notebooks_py/3_planning-prediction/scraper-wow/package.json b/notebooks_py/3_planning-prediction/scraper-wow/package.json deleted file mode 100644 index f5fee8b67..000000000 --- a/notebooks_py/3_planning-prediction/scraper-wow/package.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "name": "scraper-wow", - "version": "1.0.0", - "description": "", - "main": "index.js", - "scripts": {}, - "keywords": [], - "author": "", - "license": "ISC", - "devDependencies": { - "@playwright/test": "^1.44.0", - "@types/node": "^20.12.12" - } -} diff --git a/notebooks_py/3_planning-prediction/scraper-wow/playwright.config.ts b/notebooks_py/3_planning-prediction/scraper-wow/playwright.config.ts deleted file mode 100644 index 301801ee1..000000000 --- a/notebooks_py/3_planning-prediction/scraper-wow/playwright.config.ts +++ /dev/null @@ -1,77 +0,0 @@ -import { defineConfig, devices } from '@playwright/test'; - -/** - * Read environment variables from file. - * https://github.com/motdotla/dotenv - */ -// require('dotenv').config(); - -/** - * See https://playwright.dev/docs/test-configuration. - */ -export default defineConfig({ - testDir: './tests', - /* Run tests in files in parallel */ - fullyParallel: true, - /* Fail the build on CI if you accidentally left test.only in the source code. */ - forbidOnly: !!process.env.CI, - /* Retry on CI only */ - retries: process.env.CI ? 2 : 0, - /* Opt out of parallel tests on CI. */ - workers: process.env.CI ? 1 : undefined, - /* Reporter to use. See https://playwright.dev/docs/test-reporters */ - reporter: 'html', - /* Shared settings for all the projects below. See https://playwright.dev/docs/api/class-testoptions. */ - use: { - /* Base URL to use in actions like `await page.goto('/')`. */ - // baseURL: 'http://127.0.0.1:3000', - - /* Collect trace when retrying the failed test. See https://playwright.dev/docs/trace-viewer */ - trace: 'on-first-retry', - }, - - /* Configure projects for major browsers */ - projects: [ - { - name: 'chromium', - use: { ...devices['Desktop Chrome'] }, - }, - - { - name: 'firefox', - use: { ...devices['Desktop Firefox'] }, - }, - - { - name: 'webkit', - use: { ...devices['Desktop Safari'] }, - }, - - /* Test against mobile viewports. */ - // { - // name: 'Mobile Chrome', - // use: { ...devices['Pixel 5'] }, - // }, - // { - // name: 'Mobile Safari', - // use: { ...devices['iPhone 12'] }, - // }, - - /* Test against branded browsers. */ - // { - // name: 'Microsoft Edge', - // use: { ...devices['Desktop Edge'], channel: 'msedge' }, - // }, - // { - // name: 'Google Chrome', - // use: { ...devices['Desktop Chrome'], channel: 'chrome' }, - // }, - ], - - /* Run your local dev server before starting the tests */ - // webServer: { - // command: 'npm run start', - // url: 'http://127.0.0.1:3000', - // reuseExistingServer: !process.env.CI, - // }, -}); diff --git a/notebooks_py/3_planning-prediction/scraper-wow/tests-examples/demo-todo-app.spec.ts b/notebooks_py/3_planning-prediction/scraper-wow/tests-examples/demo-todo-app.spec.ts deleted file mode 100644 index 2fd6016fe..000000000 --- a/notebooks_py/3_planning-prediction/scraper-wow/tests-examples/demo-todo-app.spec.ts +++ /dev/null @@ -1,437 +0,0 @@ -import { test, expect, type Page } from '@playwright/test'; - -test.beforeEach(async ({ page }) => { - await page.goto('https://demo.playwright.dev/todomvc'); -}); - -const TODO_ITEMS = [ - 'buy some cheese', - 'feed the cat', - 'book a doctors appointment' -]; - -test.describe('New Todo', () => { - test('should allow me to add todo items', async ({ page }) => { - // create a new todo locator - const newTodo = page.getByPlaceholder('What needs to be done?'); - - // Create 1st todo. - await newTodo.fill(TODO_ITEMS[0]); - await newTodo.press('Enter'); - - // Make sure the list only has one todo item. - await expect(page.getByTestId('todo-title')).toHaveText([ - TODO_ITEMS[0] - ]); - - // Create 2nd todo. - await newTodo.fill(TODO_ITEMS[1]); - await newTodo.press('Enter'); - - // Make sure the list now has two todo items. - await expect(page.getByTestId('todo-title')).toHaveText([ - TODO_ITEMS[0], - TODO_ITEMS[1] - ]); - - await checkNumberOfTodosInLocalStorage(page, 2); - }); - - test('should clear text input field when an item is added', async ({ page }) => { - // create a new todo locator - const newTodo = page.getByPlaceholder('What needs to be done?'); - - // Create one todo item. - await newTodo.fill(TODO_ITEMS[0]); - await newTodo.press('Enter'); - - // Check that input is empty. - await expect(newTodo).toBeEmpty(); - await checkNumberOfTodosInLocalStorage(page, 1); - }); - - test('should append new items to the bottom of the list', async ({ page }) => { - // Create 3 items. - await createDefaultTodos(page); - - // create a todo count locator - const todoCount = page.getByTestId('todo-count') - - // Check test using different methods. - await expect(page.getByText('3 items left')).toBeVisible(); - await expect(todoCount).toHaveText('3 items left'); - await expect(todoCount).toContainText('3'); - await expect(todoCount).toHaveText(/3/); - - // Check all items in one call. - await expect(page.getByTestId('todo-title')).toHaveText(TODO_ITEMS); - await checkNumberOfTodosInLocalStorage(page, 3); - }); -}); - -test.describe('Mark all as completed', () => { - test.beforeEach(async ({ page }) => { - await createDefaultTodos(page); - await checkNumberOfTodosInLocalStorage(page, 3); - }); - - test.afterEach(async ({ page }) => { - await checkNumberOfTodosInLocalStorage(page, 3); - }); - - test('should allow me to mark all items as completed', async ({ page }) => { - // Complete all todos. - await page.getByLabel('Mark all as complete').check(); - - // Ensure all todos have 'completed' class. - await expect(page.getByTestId('todo-item')).toHaveClass(['completed', 'completed', 'completed']); - await checkNumberOfCompletedTodosInLocalStorage(page, 3); - }); - - test('should allow me to clear the complete state of all items', async ({ page }) => { - const toggleAll = page.getByLabel('Mark all as complete'); - // Check and then immediately uncheck. - await toggleAll.check(); - await toggleAll.uncheck(); - - // Should be no completed classes. - await expect(page.getByTestId('todo-item')).toHaveClass(['', '', '']); - }); - - test('complete all checkbox should update state when items are completed / cleared', async ({ page }) => { - const toggleAll = page.getByLabel('Mark all as complete'); - await toggleAll.check(); - await expect(toggleAll).toBeChecked(); - await checkNumberOfCompletedTodosInLocalStorage(page, 3); - - // Uncheck first todo. - const firstTodo = page.getByTestId('todo-item').nth(0); - await firstTodo.getByRole('checkbox').uncheck(); - - // Reuse toggleAll locator and make sure its not checked. - await expect(toggleAll).not.toBeChecked(); - - await firstTodo.getByRole('checkbox').check(); - await checkNumberOfCompletedTodosInLocalStorage(page, 3); - - // Assert the toggle all is checked again. - await expect(toggleAll).toBeChecked(); - }); -}); - -test.describe('Item', () => { - - test('should allow me to mark items as complete', async ({ page }) => { - // create a new todo locator - const newTodo = page.getByPlaceholder('What needs to be done?'); - - // Create two items. - for (const item of TODO_ITEMS.slice(0, 2)) { - await newTodo.fill(item); - await newTodo.press('Enter'); - } - - // Check first item. - const firstTodo = page.getByTestId('todo-item').nth(0); - await firstTodo.getByRole('checkbox').check(); - await expect(firstTodo).toHaveClass('completed'); - - // Check second item. - const secondTodo = page.getByTestId('todo-item').nth(1); - await expect(secondTodo).not.toHaveClass('completed'); - await secondTodo.getByRole('checkbox').check(); - - // Assert completed class. - await expect(firstTodo).toHaveClass('completed'); - await expect(secondTodo).toHaveClass('completed'); - }); - - test('should allow me to un-mark items as complete', async ({ page }) => { - // create a new todo locator - const newTodo = page.getByPlaceholder('What needs to be done?'); - - // Create two items. - for (const item of TODO_ITEMS.slice(0, 2)) { - await newTodo.fill(item); - await newTodo.press('Enter'); - } - - const firstTodo = page.getByTestId('todo-item').nth(0); - const secondTodo = page.getByTestId('todo-item').nth(1); - const firstTodoCheckbox = firstTodo.getByRole('checkbox'); - - await firstTodoCheckbox.check(); - await expect(firstTodo).toHaveClass('completed'); - await expect(secondTodo).not.toHaveClass('completed'); - await checkNumberOfCompletedTodosInLocalStorage(page, 1); - - await firstTodoCheckbox.uncheck(); - await expect(firstTodo).not.toHaveClass('completed'); - await expect(secondTodo).not.toHaveClass('completed'); - await checkNumberOfCompletedTodosInLocalStorage(page, 0); - }); - - test('should allow me to edit an item', async ({ page }) => { - await createDefaultTodos(page); - - const todoItems = page.getByTestId('todo-item'); - const secondTodo = todoItems.nth(1); - await secondTodo.dblclick(); - await expect(secondTodo.getByRole('textbox', { name: 'Edit' })).toHaveValue(TODO_ITEMS[1]); - await secondTodo.getByRole('textbox', { name: 'Edit' }).fill('buy some sausages'); - await secondTodo.getByRole('textbox', { name: 'Edit' }).press('Enter'); - - // Explicitly assert the new text value. - await expect(todoItems).toHaveText([ - TODO_ITEMS[0], - 'buy some sausages', - TODO_ITEMS[2] - ]); - await checkTodosInLocalStorage(page, 'buy some sausages'); - }); -}); - -test.describe('Editing', () => { - test.beforeEach(async ({ page }) => { - await createDefaultTodos(page); - await checkNumberOfTodosInLocalStorage(page, 3); - }); - - test('should hide other controls when editing', async ({ page }) => { - const todoItem = page.getByTestId('todo-item').nth(1); - await todoItem.dblclick(); - await expect(todoItem.getByRole('checkbox')).not.toBeVisible(); - await expect(todoItem.locator('label', { - hasText: TODO_ITEMS[1], - })).not.toBeVisible(); - await checkNumberOfTodosInLocalStorage(page, 3); - }); - - test('should save edits on blur', async ({ page }) => { - const todoItems = page.getByTestId('todo-item'); - await todoItems.nth(1).dblclick(); - await todoItems.nth(1).getByRole('textbox', { name: 'Edit' }).fill('buy some sausages'); - await todoItems.nth(1).getByRole('textbox', { name: 'Edit' }).dispatchEvent('blur'); - - await expect(todoItems).toHaveText([ - TODO_ITEMS[0], - 'buy some sausages', - TODO_ITEMS[2], - ]); - await checkTodosInLocalStorage(page, 'buy some sausages'); - }); - - test('should trim entered text', async ({ page }) => { - const todoItems = page.getByTestId('todo-item'); - await todoItems.nth(1).dblclick(); - await todoItems.nth(1).getByRole('textbox', { name: 'Edit' }).fill(' buy some sausages '); - await todoItems.nth(1).getByRole('textbox', { name: 'Edit' }).press('Enter'); - - await expect(todoItems).toHaveText([ - TODO_ITEMS[0], - 'buy some sausages', - TODO_ITEMS[2], - ]); - await checkTodosInLocalStorage(page, 'buy some sausages'); - }); - - test('should remove the item if an empty text string was entered', async ({ page }) => { - const todoItems = page.getByTestId('todo-item'); - await todoItems.nth(1).dblclick(); - await todoItems.nth(1).getByRole('textbox', { name: 'Edit' }).fill(''); - await todoItems.nth(1).getByRole('textbox', { name: 'Edit' }).press('Enter'); - - await expect(todoItems).toHaveText([ - TODO_ITEMS[0], - TODO_ITEMS[2], - ]); - }); - - test('should cancel edits on escape', async ({ page }) => { - const todoItems = page.getByTestId('todo-item'); - await todoItems.nth(1).dblclick(); - await todoItems.nth(1).getByRole('textbox', { name: 'Edit' }).fill('buy some sausages'); - await todoItems.nth(1).getByRole('textbox', { name: 'Edit' }).press('Escape'); - await expect(todoItems).toHaveText(TODO_ITEMS); - }); -}); - -test.describe('Counter', () => { - test('should display the current number of todo items', async ({ page }) => { - // create a new todo locator - const newTodo = page.getByPlaceholder('What needs to be done?'); - - // create a todo count locator - const todoCount = page.getByTestId('todo-count') - - await newTodo.fill(TODO_ITEMS[0]); - await newTodo.press('Enter'); - - await expect(todoCount).toContainText('1'); - - await newTodo.fill(TODO_ITEMS[1]); - await newTodo.press('Enter'); - await expect(todoCount).toContainText('2'); - - await checkNumberOfTodosInLocalStorage(page, 2); - }); -}); - -test.describe('Clear completed button', () => { - test.beforeEach(async ({ page }) => { - await createDefaultTodos(page); - }); - - test('should display the correct text', async ({ page }) => { - await page.locator('.todo-list li .toggle').first().check(); - await expect(page.getByRole('button', { name: 'Clear completed' })).toBeVisible(); - }); - - test('should remove completed items when clicked', async ({ page }) => { - const todoItems = page.getByTestId('todo-item'); - await todoItems.nth(1).getByRole('checkbox').check(); - await page.getByRole('button', { name: 'Clear completed' }).click(); - await expect(todoItems).toHaveCount(2); - await expect(todoItems).toHaveText([TODO_ITEMS[0], TODO_ITEMS[2]]); - }); - - test('should be hidden when there are no items that are completed', async ({ page }) => { - await page.locator('.todo-list li .toggle').first().check(); - await page.getByRole('button', { name: 'Clear completed' }).click(); - await expect(page.getByRole('button', { name: 'Clear completed' })).toBeHidden(); - }); -}); - -test.describe('Persistence', () => { - test('should persist its data', async ({ page }) => { - // create a new todo locator - const newTodo = page.getByPlaceholder('What needs to be done?'); - - for (const item of TODO_ITEMS.slice(0, 2)) { - await newTodo.fill(item); - await newTodo.press('Enter'); - } - - const todoItems = page.getByTestId('todo-item'); - const firstTodoCheck = todoItems.nth(0).getByRole('checkbox'); - await firstTodoCheck.check(); - await expect(todoItems).toHaveText([TODO_ITEMS[0], TODO_ITEMS[1]]); - await expect(firstTodoCheck).toBeChecked(); - await expect(todoItems).toHaveClass(['completed', '']); - - // Ensure there is 1 completed item. - await checkNumberOfCompletedTodosInLocalStorage(page, 1); - - // Now reload. - await page.reload(); - await expect(todoItems).toHaveText([TODO_ITEMS[0], TODO_ITEMS[1]]); - await expect(firstTodoCheck).toBeChecked(); - await expect(todoItems).toHaveClass(['completed', '']); - }); -}); - -test.describe('Routing', () => { - test.beforeEach(async ({ page }) => { - await createDefaultTodos(page); - // make sure the app had a chance to save updated todos in storage - // before navigating to a new view, otherwise the items can get lost :( - // in some frameworks like Durandal - await checkTodosInLocalStorage(page, TODO_ITEMS[0]); - }); - - test('should allow me to display active items', async ({ page }) => { - const todoItem = page.getByTestId('todo-item'); - await page.getByTestId('todo-item').nth(1).getByRole('checkbox').check(); - - await checkNumberOfCompletedTodosInLocalStorage(page, 1); - await page.getByRole('link', { name: 'Active' }).click(); - await expect(todoItem).toHaveCount(2); - await expect(todoItem).toHaveText([TODO_ITEMS[0], TODO_ITEMS[2]]); - }); - - test('should respect the back button', async ({ page }) => { - const todoItem = page.getByTestId('todo-item'); - await page.getByTestId('todo-item').nth(1).getByRole('checkbox').check(); - - await checkNumberOfCompletedTodosInLocalStorage(page, 1); - - await test.step('Showing all items', async () => { - await page.getByRole('link', { name: 'All' }).click(); - await expect(todoItem).toHaveCount(3); - }); - - await test.step('Showing active items', async () => { - await page.getByRole('link', { name: 'Active' }).click(); - }); - - await test.step('Showing completed items', async () => { - await page.getByRole('link', { name: 'Completed' }).click(); - }); - - await expect(todoItem).toHaveCount(1); - await page.goBack(); - await expect(todoItem).toHaveCount(2); - await page.goBack(); - await expect(todoItem).toHaveCount(3); - }); - - test('should allow me to display completed items', async ({ page }) => { - await page.getByTestId('todo-item').nth(1).getByRole('checkbox').check(); - await checkNumberOfCompletedTodosInLocalStorage(page, 1); - await page.getByRole('link', { name: 'Completed' }).click(); - await expect(page.getByTestId('todo-item')).toHaveCount(1); - }); - - test('should allow me to display all items', async ({ page }) => { - await page.getByTestId('todo-item').nth(1).getByRole('checkbox').check(); - await checkNumberOfCompletedTodosInLocalStorage(page, 1); - await page.getByRole('link', { name: 'Active' }).click(); - await page.getByRole('link', { name: 'Completed' }).click(); - await page.getByRole('link', { name: 'All' }).click(); - await expect(page.getByTestId('todo-item')).toHaveCount(3); - }); - - test('should highlight the currently applied filter', async ({ page }) => { - await expect(page.getByRole('link', { name: 'All' })).toHaveClass('selected'); - - //create locators for active and completed links - const activeLink = page.getByRole('link', { name: 'Active' }); - const completedLink = page.getByRole('link', { name: 'Completed' }); - await activeLink.click(); - - // Page change - active items. - await expect(activeLink).toHaveClass('selected'); - await completedLink.click(); - - // Page change - completed items. - await expect(completedLink).toHaveClass('selected'); - }); -}); - -async function createDefaultTodos(page: Page) { - // create a new todo locator - const newTodo = page.getByPlaceholder('What needs to be done?'); - - for (const item of TODO_ITEMS) { - await newTodo.fill(item); - await newTodo.press('Enter'); - } -} - -async function checkNumberOfTodosInLocalStorage(page: Page, expected: number) { - return await page.waitForFunction(e => { - return JSON.parse(localStorage['react-todos']).length === e; - }, expected); -} - -async function checkNumberOfCompletedTodosInLocalStorage(page: Page, expected: number) { - return await page.waitForFunction(e => { - return JSON.parse(localStorage['react-todos']).filter((todo: any) => todo.completed).length === e; - }, expected); -} - -async function checkTodosInLocalStorage(page: Page, title: string) { - return await page.waitForFunction(t => { - return JSON.parse(localStorage['react-todos']).map((todo: any) => todo.title).includes(t); - }, title); -} diff --git a/notebooks_py/3_planning-prediction/scraper-wow/tests/example.spec.ts b/notebooks_py/3_planning-prediction/scraper-wow/tests/example.spec.ts deleted file mode 100644 index 54a906a4e..000000000 --- a/notebooks_py/3_planning-prediction/scraper-wow/tests/example.spec.ts +++ /dev/null @@ -1,18 +0,0 @@ -import { test, expect } from '@playwright/test'; - -test('has title', async ({ page }) => { - await page.goto('https://playwright.dev/'); - - // Expect a title "to contain" a substring. - await expect(page).toHaveTitle(/Playwright/); -}); - -test('get started link', async ({ page }) => { - await page.goto('https://playwright.dev/'); - - // Click the get started link. - await page.getByRole('link', { name: 'Get started' }).click(); - - // Expects page to have a heading with the name of Installation. - await expect(page.getByRole('heading', { name: 'Installation' })).toBeVisible(); -}); diff --git a/notebooks_py/3_planning-prediction/start.sh b/notebooks_py/3_planning-prediction/start.sh deleted file mode 100644 index f956d4d71..000000000 --- a/notebooks_py/3_planning-prediction/start.sh +++ /dev/null @@ -1,3 +0,0 @@ -#!/usr/bin/env sh - -#jupyter notebook & diff --git a/notebooks_py/diffusion_/Jeremy's Notebook - Microsoft OneNote Online.pdf b/notebooks_py/diffusion_/Jeremy's Notebook - Microsoft OneNote Online.pdf deleted file mode 100644 index 775a3c4c6..000000000 Binary files a/notebooks_py/diffusion_/Jeremy's Notebook - Microsoft OneNote Online.pdf and /dev/null differ diff --git a/notebooks_py/diffusion_/Untitled.ipynb b/notebooks_py/diffusion_/Untitled.ipynb deleted file mode 100644 index 6c74f41ec..000000000 --- a/notebooks_py/diffusion_/Untitled.ipynb +++ /dev/null @@ -1,43 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 1, - "id": "ba047ad4-676b-4026-bc5d-6f76872e0e83", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "jeremey-hwaord-stablue-diffusion-diffusion-3d-policy\n" - ] - } - ], - "source": [ - "print('jeremey-hwaord-stablue-diffusion-diffusion-3d-policy')" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks_py/diffusion_/diffusion-nbs/.gitignore b/notebooks_py/diffusion_/diffusion-nbs/.gitignore deleted file mode 100644 index d5a023d48..000000000 --- a/notebooks_py/diffusion_/diffusion-nbs/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -sd-concept-output/ -.ipynb_checkpoints/ diff --git a/notebooks_py/diffusion_/diffusion-nbs/LICENSE b/notebooks_py/diffusion_/diffusion-nbs/LICENSE deleted file mode 100644 index 261eeb9e9..000000000 --- a/notebooks_py/diffusion_/diffusion-nbs/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/notebooks_py/diffusion_/diffusion-nbs/README.md b/notebooks_py/diffusion_/diffusion-nbs/README.md deleted file mode 100644 index 7ac80c70c..000000000 --- a/notebooks_py/diffusion_/diffusion-nbs/README.md +++ /dev/null @@ -1,2 +0,0 @@ -# diffusion-nbs -Getting started with diffusion diff --git a/notebooks_py/diffusion_/diffusion-nbs/learned_embeds.bin b/notebooks_py/diffusion_/diffusion-nbs/learned_embeds.bin deleted file mode 100644 index c14b13c2e..000000000 Binary files a/notebooks_py/diffusion_/diffusion-nbs/learned_embeds.bin and /dev/null differ diff --git a/notebooks_py/diffusion_/diffusion-nbs/requirements.txt b/notebooks_py/diffusion_/diffusion-nbs/requirements.txt deleted file mode 100644 index 1ef119635..000000000 --- a/notebooks_py/diffusion_/diffusion-nbs/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -fastai -diffusers -transformers -huggingface-hub diff --git a/notebooks_py/diffusion_/diffusion-nbs/suggested_tools.md b/notebooks_py/diffusion_/diffusion-nbs/suggested_tools.md deleted file mode 100644 index 066b60e6a..000000000 --- a/notebooks_py/diffusion_/diffusion-nbs/suggested_tools.md +++ /dev/null @@ -1,35 +0,0 @@ -# Tools and Notebooks - -The AI art community has been sharing notebooks and making tools to use these models at an ever-increasing rate, to the point where lists of available options [like this one from @pharmapsychotic](https://pharmapsychotic.com/tools.html) can be fairly overwhelming! -So, here are a chosen few that you might want to check out as a starting point. - -## Easy online tools - -[Dreamstudio](https://beta.dreamstudio.ai/) is a fast and easy online interface for stable diffusion. You can do ~200 generations for free and then credits cost about $0.01 per image. There is also an API which let's you use their servers for inference, which is what powers things like [this photoshop plugin](https://exchange.adobe.com/apps/cc/114117da/stable-diffusion) for quickly generating images without a local GPU. - -[Artbreeder's new beta collage feature](https://www.artbreeder.com/beta/collage) is a fun (and free) interface for image-to-image style generations. - -There are a number of other online tools (midjourney, Wombo, NighCafe), typically with some sort of subscription or membership needed. And if you don't mind the occasional wait there are also [spaces on HuggingFace](https://huggingface.co/spaces/stabilityai/stable-diffusion) that remain free and typically keep the interface simple. - -## Fancier Interfaces - -People in the community have rushed to create new interfaces with additional advanced options for generation. For example, [this stable diffusion webui](https://github.com/sd-webui/stable-diffusion-webui) by hlky has a choice of several slick interfaces. -It has tabs for multiple tasks, batch functionality, integration of things like textual inversion, support for managing a number of model versions and much more. Settings are saved in the metadata of the resulting images, making it easy to access or share your favourite presets. -This is great for running SD locally, but there is also a [colab notebook](https://colab.research.google.com/github/altryne/sd-webui-colab/blob/main/Stable_Diffusion_WebUi_Altryne.ipynb) if you'd prefer to use Google's GPUs. The [InvokeAI Stable Diffusion toolkit](https://invoke-ai.github.io/InvokeAI/) is another new alternative worth exploring, and features a number of tools as well as optimizations for running locally with as little as 4GB RAM. - -The other notable mention in terms of local interfaces is Softology's ['Visions of Chaos'](https://softology.pro/voc.htm) tool, which makes it fairly easy to run a stagering number different notebooks/pipelines locally on Windows. - -## Colab Notebooks - -There are hundreds of notebooks floating around adding different bits of functionality, but two of the most popular and fully-featured versions are - -[Deforum](https://deforum.github.io/) - a community-built notebook with lots of features (and many more planned) including settings for animations and video. [Browsing Twitter for results](https://twitter.com/search?q=%23DeforumDiffusion&src=typeahead_click) will show the kinds of trippy animation this can make - by including a depth estimation model in the pipeline they've made it easy to get all sorts of interesting perspective warping motion into animations. - -[Disco Diffusion](https://colab.research.google.com/github/alembics/disco-diffusion/blob/main/Disco_Diffusion.ipynb). Before we had Stable Diffusion, we had to use CLIP to guide unconditional diffusion models. This notebook shows how much you can do with that concept, producing amazing imagery and animations by steering the generation of an unconditional 512x512px diffusion model. - -## Keeping up with the space - -News tools appear daily. Rather than chasing every new thing, I recommend spending some time learning a couple of tools well. -YouTube channels [like this one](https://www.youtube.com/c/NerdyRodent/videos) are a good place to see new things in action and evaluate whether you really need that new feature. -Many tools also have a community where you can get involved sharing feedback or helping to add new features - keep an eye out for Discord links and get involved :) - diff --git a/notebooks_py/diffusion_/diffusion-nbs/textual_inversion_inference.ipynb b/notebooks_py/diffusion_/diffusion-nbs/textual_inversion_inference.ipynb deleted file mode 100644 index 90551543b..000000000 --- a/notebooks_py/diffusion_/diffusion-nbs/textual_inversion_inference.ipynb +++ /dev/null @@ -1,614 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 1, - "metadata": { - "cellView": "form", - "id": "1_h0kO-VnQog" - }, - "outputs": [], - "source": [ - "import torch\n", - "from diffusers import StableDiffusionPipeline\n", - "from PIL import Image\n", - "from pathlib import Path" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": { - "cellView": "form", - "id": "1_h0kO-VnQog" - }, - "outputs": [], - "source": [ - "torch.manual_seed(1)\n", - "\n", - "def image_grid(imgs, rows=1, cols=None):\n", - " if cols is None: cols = math.ceil(len(imgs)/rows)\n", - " w, h = imgs[0].size\n", - " grid = Image.new('RGB', size=(cols*w, rows*h))\n", - " for i, img in enumerate(imgs): grid.paste(img, box=(i%cols*w, i//cols*h))\n", - " return grid" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": { - "cellView": "form", - "id": "If5Jswe526QP" - }, - "outputs": [], - "source": [ - "model_nm = \"CompVis/stable-diffusion-v1-4\"\n", - "\n", - "output_dir=\"sd-concept-output\"\n", - "pipe = StableDiffusionPipeline.from_pretrained(output_dir, torch_dtype=torch.float16).to(\"cuda\")" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [], - "source": [ - "num_samples,num_rows = 2,2" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "prompt = \"a hyper-realistic high definition render of sitting on a pink rug, 4k\"\n", - "all_images = [] \n", - "for _ in range(num_rows):\n", - " images = pipe([prompt] * num_samples, num_inference_steps=100, guidance_scale=15).images\n", - " all_images.extend(images)" - ] - }, - { - "cell_type": "code", - "execution_count": 35, - "metadata": {}, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - " 0%| | 0/101 [00:00" - ] - }, - "execution_count": 35, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "image_grid(all_images, num_samples, num_rows)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "accelerator": "GPU", - "colab": { - "collapsed_sections": [ - "D633UIuGgs6M" - ], - "machine_shape": "hm", - "provenance": [] - }, - "gpuClass": "standard", - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.9.10" - }, - "toc": { - "base_numbering": 1, - "nav_menu": {}, - "number_sections": false, - "sideBar": true, - "skip_h1_title": false, - "title_cell": "Table of Contents", - "title_sidebar": "Contents", - "toc_cell": false, - "toc_position": {}, - "toc_section_display": true, - "toc_window_display": false - }, - "vscode": { - "interpreter": { - "hash": "76721e0cd9246c299eb22246d1f3c601ec1aef6bd84d45d2547549094e7b6fb7" - } - }, - "widgets": { - "application/vnd.jupyter.widget-state+json": { - "1e8761aac9ac45afabf064c02bb29980": { - "model_module": "@jupyter-widgets/base", - "model_module_version": "1.2.0", - "model_name": "LayoutModel", - "state": { - "_model_module": "@jupyter-widgets/base", - "_model_module_version": "1.2.0", - "_model_name": "LayoutModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/base", - "_view_module_version": "1.2.0", - "_view_name": "LayoutView", - "align_content": null, - "align_items": "center", - "align_self": null, - "border": null, - "bottom": null, - "display": "flex", - "flex": null, - "flex_flow": "column", - "grid_area": null, - "grid_auto_columns": null, - "grid_auto_flow": null, - "grid_auto_rows": null, - "grid_column": null, - "grid_gap": null, - "grid_row": null, - "grid_template_areas": null, - "grid_template_columns": null, - "grid_template_rows": null, - "height": null, - "justify_content": null, - "justify_items": null, - "left": null, - "margin": null, - "max_height": null, - "max_width": null, - "min_height": null, - "min_width": null, - "object_fit": null, - "object_position": null, - "order": null, - "overflow": null, - "overflow_x": null, - "overflow_y": null, - "padding": null, - "right": null, - "top": null, - "visibility": null, - "width": "50%" - } - }, - "20cd48c875bb400487a54ee073f3a249": { - "model_module": "@jupyter-widgets/controls", - "model_module_version": "1.5.0", - "model_name": "DescriptionStyleModel", - "state": { - "_model_module": "@jupyter-widgets/controls", - "_model_module_version": "1.5.0", - "_model_name": "DescriptionStyleModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/base", - "_view_module_version": "1.2.0", - "_view_name": "StyleView", - "description_width": "" - } - }, - "3aff4ee1482a45d78b5df50b65995e0f": { - "model_module": "@jupyter-widgets/base", - "model_module_version": "1.2.0", - "model_name": "LayoutModel", - "state": { - "_model_module": "@jupyter-widgets/base", - "_model_module_version": "1.2.0", - "_model_name": "LayoutModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/base", - "_view_module_version": "1.2.0", - "_view_name": "LayoutView", - "align_content": null, - "align_items": null, - "align_self": null, - "border": null, - "bottom": null, - "display": null, - "flex": null, - "flex_flow": null, - "grid_area": null, - "grid_auto_columns": null, - "grid_auto_flow": null, - "grid_auto_rows": null, - "grid_column": null, - "grid_gap": null, - "grid_row": null, - "grid_template_areas": null, - "grid_template_columns": null, - "grid_template_rows": null, - "height": null, - "justify_content": null, - "justify_items": null, - "left": null, - "margin": null, - "max_height": null, - "max_width": null, - "min_height": null, - "min_width": null, - "object_fit": null, - "object_position": null, - "order": null, - "overflow": null, - "overflow_x": null, - "overflow_y": null, - "padding": null, - "right": null, - "top": null, - "visibility": null, - "width": null - } - }, - "4ac26360191543489a24cdb6799f3233": { - "model_module": "@jupyter-widgets/base", - "model_module_version": "1.2.0", - "model_name": "LayoutModel", - "state": { - "_model_module": "@jupyter-widgets/base", - "_model_module_version": "1.2.0", - "_model_name": "LayoutModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/base", - "_view_module_version": "1.2.0", - "_view_name": "LayoutView", - "align_content": null, - "align_items": null, - "align_self": null, - "border": null, - "bottom": null, - "display": null, - "flex": null, - "flex_flow": null, - "grid_area": null, - "grid_auto_columns": null, - "grid_auto_flow": null, - "grid_auto_rows": null, - "grid_column": null, - "grid_gap": null, - "grid_row": null, - "grid_template_areas": null, - "grid_template_columns": null, - "grid_template_rows": null, - "height": null, - "justify_content": null, - "justify_items": null, - "left": null, - "margin": null, - "max_height": null, - "max_width": null, - "min_height": null, - "min_width": null, - "object_fit": null, - "object_position": null, - "order": null, - "overflow": null, - "overflow_x": null, - "overflow_y": null, - "padding": null, - "right": null, - "top": null, - "visibility": null, - "width": null - } - }, - "55736163b4f64c258b9281692326944f": { - "model_module": "@jupyter-widgets/controls", - "model_module_version": "1.5.0", - "model_name": "PasswordModel", - "state": { - "_dom_classes": [], - "_model_module": "@jupyter-widgets/controls", - "_model_module_version": "1.5.0", - "_model_name": "PasswordModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/controls", - "_view_module_version": "1.5.0", - "_view_name": "PasswordView", - "continuous_update": true, - "description": "Token:", - "description_tooltip": null, - "disabled": false, - "layout": "IPY_MODEL_c19245ddcff145cb880c89444549ba1d", - "placeholder": "​", - "style": "IPY_MODEL_a86a1dd9a4ac4c99b2ba03e20a127ec2", - "value": "" - } - }, - "6bbfa9beb69147238bdff430ea30340f": { - "model_module": "@jupyter-widgets/controls", - "model_module_version": "1.5.0", - "model_name": "DescriptionStyleModel", - "state": { - "_model_module": "@jupyter-widgets/controls", - "_model_module_version": "1.5.0", - "_model_name": "DescriptionStyleModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/base", - "_view_module_version": "1.2.0", - "_view_name": "StyleView", - "description_width": "" - } - }, - "9f69572d50a14c159540b6c2fe781606": { - "model_module": "@jupyter-widgets/base", - "model_module_version": "1.2.0", - "model_name": "LayoutModel", - "state": { - "_model_module": "@jupyter-widgets/base", - "_model_module_version": "1.2.0", - "_model_name": "LayoutModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/base", - "_view_module_version": "1.2.0", - "_view_name": "LayoutView", - "align_content": null, - "align_items": null, - "align_self": null, - "border": null, - "bottom": null, - "display": null, - "flex": null, - "flex_flow": null, - "grid_area": null, - "grid_auto_columns": null, - "grid_auto_flow": null, - "grid_auto_rows": null, - "grid_column": null, - "grid_gap": null, - "grid_row": null, - "grid_template_areas": null, - "grid_template_columns": null, - "grid_template_rows": null, - "height": null, - "justify_content": null, - "justify_items": null, - "left": null, - "margin": null, - "max_height": null, - "max_width": null, - "min_height": null, - "min_width": null, - "object_fit": null, - "object_position": null, - "order": null, - "overflow": null, - "overflow_x": null, - "overflow_y": null, - "padding": null, - "right": null, - "top": null, - "visibility": null, - "width": null - } - }, - "a86a1dd9a4ac4c99b2ba03e20a127ec2": { - "model_module": "@jupyter-widgets/controls", - "model_module_version": "1.5.0", - "model_name": "DescriptionStyleModel", - "state": { - "_model_module": "@jupyter-widgets/controls", - "_model_module_version": "1.5.0", - "_model_name": "DescriptionStyleModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/base", - "_view_module_version": "1.2.0", - "_view_name": "StyleView", - "description_width": "" - } - }, - "aa40c19a05ef4b96aaf6696f06895617": { - "model_module": "@jupyter-widgets/controls", - "model_module_version": "1.5.0", - "model_name": "HTMLModel", - "state": { - "_dom_classes": [], - "_model_module": "@jupyter-widgets/controls", - "_model_module_version": "1.5.0", - "_model_name": "HTMLModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/controls", - "_view_module_version": "1.5.0", - "_view_name": "HTMLView", - "description": "", - "description_tooltip": null, - "layout": "IPY_MODEL_4ac26360191543489a24cdb6799f3233", - "placeholder": "​", - "style": "IPY_MODEL_20cd48c875bb400487a54ee073f3a249", - "value": "

Copy a token from your Hugging Face\ntokens page and paste it below.
Immediately click login after copying\nyour token or it might be stored in plain text in this notebook file.
" - } - }, - "b771d3a71dbd4e22ac63165f57fc3c95": { - "model_module": "@jupyter-widgets/controls", - "model_module_version": "1.5.0", - "model_name": "VBoxModel", - "state": { - "_dom_classes": [], - "_model_module": "@jupyter-widgets/controls", - "_model_module_version": "1.5.0", - "_model_name": "VBoxModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/controls", - "_view_module_version": "1.5.0", - "_view_name": "VBoxView", - "box_style": "", - "children": [ - "IPY_MODEL_aa40c19a05ef4b96aaf6696f06895617", - "IPY_MODEL_55736163b4f64c258b9281692326944f", - "IPY_MODEL_e5a2c75d3fb44a31a34e9edbf9c20e4c", - "IPY_MODEL_fbaef77bd56d4bd790015ae92f17685f" - ], - "layout": "IPY_MODEL_1e8761aac9ac45afabf064c02bb29980" - } - }, - "c19245ddcff145cb880c89444549ba1d": { - "model_module": "@jupyter-widgets/base", - "model_module_version": "1.2.0", - "model_name": "LayoutModel", - "state": { - "_model_module": "@jupyter-widgets/base", - "_model_module_version": "1.2.0", - "_model_name": "LayoutModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/base", - "_view_module_version": "1.2.0", - "_view_name": "LayoutView", - "align_content": null, - "align_items": null, - "align_self": null, - "border": null, - "bottom": null, - "display": null, - "flex": null, - "flex_flow": null, - "grid_area": null, - "grid_auto_columns": null, - "grid_auto_flow": null, - "grid_auto_rows": null, - "grid_column": null, - "grid_gap": null, - "grid_row": null, - "grid_template_areas": null, - "grid_template_columns": null, - "grid_template_rows": null, - "height": null, - "justify_content": null, - "justify_items": null, - "left": null, - "margin": null, - "max_height": null, - "max_width": null, - "min_height": null, - "min_width": null, - "object_fit": null, - "object_position": null, - "order": null, - "overflow": null, - "overflow_x": null, - "overflow_y": null, - "padding": null, - "right": null, - "top": null, - "visibility": null, - "width": null - } - }, - "e5a2c75d3fb44a31a34e9edbf9c20e4c": { - "model_module": "@jupyter-widgets/controls", - "model_module_version": "1.5.0", - "model_name": "ButtonModel", - "state": { - "_dom_classes": [], - "_model_module": "@jupyter-widgets/controls", - "_model_module_version": "1.5.0", - "_model_name": "ButtonModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/controls", - "_view_module_version": "1.5.0", - "_view_name": "ButtonView", - "button_style": "", - "description": "Login", - "disabled": false, - "icon": "", - "layout": "IPY_MODEL_9f69572d50a14c159540b6c2fe781606", - "style": "IPY_MODEL_fff76a2ab12748f6976203efca969f92", - "tooltip": "" - } - }, - "fbaef77bd56d4bd790015ae92f17685f": { - "model_module": "@jupyter-widgets/controls", - "model_module_version": "1.5.0", - "model_name": "HTMLModel", - "state": { - "_dom_classes": [], - "_model_module": "@jupyter-widgets/controls", - "_model_module_version": "1.5.0", - "_model_name": "HTMLModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/controls", - "_view_module_version": "1.5.0", - "_view_name": "HTMLView", - "description": "", - "description_tooltip": null, - "layout": "IPY_MODEL_3aff4ee1482a45d78b5df50b65995e0f", - "placeholder": "​", - "style": "IPY_MODEL_6bbfa9beb69147238bdff430ea30340f", - "value": "\nPro Tip: If you don't already have one, you can create a dedicated\n'notebooks' token with 'write' access, that you can then easily reuse for all\nnotebooks. " - } - }, - "fff76a2ab12748f6976203efca969f92": { - "model_module": "@jupyter-widgets/controls", - "model_module_version": "1.5.0", - "model_name": "ButtonStyleModel", - "state": { - "_model_module": "@jupyter-widgets/controls", - "_model_module_version": "1.5.0", - "_model_name": "ButtonStyleModel", - "_view_count": null, - "_view_module": "@jupyter-widgets/base", - "_view_module_version": "1.2.0", - "_view_name": "StyleView", - "button_color": null, - "font_weight": "" - } - } - } - } - }, - "nbformat": 4, - "nbformat_minor": 1 -} diff --git a/notebooks_py/diffusion_/diffusion-nbs/textual_inversion_training.py b/notebooks_py/diffusion_/diffusion-nbs/textual_inversion_training.py deleted file mode 100644 index 5bd710ab0..000000000 --- a/notebooks_py/diffusion_/diffusion-nbs/textual_inversion_training.py +++ /dev/null @@ -1,172 +0,0 @@ -import argparse, itertools, math, os, random, PIL -import numpy as np, torch, torch.nn.functional as F, torch.utils.checkpoint -from torch.utils.data import Dataset, DataLoader - -from accelerate import Accelerator -from accelerate.logging import get_logger -from accelerate.utils import set_seed -from diffusers import AutoencoderKL, DDPMScheduler, PNDMScheduler, StableDiffusionPipeline, UNet2DConditionModel -from diffusers.hub_utils import init_git_repo, push_to_hub -from diffusers.optimization import get_scheduler -from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker -from PIL import Image -from PIL.Image import Resampling -from torchvision import transforms -from tqdm.auto import tqdm - -from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer -import fastcore.all as fc -from huggingface_hub import notebook_login -from pathlib import Path - -import torchvision.transforms.functional as tf -import accelerate - -torch.manual_seed(1) -if not (Path.home()/'.cache/huggingface'/'token').exists(): notebook_login() - -model_nm = "CompVis/stable-diffusion-v1-4" - -path = Path.home()/'Downloads/photos/' -paths = list(path.iterdir()) -images = [Image.open(p).resize((512, 512), resample=Resampling.BICUBIC).convert("RGB") for p in paths] - -what_to_teach = "object" -placeholder_token = "" -initializer_token = "teddy" - -templates = [ - "a photo of a {}", - "a rendering of a {}", - "a cropped photo of the {}", - "the photo of a {}", - "a photo of a clean {}", - "a photo of a dirty {}", - "a dark photo of the {}", - "a photo of my {}", - "a photo of the cool {}", - "a close-up photo of a {}", - "a bright photo of the {}", - "a cropped photo of a {}", - "a photo of the {}", - "a good photo of the {}", - "a photo of one {}", - "a close-up photo of the {}", - "a rendition of the {}", - "a photo of the clean {}", - "a rendition of a {}", - "a photo of a nice {}", - "a good photo of a {}", - "a photo of the nice {}", - "a photo of the small {}", - "a photo of the weird {}", - "a photo of the large {}", - "a photo of a cool {}", - "a photo of a small {}", -] - -class TextualInversionDataset: - def __init__(self, tokenizer, images, learnable_property="object", size=512, - repeats=100, interpolation=Resampling.BICUBIC, flip_p=0.5, set="train", placeholder_token="*"): - fc.store_attr() - self.num_images = len(images) - if set == "train": self._length = self.num_images * repeats - self.templates = style_templates if learnable_property == "style" else templates - self.flip_transform = transforms.RandomHorizontalFlip(p=self.flip_p) - - def __len__(self): return self.num_images - - def __getitem__(self, i): - image = tf.to_tensor(self.images[i%self.num_images])*2-1 - text = random.choice(self.templates).format(self.placeholder_token) - ids=self.tokenizer(text, padding="max_length", truncation=True, max_length=self.tokenizer.model_max_length, return_tensors="pt") - return dict(input_ids=ids.input_ids[0], pixel_values=image) - -tokenizer = CLIPTokenizer.from_pretrained(model_nm, subfolder="tokenizer") -num_added_tokens = tokenizer.add_tokens(placeholder_token) -token_ids = tokenizer.encode(initializer_token, add_special_tokens=False) -initializer_token_id = token_ids[0] -placeholder_token_id = tokenizer.convert_tokens_to_ids(placeholder_token) - -text_encoder = CLIPTextModel.from_pretrained(model_nm, subfolder="text_encoder") -vae = AutoencoderKL.from_pretrained(model_nm, subfolder="vae") -unet = UNet2DConditionModel.from_pretrained(model_nm, subfolder="unet") -text_encoder.resize_token_embeddings(len(tokenizer)) -token_embeds = text_encoder.get_input_embeddings().weight.data -token_embeds[placeholder_token_id] = token_embeds[initializer_token_id] - -# Freeze all parameters except for the token embeddings in text encoder -tm = text_encoder.text_model -for o in (vae, unet, tm.encoder, tm.final_layer_norm, tm.embeddings.position_embedding): - for p in o.parameters(): p.requires_grad = False - -train_dataset = TextualInversionDataset( - images=images, tokenizer=tokenizer, size=512, placeholder_token=placeholder_token, - repeats=100, learnable_property=what_to_teach, set="train") - -def create_dataloader(bs=1): return DataLoader(train_dataset, batch_size=bs, shuffle=True) - -noise_scheduler = DDPMScheduler( - beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, tensor_format="pt") - -def training_function(text_encoder, vae, unet, train_batch_size, gradient_accumulation_steps, - lr, max_train_steps, scale_lr, output_dir): - accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps, mixed_precision='fp16') - train_dataloader = create_dataloader(train_batch_size) - if scale_lr: lr = (lr * gradient_accumulation_steps * train_batch_size * accelerator.num_processes) - optimizer = torch.optim.AdamW(text_encoder.get_input_embeddings().parameters(), lr=lr) - text_encoder, optimizer, train_dataloader = accelerator.prepare(text_encoder, optimizer, train_dataloader) - vae.to(accelerator.device).eval() - unet.to(accelerator.device).eval() - - num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps) - num_train_epochs = math.ceil(max_train_steps / num_update_steps_per_epoch) - total_batch_size = train_batch_size * accelerator.num_processes * gradient_accumulation_steps - progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process) - progress_bar.set_description("Steps") - global_step = 0 - - for epoch in range(num_train_epochs): - text_encoder.train() - for step, batch in enumerate(train_dataloader): - with accelerator.accumulate(text_encoder): - latents = vae.encode(batch["pixel_values"]).latent_dist.sample().detach() * 0.18215 - noise = torch.randn(latents.shape).to(latents.device) - bsz = latents.shape[0] - timesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (bsz,), device=latents.device).long() - - noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps) - encoder_hidden_states = text_encoder(batch["input_ids"])[0] - noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample - loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean() - accelerator.backward(loss) - - # We only want to optimize the concept embeddings - grads = text_encoder.get_input_embeddings().weight.grad - index_grads_to_zero = torch.arange(len(tokenizer)) != placeholder_token_id - grads.data[index_grads_to_zero, :] = grads.data[index_grads_to_zero, :].fill_(0) - optimizer.step() - optimizer.zero_grad() - - if accelerator.sync_gradients: - progress_bar.update(1) - global_step += 1 - - progress_bar.set_postfix(loss=loss.detach().item()) - if global_step >= max_train_steps: break - - pipeline = StableDiffusionPipeline( - text_encoder=accelerator.unwrap_model(text_encoder), - vae=vae, unet=unet, tokenizer=tokenizer, - scheduler=PNDMScheduler( - beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", skip_prk_steps=True), - safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker"), - feature_extractor=CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")) - pipeline.save_pretrained(output_dir) - learned_embeds = accelerator.unwrap_model(text_encoder).get_input_embeddings().weight[placeholder_token_id] - learned_embeds_dict = {placeholder_token: learned_embeds.detach().cpu()} - torch.save(learned_embeds_dict, os.path.join(output_dir, "learned_embeds.bin")) - -torch.manual_seed(42) -training_function(text_encoder, vae, unet, train_batch_size=1, gradient_accumulation_steps=4, lr=5e-04, - max_train_steps=3000, scale_lr=True, output_dir="sd-concept-output") diff --git a/notebooks_py/diffusion_/diffusion_policy/.gitignore b/notebooks_py/diffusion_/diffusion_policy/.gitignore deleted file mode 100644 index 23dbd4ec3..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/.gitignore +++ /dev/null @@ -1,140 +0,0 @@ -bin -logs -wandb -outputs -data -data_local -.vscode -_wandb - -**/.DS_Store - -fuse.cfg - -*.ai - -# Generation results -results/ - -ray/auth.json - -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class - -# C extensions -*.so - -# Distribution / packaging -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -pip-wheel-metadata/ -share/python-wheels/ -*.egg-info/ -.installed.cfg -*.egg -MANIFEST - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.nox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -*.py,cover -.hypothesis/ -.pytest_cache/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py -db.sqlite3 -db.sqlite3-journal - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -docs/_build/ - -# PyBuilder -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# IPython -profile_default/ -ipython_config.py - -# pyenv -.python-version - -# pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. -#Pipfile.lock - -# PEP 582; used by e.g. github.com/David-OConnor/pyflow -__pypackages__/ - -# Celery stuff -celerybeat-schedule -celerybeat.pid - -# SageMath parsed files -*.sage.py - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy -.mypy_cache/ -.dmypy.json -dmypy.json - -# Pyre type checker -.pyre/ diff --git a/notebooks_py/diffusion_/diffusion_policy/LICENSE b/notebooks_py/diffusion_/diffusion_policy/LICENSE deleted file mode 100644 index a2eb6cb62..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2023 Columbia Artificial Intelligence and Robotics Lab - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/notebooks_py/diffusion_/diffusion_policy/README.md b/notebooks_py/diffusion_/diffusion_policy/README.md deleted file mode 100644 index 7741de13d..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/README.md +++ /dev/null @@ -1,441 +0,0 @@ -# Diffusion Policy - -rename or just leave notebooks of ocol stuff - into ifnra?!?!? - -vow of _"sielne" - -[[Project page]](https://diffusion-policy.cs.columbia.edu/) -[[Paper]](https://diffusion-policy.cs.columbia.edu/#paper) -[[Data]](https://diffusion-policy.cs.columbia.edu/data/) -[[Colab (state)]](https://colab.research.google.com/drive/1gxdkgRVfM55zihY9TFLja97cSVZOZq2B?usp=sharing) -[[Colab (vision)]](https://colab.research.google.com/drive/18GIHeOQ5DyjMN8iIRZL2EKZ0745NLIpg?usp=sharing) - - -[Cheng Chi](http://cheng-chi.github.io/)1, -[Siyuan Feng](https://www.cs.cmu.edu/~sfeng/)2, -[Yilun Du](https://yilundu.github.io/)3, -[Zhenjia Xu](https://www.zhenjiaxu.com/)1, -[Eric Cousineau](https://www.eacousineau.com/)2, -[Benjamin Burchfiel](http://www.benburchfiel.com/)2, -[Shuran Song](https://www.cs.columbia.edu/~shurans/)1 - -1Columbia University, -2Toyota Research Institute, -3MIT - -drawing -drawing - -## 🛝 Try it out! -Our self-contained Google Colab notebooks is the easiest way to play with Diffusion Policy. We provide separate notebooks for [state-based environment](https://colab.research.google.com/drive/1gxdkgRVfM55zihY9TFLja97cSVZOZq2B?usp=sharing) and [vision-based environment](https://colab.research.google.com/drive/18GIHeOQ5DyjMN8iIRZL2EKZ0745NLIpg?usp=sharing). - -## 🧾 Checkout our experiment logs! -For each experiment used to generate Table I,II and IV in the [paper](https://diffusion-policy.cs.columbia.edu/#paper), we provide: -1. A `config.yaml` that contains all parameters needed to reproduce the experiment. -2. Detailed training/eval `logs.json.txt` for every training step. -3. Checkpoints for the best `epoch=*-test_mean_score=*.ckpt` and last `latest.ckpt` epoch of each run. - -Experiment logs are hosted on our website as nested directories in format: -`https://diffusion-policy.cs.columbia.edu/data/experiments////` - -Within each experiment directory you may find: -``` -. -├── config.yaml -├── metrics -│   └── logs.json.txt -├── train_0 -│   ├── checkpoints -│   │   ├── epoch=0300-test_mean_score=1.000.ckpt -│   │   └── latest.ckpt -│   └── logs.json.txt -├── train_1 -│   ├── checkpoints -│   │   ├── epoch=0250-test_mean_score=1.000.ckpt -│   │   └── latest.ckpt -│   └── logs.json.txt -└── train_2 - ├── checkpoints - │   ├── epoch=0250-test_mean_score=1.000.ckpt - │   └── latest.ckpt - └── logs.json.txt -``` -The `metrics/logs.json.txt` file aggregates evaluation metrics from all 3 training runs every 50 epochs using `multirun_metrics.py`. The numbers reported in the paper correspond to `max` and `k_min_train_loss` aggregation keys. - -To download all files in a subdirectory, use: - -```console -$ wget --recursive --no-parent --no-host-directories --relative --reject="index.html*" https://diffusion-policy.cs.columbia.edu/data/experiments/low_dim/square_ph/diffusion_policy_cnn/ -``` - -## 🛠️ Installation -### 🖥️ Simulation -To reproduce our simulation benchmark results, install our conda environment on a Linux machine with Nvidia GPU. On Ubuntu 20.04 you need to install the following apt packages for mujoco: -```console -$ sudo apt install -y libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf -``` - -We recommend [Mambaforge](https://github.com/conda-forge/miniforge#mambaforge) instead of the standard anaconda distribution for faster installation: -```console -$ mamba env create -f conda_environment.yaml -``` - -but you can use conda as well: -```console -$ conda env create -f conda_environment.yaml -``` - -The `conda_environment_macos.yaml` file is only for development on MacOS and does not have full support for benchmarks. - -### 🦾 Real Robot -Hardware (for Push-T): -* 1x [UR5-CB3](https://www.universal-robots.com/cb3) or [UR5e](https://www.universal-robots.com/products/ur5-robot/) ([RTDE Interface](https://www.universal-robots.com/articles/ur/interface-communication/real-time-data-exchange-rtde-guide/) is required) -* 2x [RealSense D415](https://www.intelrealsense.com/depth-camera-d415/) -* 1x [3Dconnexion SpaceMouse](https://3dconnexion.com/us/product/spacemouse-wireless/) (for teleop) -* 1x [Millibar Robotics Manual Tool Changer](https://www.millibar.com/manual-tool-changer/) (only need robot side) -* 1x 3D printed [End effector](https://cad.onshape.com/documents/a818888644a15afa6cc68ee5/w/2885b48b018cda84f425beca/e/3e8771c2124cee024edd2fed?renderMode=0&uiState=63ffcba6631ca919895e64e5) -* 1x 3D printed [T-block](https://cad.onshape.com/documents/f1140134e38f6ed6902648d5/w/a78cf81827600e4ff4058d03/e/f35f57fb7589f72e05c76caf?renderMode=0&uiState=63ffcbc9af4a881b344898ee) -* USB-C cables and screws for RealSense - -Software: -* Ubuntu 20.04.3 (tested) -* Mujoco dependencies: -`sudo apt install libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf` -* [RealSense SDK](https://github.com/IntelRealSense/librealsense/blob/master/doc/distribution_linux.md) -* Spacemouse dependencies: -`sudo apt install libspnav-dev spacenavd; sudo systemctl start spacenavd` -* Conda environment `mamba env create -f conda_environment_real.yaml` - -## 🖥️ Reproducing Simulation Benchmark Results -### Download Training Data -Under the repo root, create data subdirectory: -```console -[diffusion_policy]$ mkdir data && cd data -``` - -Download the corresponding zip file from [https://diffusion-policy.cs.columbia.edu/data/training/](https://diffusion-policy.cs.columbia.edu/data/training/) -```console -[data]$ wget https://diffusion-policy.cs.columbia.edu/data/training/pusht.zip -``` - -Extract training data: -```console -[data]$ unzip pusht.zip && rm -f pusht.zip && cd .. -``` - -Grab config file for the corresponding experiment: -```console -[diffusion_policy]$ wget -O image_pusht_diffusion_policy_cnn.yaml https://diffusion-policy.cs.columbia.edu/data/experiments/image/pusht/diffusion_policy_cnn/config.yaml -``` - -### Running for a single seed -Activate conda environment and login to [wandb](https://wandb.ai) (if you haven't already). -```console -[diffusion_policy]$ conda activate robodiff -(robodiff)[diffusion_policy]$ wandb login -``` - -Launch training with seed 42 on GPU 0. -```console -(robodiff)[diffusion_policy]$ python train.py --config-dir=. --config-name=image_pusht_diffusion_policy_cnn.yaml training.seed=42 training.device=cuda:0 hydra.run.dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}' -``` - -This will create a directory in format `data/outputs/yyyy.mm.dd/hh.mm.ss__` where configs, logs and checkpoints are written to. The policy will be evaluated every 50 epochs with the success rate logged as `test/mean_score` on wandb, as well as videos for some rollouts. -```console -(robodiff)[diffusion_policy]$ tree data/outputs/2023.03.01/20.02.03_train_diffusion_unet_hybrid_pusht_image -I wandb -data/outputs/2023.03.01/20.02.03_train_diffusion_unet_hybrid_pusht_image -├── checkpoints -│ ├── epoch=0000-test_mean_score=0.134.ckpt -│ └── latest.ckpt -├── .hydra -│ ├── config.yaml -│ ├── hydra.yaml -│ └── overrides.yaml -├── logs.json.txt -├── media -│ ├── 2k5u6wli.mp4 -│ ├── 2kvovxms.mp4 -│ ├── 2pxd9f6b.mp4 -│ ├── 2q5gjt5f.mp4 -│ ├── 2sawbf6m.mp4 -│ └── 538ubl79.mp4 -└── train.log - -3 directories, 13 files -``` - -### Running for multiple seeds -Launch local ray cluster. For large scale experiments, you might want to setup an [AWS cluster with autoscaling](https://docs.ray.io/en/master/cluster/vms/user-guides/launching-clusters/aws.html). All other commands remain the same. -```console -(robodiff)[diffusion_policy]$ export CUDA_VISIBLE_DEVICES=0,1,2 # select GPUs to be managed by the ray cluster -(robodiff)[diffusion_policy]$ ray start --head --num-gpus=3 -``` - -Launch a ray client which will start 3 training workers (3 seeds) and 1 metrics monitor worker. -```console -(robodiff)[diffusion_policy]$ python ray_train_multirun.py --config-dir=. --config-name=image_pusht_diffusion_policy_cnn.yaml --seeds=42,43,44 --monitor_key=test/mean_score -- multi_run.run_dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}' multi_run.wandb_name_base='${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}' -``` - -In addition to the wandb log written by each training worker individually, the metrics monitor worker will log to wandb project `diffusion_policy_metrics` for the metrics aggregated from all 3 training runs. Local config, logs and checkpoints will be written to `data/outputs/yyyy.mm.dd/hh.mm.ss__` in a directory structure identical to our [training logs](https://diffusion-policy.cs.columbia.edu/data/experiments/): -```console -(robodiff)[diffusion_policy]$ tree data/outputs/2023.03.01/22.13.58_train_diffusion_unet_hybrid_pusht_image -I 'wandb|media' -data/outputs/2023.03.01/22.13.58_train_diffusion_unet_hybrid_pusht_image -├── config.yaml -├── metrics -│ ├── logs.json.txt -│ ├── metrics.json -│ └── metrics.log -├── train_0 -│ ├── checkpoints -│ │ ├── epoch=0000-test_mean_score=0.174.ckpt -│ │ └── latest.ckpt -│ ├── logs.json.txt -│ └── train.log -├── train_1 -│ ├── checkpoints -│ │ ├── epoch=0000-test_mean_score=0.131.ckpt -│ │ └── latest.ckpt -│ ├── logs.json.txt -│ └── train.log -└── train_2 - ├── checkpoints - │ ├── epoch=0000-test_mean_score=0.105.ckpt - │ └── latest.ckpt - ├── logs.json.txt - └── train.log - -7 directories, 16 files -``` -### 🆕 Evaluate Pre-trained Checkpoints -Download a checkpoint from the published training log folders, such as [https://diffusion-policy.cs.columbia.edu/data/experiments/low_dim/pusht/diffusion_policy_cnn/train_0/checkpoints/epoch=0550-test_mean_score=0.969.ckpt](https://diffusion-policy.cs.columbia.edu/data/experiments/low_dim/pusht/diffusion_policy_cnn/train_0/checkpoints/epoch=0550-test_mean_score=0.969.ckpt). - -Run the evaluation script: -```console -(robodiff)[diffusion_policy]$ python eval.py --checkpoint data/0550-test_mean_score=0.969.ckpt --output_dir data/pusht_eval_output --device cuda:0 -``` - -This will generate the following directory structure: -```console -(robodiff)[diffusion_policy]$ tree data/pusht_eval_output -data/pusht_eval_output -├── eval_log.json -└── media - ├── 1fxtno84.mp4 - ├── 224l7jqd.mp4 - ├── 2fo4btlf.mp4 - ├── 2in4cn7a.mp4 - ├── 34b3o2qq.mp4 - └── 3p7jqn32.mp4 - -1 directory, 7 files -``` - -`eval_log.json` contains metrics that is logged to wandb during training: -```console -(robodiff)[diffusion_policy]$ cat data/pusht_eval_output/eval_log.json -{ - "test/mean_score": 0.9150393806777066, - "test/sim_max_reward_4300000": 1.0, - "test/sim_max_reward_4300001": 0.9872969750774386, -... - "train/sim_video_1": "data/pusht_eval_output//media/2fo4btlf.mp4" -} -``` - -## 🦾 Demo, Training and Eval on a Real Robot -Make sure your UR5 robot is running and accepting command from its network interface (emergency stop button within reach at all time), your RealSense cameras plugged in to your workstation (tested with `realsense-viewer`) and your SpaceMouse connected with the `spacenavd` daemon running (verify with `systemctl status spacenavd`). - -Start the demonstration collection script. Press "C" to start recording. Use SpaceMouse to move the robot. Press "S" to stop recording. -```console -(robodiff)[diffusion_policy]$ python demo_real_robot.py -o data/demo_pusht_real --robot_ip 192.168.0.204 -``` - -This should result in a demonstration dataset in `data/demo_pusht_real` with in the same structure as our example [real Push-T training dataset](https://diffusion-policy.cs.columbia.edu/data/training/pusht_real.zip). - -To train a Diffusion Policy, launch training with config: -```console -(robodiff)[diffusion_policy]$ python train.py --config-name=train_diffusion_unet_real_image_workspace task.dataset_path=data/demo_pusht_real -``` -Edit [`diffusion_policy/config/task/real_pusht_image.yaml`](./diffusion_policy/config/task/real_pusht_image.yaml) if your camera setup is different. - -Assuming the training has finished and you have a checkpoint at `data/outputs/blah/checkpoints/latest.ckpt`, launch the evaluation script with: -```console -python eval_real_robot.py -i data/outputs/blah/checkpoints/latest.ckpt -o data/eval_pusht_real --robot_ip 192.168.0.204 -``` -Press "C" to start evaluation (handing control over to the policy). Press "S" to stop the current episode. - -## 🗺️ Codebase Tutorial -This codebase is structured under the requirement that: -1. implementing `N` tasks and `M` methods will only require `O(N+M)` amount of code instead of `O(N*M)` -2. while retaining maximum flexibility. - -To achieve this requirement, we -1. maintained a simple unified interface between tasks and methods and -2. made the implementation of the tasks and the methods independent of each other. - -These design decisions come at the cost of code repetition between the tasks and the methods. However, we believe that the benefit of being able to add/modify task/methods without affecting the remainder and being able understand a task/method by reading the code linearly outweighs the cost of copying and pasting 😊. - -### The Split -On the task side, we have: -* `Dataset`: adapts a (third-party) dataset to the interface. -* `EnvRunner`: executes a `Policy` that accepts the interface and produce logs and metrics. -* `config/task/.yaml`: contains all information needed to construct `Dataset` and `EnvRunner`. -* (optional) `Env`: an `gym==0.21.0` compatible class that encapsulates the task environment. - -On the policy side, we have: -* `Policy`: implements inference according to the interface and part of the training process. -* `Workspace`: manages the life-cycle of training and evaluation (interleaved) of a method. -* `config/.yaml`: contains all information needed to construct `Policy` and `Workspace`. - -### The Interface -#### Low Dim -A [`LowdimPolicy`](./diffusion_policy/policy/base_lowdim_policy.py) takes observation dictionary: -- `"obs":` Tensor of shape `(B,To,Do)` - -and predicts action dictionary: -- `"action": ` Tensor of shape `(B,Ta,Da)` - -A [`LowdimDataset`](./diffusion_policy/dataset/base_dataset.py) returns a sample of dictionary: -- `"obs":` Tensor of shape `(To, Do)` -- `"action":` Tensor of shape `(Ta, Da)` - -Its `get_normalizer` method returns a [`LinearNormalizer`](./diffusion_policy/model/common/normalizer.py) with keys `"obs","action"`. - -The `Policy` handles normalization on GPU with its copy of the `LinearNormalizer`. The parameters of the `LinearNormalizer` is saved as part of the `Policy`'s weights checkpoint. - -#### Image -A [`ImagePolicy`](./diffusion_policy/policy/base_image_policy.py) takes observation dictionary: -- `"key0":` Tensor of shape `(B,To,*)` -- `"key1":` Tensor of shape e.g. `(B,To,H,W,3)` ([0,1] float32) - -and predicts action dictionary: -- `"action": ` Tensor of shape `(B,Ta,Da)` - -A [`ImageDataset`](./diffusion_policy/dataset/base_dataset.py) returns a sample of dictionary: -- `"obs":` Dict of - - `"key0":` Tensor of shape `(To, *)` - - `"key1":` Tensor fo shape `(To,H,W,3)` -- `"action":` Tensor of shape `(Ta, Da)` - -Its `get_normalizer` method returns a [`LinearNormalizer`](./diffusion_policy/model/common/normalizer.py) with keys `"key0","key1","action"`. - -#### Example -``` -To = 3 -Ta = 4 -T = 6 -|o|o|o| -| | |a|a|a|a| -|o|o| -| |a|a|a|a|a| -| | | | |a|a| -``` -Terminology in the paper: `varname` in the codebase -- Observation Horizon: `To|n_obs_steps` -- Action Horizon: `Ta|n_action_steps` -- Prediction Horizon: `T|horizon` - -The classical (e.g. MDP) single step observation/action formulation is included as a special case where `To=1` and `Ta=1`. - -## 🔩 Key Components -### `Workspace` -A `Workspace` object encapsulates all states and code needed to run an experiment. -* Inherits from [`BaseWorkspace`](./diffusion_policy/workspace/base_workspace.py). -* A single `OmegaConf` config object generated by `hydra` should contain all information needed to construct the Workspace object and running experiments. This config correspond to `config/.yaml` + hydra overrides. -* The `run` method contains the entire pipeline for the experiment. -* Checkpoints happen at the `Workspace` level. All training states implemented as object attributes are automatically saved by the `save_checkpoint` method. -* All other states for the experiment should be implemented as local variables in the `run` method. - -The entrypoint for training is `train.py` which uses `@hydra.main` decorator. Read [hydra](https://hydra.cc/)'s official documentation for command line arguments and config overrides. For example, the argument `task=` will replace the `task` subtree of the config with the content of `config/task/.yaml`, thereby selecting the task to run for this experiment. - -### `Dataset` -A `Dataset` object: -* Inherits from `torch.utils.data.Dataset`. -* Returns a sample conforming to [the interface](#the-interface) depending on whether the task has Low Dim or Image observations. -* Has a method `get_normalizer` that returns a `LinearNormalizer` conforming to [the interface](#the-interface). - -Normalization is a very common source of bugs during project development. It is sometimes helpful to print out the specific `scale` and `bias` vectors used for each key in the `LinearNormalizer`. - -Most of our implementations of `Dataset` uses a combination of [`ReplayBuffer`](#replaybuffer) and [`SequenceSampler`](./diffusion_policy/common/sampler.py) to generate samples. Correctly handling padding at the beginning and the end of each demonstration episode according to `To` and `Ta` is important for good performance. Please read our [`SequenceSampler`](./diffusion_policy/common/sampler.py) before implementing your own sampling method. - -### `Policy` -A `Policy` object: -* Inherits from `BaseLowdimPolicy` or `BaseImagePolicy`. -* Has a method `predict_action` that given observation dict, predicts actions conforming to [the interface](#the-interface). -* Has a method `set_normalizer` that takes in a `LinearNormalizer` and handles observation/action normalization internally in the policy. -* (optional) Might has a method `compute_loss` that takes in a batch and returns the loss to be optimized. -* (optional) Usually each `Policy` class correspond to a `Workspace` class due to the differences of training and evaluation process between methods. - -### `EnvRunner` -A `EnvRunner` object abstracts away the subtle differences between different task environments. -* Has a method `run` that takes a `Policy` object for evaluation, and returns a dict of logs and metrics. Each value should be compatible with `wandb.log`. - -To maximize evaluation speed, we usually vectorize environments using our modification of [`gym.vector.AsyncVectorEnv`](./diffusion_policy/gym_util/async_vector_env.py) which runs each individual environment in a separate process (workaround python GIL). - -⚠️ Since subprocesses are launched using `fork` on linux, you need to be specially careful for environments that creates its OpenGL context during initialization (e.g. robosuite) which, once inherited by the child process memory space, often causes obscure bugs like segmentation fault. As a workaround, you can provide a `dummy_env_fn` that constructs an environment without initializing OpenGL. - -### `ReplayBuffer` -The [`ReplayBuffer`](./diffusion_policy/common/replay_buffer.py) is a key data structure for storing a demonstration dataset both in-memory and on-disk with chunking and compression. It makes heavy use of the [`zarr`](https://zarr.readthedocs.io/en/stable/index.html) format but also has a `numpy` backend for lower access overhead. - -On disk, it can be stored as a nested directory (e.g. `data/pusht_cchi_v7_replay.zarr`) or a zip file (e.g. `data/robomimic/datasets/square/mh/image_abs.hdf5.zarr.zip`). - -Due to the relative small size of our datasets, it's often possible to store the entire image-based dataset in RAM with [`Jpeg2000` compression](./diffusion_policy/codecs/imagecodecs_numcodecs.py) which eliminates disk IO during training at the expense increasing of CPU workload. - -Example: -``` -data/pusht_cchi_v7_replay.zarr - ├── data - │ ├── action (25650, 2) float32 - │ ├── img (25650, 96, 96, 3) float32 - │ ├── keypoint (25650, 9, 2) float32 - │ ├── n_contacts (25650, 1) float32 - │ └── state (25650, 5) float32 - └── meta - └── episode_ends (206,) int64 -``` - -Each array in `data` stores one data field from all episodes concatenated along the first dimension (time). The `meta/episode_ends` array stores the end index for each episode along the fist dimension. - -### `SharedMemoryRingBuffer` -The [`SharedMemoryRingBuffer`](./diffusion_policy/shared_memory/shared_memory_ring_buffer.py) is a lock-free FILO data structure used extensively in our [real robot implementation](./diffusion_policy/real_world) to utilize multiple CPU cores while avoiding pickle serialization and locking overhead for `multiprocessing.Queue`. - -As an example, we would like to get the most recent `To` frames from 5 RealSense cameras. We launch 1 realsense SDK/pipeline per process using [`SingleRealsense`](./diffusion_policy/real_world/single_realsense.py), each continuously writes the captured images into a `SharedMemoryRingBuffer` shared with the main process. We can very quickly get the last `To` frames in the main process due to the FILO nature of `SharedMemoryRingBuffer`. - -We also implemented [`SharedMemoryQueue`](./diffusion_policy/shared_memory/shared_memory_queue.py) for FIFO, which is used in [`RTDEInterpolationController`](./diffusion_policy/real_world/rtde_interpolation_controller.py). - -### `RealEnv` -In contrast to [OpenAI Gym](https://gymnasium.farama.org/), our polices interact with the environment asynchronously. In [`RealEnv`](./diffusion_policy/real_world/real_env.py), the `step` method in `gym` is split into two methods: `get_obs` and `exec_actions`. - -The `get_obs` method returns the latest observation from `SharedMemoryRingBuffer` as well as their corresponding timestamps. This method can be call at any time during an evaluation episode. - -The `exec_actions` method accepts a sequence of actions and timestamps for the expected time of execution for each step. Once called, the actions are simply enqueued to the `RTDEInterpolationController`, and the method returns without blocking for execution. - -## 🩹 Adding a Task -Read and imitate: -* `diffusion_policy/dataset/pusht_image_dataset.py` -* `diffusion_policy/env_runner/pusht_image_runner.py` -* `diffusion_policy/config/task/pusht_image.yaml` - -Make sure that `shape_meta` correspond to input and output shapes for your task. Make sure `env_runner._target_` and `dataset._target_` point to the new classes you have added. When training, add `task=` to `train.py`'s arguments. - -## 🩹 Adding a Method -Read and imitate: -* `diffusion_policy/workspace/train_diffusion_unet_image_workspace.py` -* `diffusion_policy/policy/diffusion_unet_image_policy.py` -* `diffusion_policy/config/train_diffusion_unet_image_workspace.yaml` - -Make sure your workspace yaml's `_target_` points to the new workspace class you created. - -## 🏷️ License -This repository is released under the MIT license. See [LICENSE](LICENSE) for additional details. - -## 🙏 Acknowledgement -* Our [`ConditionalUnet1D`](./diffusion_policy/model/diffusion/conditional_unet1d.py) implementation is adapted from [Planning with Diffusion](https://github.com/jannerm/diffuser). -* Our [`TransformerForDiffusion`](./diffusion_policy/model/diffusion/transformer_for_diffusion.py) implementation is adapted from [MinGPT](https://github.com/karpathy/minGPT). -* The [BET](./diffusion_policy/model/bet) baseline is adapted from [its original repo](https://github.com/notmahi/bet). -* The [IBC](./diffusion_policy/policy/ibc_dfo_lowdim_policy.py) baseline is adapted from [Kevin Zakka's reimplementation](https://github.com/kevinzakka/ibc). -* The [Robomimic](https://github.com/ARISE-Initiative/robomimic) tasks and [`ObservationEncoder`](https://github.com/ARISE-Initiative/robomimic/blob/master/robomimic/models/obs_nets.py) are used extensively in this project. -* The [Push-T](./diffusion_policy/env/pusht) task is adapted from [IBC](https://github.com/google-research/ibc). -* The [Block Pushing](./diffusion_policy/env/block_pushing) task is adapted from [BET](https://github.com/notmahi/bet) and [IBC](https://github.com/google-research/ibc). -* The [Kitchen](./diffusion_policy/env/kitchen) task is adapted from [BET](https://github.com/notmahi/bet) and [Relay Policy Learning](https://github.com/google-research/relay-policy-learning). -* Our [shared_memory](./diffusion_policy/shared_memory) data structures are heavily inspired by [shared-ndarray2](https://gitlab.com/osu-nrsg/shared-ndarray2). diff --git a/notebooks_py/diffusion_/diffusion_policy/conda_environment.yaml b/notebooks_py/diffusion_/diffusion_policy/conda_environment.yaml deleted file mode 100644 index c540b5bdb..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/conda_environment.yaml +++ /dev/null @@ -1,65 +0,0 @@ -name: robodiff -channels: - - pytorch - - pytorch3d - - nvidia - - conda-forge -dependencies: - - python=3.9 - - pip=22.2.2 - - cudatoolkit=11.6 - - pytorch=1.12.1 - - torchvision=0.13.1 - - pytorch3d=0.7.0 - - numpy=1.23.3 - - numba==0.56.4 - - scipy==1.9.1 - - py-opencv=4.6.0 - - cffi=1.15.1 - - ipykernel=6.16 - - matplotlib=3.6.1 - - zarr=2.12.0 - - numcodecs=0.10.2 - - h5py=3.7.0 - - hydra-core=1.2.0 - - einops=0.4.1 - - tqdm=4.64.1 - - dill=0.3.5.1 - - scikit-video=1.1.11 - - scikit-image=0.19.3 - - gym=0.21.0 - - pymunk=6.2.1 - - wandb=0.13.3 - - threadpoolctl=3.1.0 - - shapely=1.8.4 - - cython=0.29.32 - - imageio=2.22.0 - - imageio-ffmpeg=0.4.7 - - termcolor=2.0.1 - - tensorboard=2.10.1 - - tensorboardx=2.5.1 - - psutil=5.9.2 - - click=8.0.4 - - boto3=1.24.96 - - accelerate=0.13.2 - - datasets=2.6.1 - - diffusers=0.11.1 - - av=10.0.0 - - cmake=3.24.3 - # trick to avoid cpu affinity issue described in https://github.com/pytorch/pytorch/issues/99625 - - llvm-openmp=14 - # trick to force reinstall imagecodecs via pip - - imagecodecs==2022.8.8 - - pip: - - ray[default,tune]==2.2.0 - # requires mujoco py dependencies libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf - - free-mujoco-py==2.1.6 - - pygame==2.1.2 - - pybullet-svl==3.1.6.4 - - robosuite @ https://github.com/cheng-chi/robosuite/archive/277ab9588ad7a4f4b55cf75508b44aa67ec171f0.tar.gz - - robomimic==0.2.0 - - pytorchvideo==0.1.5 - # pip package required for jpeg-xl - - imagecodecs==2022.9.26 - - r3m @ https://github.com/facebookresearch/r3m/archive/b2334e726887fa0206962d7984c69c5fb09cceab.tar.gz - - dm-control==1.0.9 diff --git a/notebooks_py/diffusion_/diffusion_policy/conda_environment_macos.yaml b/notebooks_py/diffusion_/diffusion_policy/conda_environment_macos.yaml deleted file mode 100644 index 0d2676724..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/conda_environment_macos.yaml +++ /dev/null @@ -1,55 +0,0 @@ -name: robodiff -channels: - - pytorch - - conda-forge -dependencies: - - python=3.9 - - pip=22.2.2 - - pytorch=1.12.1 - - torchvision=0.13.1 - - numpy=1.23.3 - - numba==0.56.4 - - scipy==1.9.1 - - py-opencv=4.6.0 - - cffi=1.15.1 - - ipykernel=6.16 - - matplotlib=3.6.1 - - zarr=2.12.0 - - numcodecs=0.10.2 - - h5py=3.7.0 - - hydra-core=1.2.0 - - einops=0.4.1 - - tqdm=4.64.1 - - dill=0.3.5.1 - - scikit-video=1.1.11 - - scikit-image=0.19.3 - - gym=0.21.0 - - pymunk=6.2.1 - - wandb=0.13.3 - - threadpoolctl=3.1.0 - - shapely=1.8.4 - - cython=0.29.32 - - imageio=2.22.0 - - imageio-ffmpeg=0.4.7 - - termcolor=2.0.1 - - tensorboard=2.10.1 - - tensorboardx=2.5.1 - - psutil=5.9.2 - - click=8.0.4 - - boto3=1.24.96 - - accelerate=0.13.2 - - datasets=2.6.1 - - diffusers=0.11.1 - - av=10.0.0 - - cmake=3.24.3 - # trick to force reinstall imagecodecs via pip - - imagecodecs==2022.8.8 - - pip: - - ray[default,tune]==2.2.0 - - pygame==2.1.2 - - robomimic==0.2.0 - - pytorchvideo==0.1.5 - - atomics==1.0.2 - # No support for jpeg-xl for MacOS - - imagecodecs==2022.9.26 - - r3m @ https://github.com/facebookresearch/r3m/archive/b2334e726887fa0206962d7984c69c5fb09cceab.tar.gz diff --git a/notebooks_py/diffusion_/diffusion_policy/conda_environment_real.yaml b/notebooks_py/diffusion_/diffusion_policy/conda_environment_real.yaml deleted file mode 100644 index 763636079..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/conda_environment_real.yaml +++ /dev/null @@ -1,73 +0,0 @@ -name: robodiff -channels: - - pytorch - - pytorch3d - - nvidia - - conda-forge -dependencies: - - python=3.9 - - pip=22.2.2 - - cudatoolkit=11.6 - - pytorch=1.12.1 - - torchvision=0.13.1 - - pytorch3d=0.7.0 - - numpy=1.23.3 - - numba==0.56.4 - - scipy==1.9.1 - - py-opencv=4.6.0 - - cffi=1.15.1 - - ipykernel=6.16 - - matplotlib=3.6.1 - - zarr=2.12.0 - - numcodecs=0.10.2 - - h5py=3.7.0 - - hydra-core=1.2.0 - - einops=0.4.1 - - tqdm=4.64.1 - - dill=0.3.5.1 - - scikit-video=1.1.11 - - scikit-image=0.19.3 - - gym=0.21.0 - - pymunk=6.2.1 - - wandb=0.13.3 - - threadpoolctl=3.1.0 - - shapely=1.8.4 - - cython=0.29.32 - - imageio=2.22.0 - - imageio-ffmpeg=0.4.7 - - termcolor=2.0.1 - - tensorboard=2.10.1 - - tensorboardx=2.5.1 - - psutil=5.9.2 - - click=8.0.4 - - boto3=1.24.96 - - accelerate=0.13.2 - - datasets=2.6.1 - - diffusers=0.11.1 - - av=10.0.0 - - cmake=3.24.3 - # trick to avoid cpu affinity issue described in https://github.com/pytorch/pytorch/issues/99625 - - llvm-openmp=14 - # trick to force reinstall imagecodecs via pip - - imagecodecs==2022.8.8 - - pip: - - ray[default,tune]==2.2.0 - # requires mujoco py dependencies libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf - - free-mujoco-py==2.1.6 - - pygame==2.1.2 - - pybullet-svl==3.1.6.4 - - robosuite @ https://github.com/cheng-chi/robosuite/archive/277ab9588ad7a4f4b55cf75508b44aa67ec171f0.tar.gz - - robomimic==0.2.0 - - pytorchvideo==0.1.5 - # requires librealsense https://github.com/IntelRealSense/librealsense/blob/master/doc/distribution_linux.md - - pyrealsense2==2.51.1.4348 - # requires apt install libspnav-dev spacenavd; systemctl start spacenavd - - spnav @ https://github.com/cheng-chi/spnav/archive/c1c938ebe3cc542db4685e0d13850ff1abfdb943.tar.gz - - ur-rtde==1.5.5 - - atomics==1.0.2 - # pip package required for jpeg-xl - - imagecodecs==2022.9.26 - - r3m @ https://github.com/facebookresearch/r3m/archive/b2334e726887fa0206962d7984c69c5fb09cceab.tar.gz - - dm-control==1.0.9 - - pynput==1.7.6 - \ No newline at end of file diff --git a/notebooks_py/diffusion_/diffusion_policy/demo_pusht.py b/notebooks_py/diffusion_/diffusion_policy/demo_pusht.py deleted file mode 100644 index 12fe1949e..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/demo_pusht.py +++ /dev/null @@ -1,124 +0,0 @@ -import sys -import os -print(sys.path) -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) -import numpy as np -import click -from diffusion_policy.common.replay_buffer import ReplayBuffer -from diffusion_policy.env.pusht.pusht_keypoints_env import PushTKeypointsEnv -import pygame - -@click.command() -@click.option('-o', '--output', required=True) -@click.option('-rs', '--render_size', default=96, type=int) -@click.option('-hz', '--control_hz', default=10, type=int) -def main(output, render_size, control_hz): - """ - Collect demonstration for the Push-T task. - - Usage: python demo_pusht.py -o data/pusht_demo.zarr - - This script is compatible with both Linux and MacOS. - Hover mouse close to the blue circle to start. - Push the T block into the green area. - The episode will automatically terminate if the task is succeeded. - Press "Q" to exit. - Press "R" to retry. - Hold "Space" to pause. - """ - - #create replay buffer in read-write mode - replay_buffer = ReplayBuffer.create_from_path(output, mode='a') - - # create PushT env with keypoints - kp_kwargs = PushTKeypointsEnv.genenerate_keypoint_manager_params() - env = PushTKeypointsEnv(render_size=render_size, render_action=False, **kp_kwargs) - agent = env.teleop_agent() - clock = pygame.time.Clock() - - # episode-level while loop - while True: - episode = list() - # record in seed order, starting with 0 - seed = replay_buffer.n_episodes - print(f'starting seed {seed}') - - # set seed for env - env.seed(seed) - - # reset env and get observations (including info and render for recording) - obs = env.reset() - info = env._get_info() - img = env.render(mode='human') - - # loop state - retry = False - pause = False - done = False - plan_idx = 0 - pygame.display.set_caption(f'plan_idx:{plan_idx}') - # step-level while loop - while not done: - # process keypress events - for event in pygame.event.get(): - if event.type == pygame.KEYDOWN: - if event.key == pygame.K_SPACE: - # hold Space to pause - plan_idx += 1 - pygame.display.set_caption(f'plan_idx:{plan_idx}') - pause = True - elif event.key == pygame.K_r: - # press "R" to retry - retry=True - elif event.key == pygame.K_q: - # press "Q" to exit - exit(0) - if event.type == pygame.KEYUP: - if event.key == pygame.K_SPACE: - pause = False - - # handle control flow - if retry: - break - if pause: - continue - - # get action from mouse - # None if mouse is not close to the agent - act = agent.act(obs) - if not act is None: - # teleop started - # state dim 2+3 - state = np.concatenate([info['pos_agent'], info['block_pose']]) - # discard unused information such as visibility mask and agent pos - # for compatibility - keypoint = obs.reshape(2,-1)[0].reshape(-1,2)[:9] - data = { - 'img': img, - 'state': np.float32(state), - 'keypoint': np.float32(keypoint), - 'action': np.float32(act), - 'n_contacts': np.float32([info['n_contacts']]) - } - episode.append(data) - - # step env and render - obs, reward, done, info = env.step(act) - img = env.render(mode='human') - - # regulate control frequency - clock.tick(control_hz) - if not retry: - # save episode buffer to replay buffer (on disk) - data_dict = dict() - for key in episode[0].keys(): - data_dict[key] = np.stack( - [x[key] for x in episode]) - replay_buffer.add_episode(data_dict, compressors='disk') - print(f'saved seed {seed}') - else: - print(f'retry seed {seed}') - - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/demo_real_robot.py b/notebooks_py/diffusion_/diffusion_policy/demo_real_robot.py deleted file mode 100644 index 846badefa..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/demo_real_robot.py +++ /dev/null @@ -1,160 +0,0 @@ -""" -Usage: -(robodiff)$ python demo_real_robot.py -o --robot_ip - -Robot movement: -Move your SpaceMouse to move the robot EEF (locked in xy plane). -Press SpaceMouse right button to unlock z axis. -Press SpaceMouse left button to enable rotation axes. - -Recording control: -Click the opencv window (make sure it's in focus). -Press "C" to start recording. -Press "S" to stop recording. -Press "Q" to exit program. -Press "Backspace" to delete the previously recorded episode. -""" - -# %% -import time -from multiprocessing.managers import SharedMemoryManager -import click -import cv2 -import numpy as np -import scipy.spatial.transform as st -from diffusion_policy.real_world.real_env import RealEnv -from diffusion_policy.real_world.spacemouse_shared_memory import Spacemouse -from diffusion_policy.common.precise_sleep import precise_wait -from diffusion_policy.real_world.keystroke_counter import ( - KeystrokeCounter, Key, KeyCode -) - -@click.command() -@click.option('--output', '-o', required=True, help="Directory to save demonstration dataset.") -@click.option('--robot_ip', '-ri', required=True, help="UR5's IP address e.g. 192.168.0.204") -@click.option('--vis_camera_idx', default=0, type=int, help="Which RealSense camera to visualize.") -@click.option('--init_joints', '-j', is_flag=True, default=False, help="Whether to initialize robot joint configuration in the beginning.") -@click.option('--frequency', '-f', default=10, type=float, help="Control frequency in Hz.") -@click.option('--command_latency', '-cl', default=0.01, type=float, help="Latency between receiving SapceMouse command to executing on Robot in Sec.") -def main(output, robot_ip, vis_camera_idx, init_joints, frequency, command_latency): - dt = 1/frequency - with SharedMemoryManager() as shm_manager: - with KeystrokeCounter() as key_counter, \ - Spacemouse(shm_manager=shm_manager) as sm, \ - RealEnv( - output_dir=output, - robot_ip=robot_ip, - # recording resolution - obs_image_resolution=(1280,720), - frequency=frequency, - init_joints=init_joints, - enable_multi_cam_vis=True, - record_raw_video=True, - # number of threads per camera view for video recording (H.264) - thread_per_video=3, - # video recording quality, lower is better (but slower). - video_crf=21, - shm_manager=shm_manager - ) as env: - cv2.setNumThreads(1) - - # realsense exposure - env.realsense.set_exposure(exposure=120, gain=0) - # realsense white balance - env.realsense.set_white_balance(white_balance=5900) - - time.sleep(1.0) - print('Ready!') - state = env.get_robot_state() - target_pose = state['TargetTCPPose'] - t_start = time.monotonic() - iter_idx = 0 - stop = False - is_recording = False - while not stop: - # calculate timing - t_cycle_end = t_start + (iter_idx + 1) * dt - t_sample = t_cycle_end - command_latency - t_command_target = t_cycle_end + dt - - # pump obs - obs = env.get_obs() - - # handle key presses - press_events = key_counter.get_press_events() - for key_stroke in press_events: - if key_stroke == KeyCode(char='q'): - # Exit program - stop = True - elif key_stroke == KeyCode(char='c'): - # Start recording - env.start_episode(t_start + (iter_idx + 2) * dt - time.monotonic() + time.time()) - key_counter.clear() - is_recording = True - print('Recording!') - elif key_stroke == KeyCode(char='s'): - # Stop recording - env.end_episode() - key_counter.clear() - is_recording = False - print('Stopped.') - elif key_stroke == Key.backspace: - # Delete the most recent recorded episode - if click.confirm('Are you sure to drop an episode?'): - env.drop_episode() - key_counter.clear() - is_recording = False - # delete - stage = key_counter[Key.space] - - # visualize - vis_img = obs[f'camera_{vis_camera_idx}'][-1,:,:,::-1].copy() - episode_id = env.replay_buffer.n_episodes - text = f'Episode: {episode_id}, Stage: {stage}' - if is_recording: - text += ', Recording!' - cv2.putText( - vis_img, - text, - (10,30), - fontFace=cv2.FONT_HERSHEY_SIMPLEX, - fontScale=1, - thickness=2, - color=(255,255,255) - ) - - cv2.imshow('default', vis_img) - cv2.pollKey() - - precise_wait(t_sample) - # get teleop command - sm_state = sm.get_motion_state_transformed() - # print(sm_state) - dpos = sm_state[:3] * (env.max_pos_speed / frequency) - drot_xyz = sm_state[3:] * (env.max_rot_speed / frequency) - - if not sm.is_button_pressed(0): - # translation mode - drot_xyz[:] = 0 - else: - dpos[:] = 0 - if not sm.is_button_pressed(1): - # 2D translation mode - dpos[2] = 0 - - drot = st.Rotation.from_euler('xyz', drot_xyz) - target_pose[:3] += dpos - target_pose[3:] = (drot * st.Rotation.from_rotvec( - target_pose[3:])).as_rotvec() - - # execute teleop command - env.exec_actions( - actions=[target_pose], - timestamps=[t_command_target-time.monotonic()+time.time()], - stages=[stage]) - precise_wait(t_cycle_end) - iter_idx += 1 - -# %% -if __name__ == '__main__': - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/codecs/imagecodecs_numcodecs.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/codecs/imagecodecs_numcodecs.py deleted file mode 100644 index cf55b97b4..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/codecs/imagecodecs_numcodecs.py +++ /dev/null @@ -1,1386 +0,0 @@ - -# imagecodecs/numcodecs.py - -# Copyright (c) 2021-2022, Christoph Gohlke -# All rights reserved. -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, -# this list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE -# POSSIBILITY OF SUCH DAMAGE. - -"""Additional numcodecs implemented using imagecodecs.""" - -__version__ = '2022.9.26' - -__all__ = ('register_codecs',) - -import numpy -from numcodecs.abc import Codec -from numcodecs.registry import register_codec, get_codec - -import imagecodecs - - -def protective_squeeze(x: numpy.ndarray): - """ - Squeeze dim only if it's not the last dim. - Image dim expected to be *, H, W, C - """ - img_shape = x.shape[-3:] - if len(x.shape) > 3: - n_imgs = numpy.prod(x.shape[:-3]) - if n_imgs > 1: - img_shape = (-1,) + img_shape - return x.reshape(img_shape) - -def get_default_image_compressor(**kwargs): - if imagecodecs.JPEGXL: - # has JPEGXL - this_kwargs = { - 'effort': 3, - 'distance': 0.3, - # bug in libjxl, invalid codestream for non-lossless - # when decoding speed > 1 - 'decodingspeed': 1 - } - this_kwargs.update(kwargs) - return JpegXl(**this_kwargs) - else: - this_kwargs = { - 'level': 50 - } - this_kwargs.update(kwargs) - return Jpeg2k(**this_kwargs) - -class Aec(Codec): - """AEC codec for numcodecs.""" - - codec_id = 'imagecodecs_aec' - - def __init__( - self, bitspersample=None, flags=None, blocksize=None, rsi=None - ): - self.bitspersample = bitspersample - self.flags = flags - self.blocksize = blocksize - self.rsi = rsi - - def encode(self, buf): - return imagecodecs.aec_encode( - buf, - bitspersample=self.bitspersample, - flags=self.flags, - blocksize=self.blocksize, - rsi=self.rsi, - ) - - def decode(self, buf, out=None): - return imagecodecs.aec_decode( - buf, - bitspersample=self.bitspersample, - flags=self.flags, - blocksize=self.blocksize, - rsi=self.rsi, - out=_flat(out), - ) - - -class Apng(Codec): - """APNG codec for numcodecs.""" - - codec_id = 'imagecodecs_apng' - - def __init__(self, level=None, photometric=None, delay=None): - self.level = level - self.photometric = photometric - self.delay = delay - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.apng_encode( - buf, - level=self.level, - photometric=self.photometric, - delay=self.delay, - ) - - def decode(self, buf, out=None): - return imagecodecs.apng_decode(buf, out=out) - - -class Avif(Codec): - """AVIF codec for numcodecs.""" - - codec_id = 'imagecodecs_avif' - - def __init__( - self, - level=None, - speed=None, - tilelog2=None, - bitspersample=None, - pixelformat=None, - numthreads=None, - index=None, - ): - self.level = level - self.speed = speed - self.tilelog2 = tilelog2 - self.bitspersample = bitspersample - self.pixelformat = pixelformat - self.numthreads = numthreads - self.index = index - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.avif_encode( - buf, - level=self.level, - speed=self.speed, - tilelog2=self.tilelog2, - bitspersample=self.bitspersample, - pixelformat=self.pixelformat, - numthreads=self.numthreads, - ) - - def decode(self, buf, out=None): - return imagecodecs.avif_decode( - buf, index=self.index, numthreads=self.numthreads, out=out - ) - - -class Bitorder(Codec): - """Bitorder codec for numcodecs.""" - - codec_id = 'imagecodecs_bitorder' - - def encode(self, buf): - return imagecodecs.bitorder_encode(buf) - - def decode(self, buf, out=None): - return imagecodecs.bitorder_decode(buf, out=_flat(out)) - - -class Bitshuffle(Codec): - """Bitshuffle codec for numcodecs.""" - - codec_id = 'imagecodecs_bitshuffle' - - def __init__(self, itemsize=1, blocksize=0): - self.itemsize = itemsize - self.blocksize = blocksize - - def encode(self, buf): - return imagecodecs.bitshuffle_encode( - buf, itemsize=self.itemsize, blocksize=self.blocksize - ).tobytes() - - def decode(self, buf, out=None): - return imagecodecs.bitshuffle_decode( - buf, - itemsize=self.itemsize, - blocksize=self.blocksize, - out=_flat(out), - ) - - -class Blosc(Codec): - """Blosc codec for numcodecs.""" - - codec_id = 'imagecodecs_blosc' - - def __init__( - self, - level=None, - compressor=None, - typesize=None, - blocksize=None, - shuffle=None, - numthreads=None, - ): - self.level = level - self.compressor = compressor - self.typesize = typesize - self.blocksize = blocksize - self.shuffle = shuffle - self.numthreads = numthreads - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.blosc_encode( - buf, - level=self.level, - compressor=self.compressor, - typesize=self.typesize, - blocksize=self.blocksize, - shuffle=self.shuffle, - numthreads=self.numthreads, - ) - - def decode(self, buf, out=None): - return imagecodecs.blosc_decode( - buf, numthreads=self.numthreads, out=_flat(out) - ) - - -class Blosc2(Codec): - """Blosc2 codec for numcodecs.""" - - codec_id = 'imagecodecs_blosc2' - - def __init__( - self, - level=None, - compressor=None, - typesize=None, - blocksize=None, - shuffle=None, - numthreads=None, - ): - self.level = level - self.compressor = compressor - self.typesize = typesize - self.blocksize = blocksize - self.shuffle = shuffle - self.numthreads = numthreads - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.blosc2_encode( - buf, - level=self.level, - compressor=self.compressor, - typesize=self.typesize, - blocksize=self.blocksize, - shuffle=self.shuffle, - numthreads=self.numthreads, - ) - - def decode(self, buf, out=None): - return imagecodecs.blosc2_decode( - buf, numthreads=self.numthreads, out=_flat(out) - ) - - -class Brotli(Codec): - """Brotli codec for numcodecs.""" - - codec_id = 'imagecodecs_brotli' - - def __init__(self, level=None, mode=None, lgwin=None): - self.level = level - self.mode = mode - self.lgwin = lgwin - - def encode(self, buf): - return imagecodecs.brotli_encode( - buf, level=self.level, mode=self.mode, lgwin=self.lgwin - ) - - def decode(self, buf, out=None): - return imagecodecs.brotli_decode(buf, out=_flat(out)) - - -class ByteShuffle(Codec): - """ByteShuffle codec for numcodecs.""" - - codec_id = 'imagecodecs_byteshuffle' - - def __init__( - self, shape, dtype, axis=-1, dist=1, delta=False, reorder=False - ): - self.shape = tuple(shape) - self.dtype = numpy.dtype(dtype).str - self.axis = axis - self.dist = dist - self.delta = bool(delta) - self.reorder = bool(reorder) - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - assert buf.shape == self.shape - assert buf.dtype == self.dtype - return imagecodecs.byteshuffle_encode( - buf, - axis=self.axis, - dist=self.dist, - delta=self.delta, - reorder=self.reorder, - ).tobytes() - - def decode(self, buf, out=None): - if not isinstance(buf, numpy.ndarray): - buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape) - return imagecodecs.byteshuffle_decode( - buf, - axis=self.axis, - dist=self.dist, - delta=self.delta, - reorder=self.reorder, - out=out, - ) - - -class Bz2(Codec): - """Bz2 codec for numcodecs.""" - - codec_id = 'imagecodecs_bz2' - - def __init__(self, level=None): - self.level = level - - def encode(self, buf): - return imagecodecs.bz2_encode(buf, level=self.level) - - def decode(self, buf, out=None): - return imagecodecs.bz2_decode(buf, out=_flat(out)) - - -class Cms(Codec): - """CMS codec for numcodecs.""" - - codec_id = 'imagecodecs_cms' - - def __init__(self, *args, **kwargs): - pass - - def encode(self, buf, out=None): - # return imagecodecs.cms_transform(buf) - raise NotImplementedError - - def decode(self, buf, out=None): - # return imagecodecs.cms_transform(buf) - raise NotImplementedError - - -class Deflate(Codec): - """Deflate codec for numcodecs.""" - - codec_id = 'imagecodecs_deflate' - - def __init__(self, level=None, raw=False): - self.level = level - self.raw = bool(raw) - - def encode(self, buf): - return imagecodecs.deflate_encode(buf, level=self.level, raw=self.raw) - - def decode(self, buf, out=None): - return imagecodecs.deflate_decode(buf, out=_flat(out), raw=self.raw) - - -class Delta(Codec): - """Delta codec for numcodecs.""" - - codec_id = 'imagecodecs_delta' - - def __init__(self, shape=None, dtype=None, axis=-1, dist=1): - self.shape = None if shape is None else tuple(shape) - self.dtype = None if dtype is None else numpy.dtype(dtype).str - self.axis = axis - self.dist = dist - - def encode(self, buf): - if self.shape is not None or self.dtype is not None: - buf = protective_squeeze(numpy.asarray(buf)) - assert buf.shape == self.shape - assert buf.dtype == self.dtype - return imagecodecs.delta_encode( - buf, axis=self.axis, dist=self.dist - ).tobytes() - - def decode(self, buf, out=None): - if self.shape is not None or self.dtype is not None: - buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape) - return imagecodecs.delta_decode( - buf, axis=self.axis, dist=self.dist, out=out - ) - - -class Float24(Codec): - """Float24 codec for numcodecs.""" - - codec_id = 'imagecodecs_float24' - - def __init__(self, byteorder=None, rounding=None): - self.byteorder = byteorder - self.rounding = rounding - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.float24_encode( - buf, byteorder=self.byteorder, rounding=self.rounding - ) - - def decode(self, buf, out=None): - return imagecodecs.float24_decode( - buf, byteorder=self.byteorder, out=out - ) - - -class FloatPred(Codec): - """Floating Point Predictor codec for numcodecs.""" - - codec_id = 'imagecodecs_floatpred' - - def __init__(self, shape, dtype, axis=-1, dist=1): - self.shape = tuple(shape) - self.dtype = numpy.dtype(dtype).str - self.axis = axis - self.dist = dist - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - assert buf.shape == self.shape - assert buf.dtype == self.dtype - return imagecodecs.floatpred_encode( - buf, axis=self.axis, dist=self.dist - ).tobytes() - - def decode(self, buf, out=None): - if not isinstance(buf, numpy.ndarray): - buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape) - return imagecodecs.floatpred_decode( - buf, axis=self.axis, dist=self.dist, out=out - ) - - -class Gif(Codec): - """GIF codec for numcodecs.""" - - codec_id = 'imagecodecs_gif' - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.gif_encode(buf) - - def decode(self, buf, out=None): - return imagecodecs.gif_decode(buf, asrgb=False, out=out) - - -class Heif(Codec): - """HEIF codec for numcodecs.""" - - codec_id = 'imagecodecs_heif' - - def __init__( - self, - level=None, - bitspersample=None, - photometric=None, - compression=None, - numthreads=None, - index=None, - ): - self.level = level - self.bitspersample = bitspersample - self.photometric = photometric - self.compression = compression - self.numthreads = numthreads - self.index = index - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.heif_encode( - buf, - level=self.level, - bitspersample=self.bitspersample, - photometric=self.photometric, - compression=self.compression, - numthreads=self.numthreads, - ) - - def decode(self, buf, out=None): - return imagecodecs.heif_decode( - buf, - index=self.index, - photometric=self.photometric, - numthreads=self.numthreads, - out=out, - ) - - -class Jetraw(Codec): - """Jetraw codec for numcodecs.""" - - codec_id = 'imagecodecs_jetraw' - - def __init__( - self, - shape, - identifier, - parameters=None, - verbosity=None, - errorbound=None, - ): - self.shape = shape - self.identifier = identifier - self.errorbound = errorbound - imagecodecs.jetraw_init(parameters, verbosity) - - def encode(self, buf): - return imagecodecs.jetraw_encode( - buf, identifier=self.identifier, errorbound=self.errorbound - ) - - def decode(self, buf, out=None): - if out is None: - out = numpy.empty(self.shape, numpy.uint16) - return imagecodecs.jetraw_decode(buf, out=out) - - -class Jpeg(Codec): - """JPEG codec for numcodecs.""" - - codec_id = 'imagecodecs_jpeg' - - def __init__( - self, - bitspersample=None, - tables=None, - header=None, - colorspace_data=None, - colorspace_jpeg=None, - level=None, - subsampling=None, - optimize=None, - smoothing=None, - ): - self.tables = tables - self.header = header - self.bitspersample = bitspersample - self.colorspace_data = colorspace_data - self.colorspace_jpeg = colorspace_jpeg - self.level = level - self.subsampling = subsampling - self.optimize = optimize - self.smoothing = smoothing - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.jpeg_encode( - buf, - level=self.level, - colorspace=self.colorspace_data, - outcolorspace=self.colorspace_jpeg, - subsampling=self.subsampling, - optimize=self.optimize, - smoothing=self.smoothing, - ) - - def decode(self, buf, out=None): - out_shape = None - if out is not None: - out_shape = out.shape - out = protective_squeeze(out) - img = imagecodecs.jpeg_decode( - buf, - bitspersample=self.bitspersample, - tables=self.tables, - header=self.header, - colorspace=self.colorspace_jpeg, - outcolorspace=self.colorspace_data, - out=out, - ) - if out_shape is not None: - img = img.reshape(out_shape) - return img - - def get_config(self): - """Return dictionary holding configuration parameters.""" - config = dict(id=self.codec_id) - for key in self.__dict__: - if not key.startswith('_'): - value = getattr(self, key) - if value is not None and key in ('header', 'tables'): - import base64 - - value = base64.b64encode(value).decode() - config[key] = value - return config - - @classmethod - def from_config(cls, config): - """Instantiate codec from configuration object.""" - for key in ('header', 'tables'): - value = config.get(key, None) - if value is not None and isinstance(value, str): - import base64 - - config[key] = base64.b64decode(value.encode()) - return cls(**config) - - -class Jpeg2k(Codec): - """JPEG 2000 codec for numcodecs.""" - - codec_id = 'imagecodecs_jpeg2k' - - def __init__( - self, - level=None, - codecformat=None, - colorspace=None, - tile=None, - reversible=None, - bitspersample=None, - resolutions=None, - numthreads=None, - verbose=0, - ): - self.level = level - self.codecformat = codecformat - self.colorspace = colorspace - self.tile = None if tile is None else tuple(tile) - self.reversible = reversible - self.bitspersample = bitspersample - self.resolutions = resolutions - self.numthreads = numthreads - self.verbose = verbose - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.jpeg2k_encode( - buf, - level=self.level, - codecformat=self.codecformat, - colorspace=self.colorspace, - tile=self.tile, - reversible=self.reversible, - bitspersample=self.bitspersample, - resolutions=self.resolutions, - numthreads=self.numthreads, - verbose=self.verbose, - ) - - def decode(self, buf, out=None): - return imagecodecs.jpeg2k_decode( - buf, verbose=self.verbose, numthreads=self.numthreads, out=out - ) - - -class JpegLs(Codec): - """JPEG LS codec for numcodecs.""" - - codec_id = 'imagecodecs_jpegls' - - def __init__(self, level=None): - self.level = level - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.jpegls_encode(buf, level=self.level) - - def decode(self, buf, out=None): - return imagecodecs.jpegls_decode(buf, out=out) - - -class JpegXl(Codec): - """JPEG XL codec for numcodecs.""" - - codec_id = 'imagecodecs_jpegxl' - - def __init__( - self, - # encode - level=None, - effort=None, - distance=None, - lossless=None, - decodingspeed=None, - photometric=None, - planar=None, - usecontainer=None, - # decode - index=None, - keeporientation=None, - # both - numthreads=None, - ): - """ - Return JPEG XL image from numpy array. - Float must be in nominal range 0..1. - - Currently L, LA, RGB, RGBA images are supported in contig mode. - Extra channels are only supported for grayscale images in planar mode. - - Parameters - ---------- - level : Default to None, i.e. not overwriting lossess and decodingspeed options. - When < 0: Use lossless compression - When in [0,1,2,3,4]: Sets the decoding speed tier for the provided options. - Minimum is 0 (slowest to decode, best quality/density), and maximum - is 4 (fastest to decode, at the cost of some quality/density). - effort : Default to 3. - Sets encoder effort/speed level without affecting decoding speed. - Valid values are, from faster to slower speed: 1:lightning 2:thunder - 3:falcon 4:cheetah 5:hare 6:wombat 7:squirrel 8:kitten 9:tortoise. - Speed: lightning, thunder, falcon, cheetah, hare, wombat, squirrel, kitten, tortoise - control the encoder effort in ascending order. - This also affects memory usage: using lower effort will typically reduce memory - consumption during encoding. - lightning and thunder are fast modes useful for lossless mode (modular). - falcon disables all of the following tools. - cheetah enables coefficient reordering, context clustering, and heuristics for selecting DCT sizes and quantization steps. - hare enables Gaborish filtering, chroma from luma, and an initial estimate of quantization steps. - wombat enables error diffusion quantization and full DCT size selection heuristics. - squirrel (default) enables dots, patches, and spline detection, and full context clustering. - kitten optimizes the adaptive quantization for a psychovisual metric. - tortoise enables a more thorough adaptive quantization search. - distance : Default to 1.0 - Sets the distance level for lossy compression: target max butteraugli distance, - lower = higher quality. Range: 0 .. 15. 0.0 = mathematically lossless - (however, use JxlEncoderSetFrameLossless instead to use true lossless, - as setting distance to 0 alone is not the only requirement). - 1.0 = visually lossless. Recommended range: 0.5 .. 3.0. - lossess : Default to False. - Use lossess encoding. - decodingspeed : Default to 0. - Duplicate to level. [0,4] - photometric : Return JxlColorSpace value. - Default logic is quite complicated but works most of the time. - Accepted value: - int: [-1,3] - str: ['RGB', - 'WHITEISZERO', 'MINISWHITE', - 'BLACKISZERO', 'MINISBLACK', 'GRAY', - 'XYB', 'KNOWN'] - planar : Enable multi-channel mode. - Default to false. - usecontainer : - Forces the encoder to use the box-based container format (BMFF) - even when not necessary. - When using JxlEncoderUseBoxes, JxlEncoderStoreJPEGMetadata or - JxlEncoderSetCodestreamLevel with level 10, the encoder will - automatically also use the container format, it is not necessary - to use JxlEncoderUseContainer for those use cases. - By default this setting is disabled. - index : Selectively decode frames for animation. - Default to 0, decode all frames. - When set to > 0, decode that frame index only. - keeporientation : - Enables or disables preserving of as-in-bitstream pixeldata orientation. - Some images are encoded with an Orientation tag indicating that the - decoder must perform a rotation and/or mirroring to the encoded image data. - - If skip_reorientation is JXL_FALSE (the default): the decoder will apply - the transformation from the orientation setting, hence rendering the image - according to its specified intent. When producing a JxlBasicInfo, the decoder - will always set the orientation field to JXL_ORIENT_IDENTITY (matching the - returned pixel data) and also align xsize and ysize so that they correspond - to the width and the height of the returned pixel data. - - If skip_reorientation is JXL_TRUE: the decoder will skip applying the - transformation from the orientation setting, returning the image in - the as-in-bitstream pixeldata orientation. This may be faster to decode - since the decoder doesnt have to apply the transformation, but can - cause wrong display of the image if the orientation tag is not correctly - taken into account by the user. - - By default, this option is disabled, and the returned pixel data is - re-oriented according to the images Orientation setting. - threads : Default to 1. - If <= 0, use all cores. - If > 32, clipped to 32. - """ - - self.level = level - self.effort = effort - self.distance = distance - self.lossless = bool(lossless) - self.decodingspeed = decodingspeed - self.photometric = photometric - self.planar = planar - self.usecontainer = usecontainer - self.index = index - self.keeporientation = keeporientation - self.numthreads = numthreads - - def encode(self, buf): - # TODO: only squeeze all but last dim - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.jpegxl_encode( - buf, - level=self.level, - effort=self.effort, - distance=self.distance, - lossless=self.lossless, - decodingspeed=self.decodingspeed, - photometric=self.photometric, - planar=self.planar, - usecontainer=self.usecontainer, - numthreads=self.numthreads, - ) - - def decode(self, buf, out=None): - return imagecodecs.jpegxl_decode( - buf, - index=self.index, - keeporientation=self.keeporientation, - numthreads=self.numthreads, - out=out, - ) - - -class JpegXr(Codec): - """JPEG XR codec for numcodecs.""" - - codec_id = 'imagecodecs_jpegxr' - - def __init__( - self, - level=None, - photometric=None, - hasalpha=None, - resolution=None, - fp2int=None, - ): - self.level = level - self.photometric = photometric - self.hasalpha = hasalpha - self.resolution = resolution - self.fp2int = fp2int - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.jpegxr_encode( - buf, - level=self.level, - photometric=self.photometric, - hasalpha=self.hasalpha, - resolution=self.resolution, - ) - - def decode(self, buf, out=None): - return imagecodecs.jpegxr_decode(buf, fp2int=self.fp2int, out=out) - - -class Lerc(Codec): - """LERC codec for numcodecs.""" - - codec_id = 'imagecodecs_lerc' - - def __init__(self, level=None, version=None, planar=None): - self.level = level - self.version = version - self.planar = bool(planar) - # TODO: support mask? - # self.mask = None - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.lerc_encode( - buf, - level=self.level, - version=self.version, - planar=self.planar, - ) - - def decode(self, buf, out=None): - return imagecodecs.lerc_decode(buf, out=out) - - -class Ljpeg(Codec): - """LJPEG codec for numcodecs.""" - - codec_id = 'imagecodecs_ljpeg' - - def __init__(self, bitspersample=None): - self.bitspersample = bitspersample - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.ljpeg_encode(buf, bitspersample=self.bitspersample) - - def decode(self, buf, out=None): - return imagecodecs.ljpeg_decode(buf, out=out) - - -class Lz4(Codec): - """LZ4 codec for numcodecs.""" - - codec_id = 'imagecodecs_lz4' - - def __init__(self, level=None, hc=False, header=True): - self.level = level - self.hc = hc - self.header = bool(header) - - def encode(self, buf): - return imagecodecs.lz4_encode( - buf, level=self.level, hc=self.hc, header=self.header - ) - - def decode(self, buf, out=None): - return imagecodecs.lz4_decode(buf, header=self.header, out=_flat(out)) - - -class Lz4f(Codec): - """LZ4F codec for numcodecs.""" - - codec_id = 'imagecodecs_lz4f' - - def __init__( - self, - level=None, - blocksizeid=False, - contentchecksum=None, - blockchecksum=None, - ): - self.level = level - self.blocksizeid = blocksizeid - self.contentchecksum = contentchecksum - self.blockchecksum = blockchecksum - - def encode(self, buf): - return imagecodecs.lz4f_encode( - buf, - level=self.level, - blocksizeid=self.blocksizeid, - contentchecksum=self.contentchecksum, - blockchecksum=self.blockchecksum, - ) - - def decode(self, buf, out=None): - return imagecodecs.lz4f_decode(buf, out=_flat(out)) - - -class Lzf(Codec): - """LZF codec for numcodecs.""" - - codec_id = 'imagecodecs_lzf' - - def __init__(self, header=True): - self.header = bool(header) - - def encode(self, buf): - return imagecodecs.lzf_encode(buf, header=self.header) - - def decode(self, buf, out=None): - return imagecodecs.lzf_decode(buf, header=self.header, out=_flat(out)) - - -class Lzma(Codec): - """LZMA codec for numcodecs.""" - - codec_id = 'imagecodecs_lzma' - - def __init__(self, level=None): - self.level = level - - def encode(self, buf): - return imagecodecs.lzma_encode(buf, level=self.level) - - def decode(self, buf, out=None): - return imagecodecs.lzma_decode(buf, out=_flat(out)) - - -class Lzw(Codec): - """LZW codec for numcodecs.""" - - codec_id = 'imagecodecs_lzw' - - def encode(self, buf): - return imagecodecs.lzw_encode(buf) - - def decode(self, buf, out=None): - return imagecodecs.lzw_decode(buf, out=_flat(out)) - - -class PackBits(Codec): - """PackBits codec for numcodecs.""" - - codec_id = 'imagecodecs_packbits' - - def __init__(self, axis=None): - self.axis = axis - - def encode(self, buf): - if not isinstance(buf, (bytes, bytearray)): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.packbits_encode(buf, axis=self.axis) - - def decode(self, buf, out=None): - return imagecodecs.packbits_decode(buf, out=_flat(out)) - - -class Pglz(Codec): - """PGLZ codec for numcodecs.""" - - codec_id = 'imagecodecs_pglz' - - def __init__(self, header=True, strategy=None): - self.header = bool(header) - self.strategy = strategy - - def encode(self, buf): - return imagecodecs.pglz_encode( - buf, strategy=self.strategy, header=self.header - ) - - def decode(self, buf, out=None): - return imagecodecs.pglz_decode(buf, header=self.header, out=_flat(out)) - - -class Png(Codec): - """PNG codec for numcodecs.""" - - codec_id = 'imagecodecs_png' - - def __init__(self, level=None): - self.level = level - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.png_encode(buf, level=self.level) - - def decode(self, buf, out=None): - return imagecodecs.png_decode(buf, out=out) - - -class Qoi(Codec): - """QOI codec for numcodecs.""" - - codec_id = 'imagecodecs_qoi' - - def __init__(self): - pass - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.qoi_encode(buf) - - def decode(self, buf, out=None): - return imagecodecs.qoi_decode(buf, out=out) - - -class Rgbe(Codec): - """RGBE codec for numcodecs.""" - - codec_id = 'imagecodecs_rgbe' - - def __init__(self, header=False, shape=None, rle=None): - if not header and shape is None: - raise ValueError('must specify data shape if no header') - if shape and shape[-1] != 3: - raise ValueError('invalid shape') - self.shape = shape - self.header = bool(header) - self.rle = None if rle is None else bool(rle) - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.rgbe_encode(buf, header=self.header, rle=self.rle) - - def decode(self, buf, out=None): - if out is None and not self.header: - out = numpy.empty(self.shape, numpy.float32) - return imagecodecs.rgbe_decode( - buf, header=self.header, rle=self.rle, out=out - ) - - -class Rcomp(Codec): - """Rcomp codec for numcodecs.""" - - codec_id = 'imagecodecs_rcomp' - - def __init__(self, shape, dtype, nblock=None): - self.shape = tuple(shape) - self.dtype = numpy.dtype(dtype).str - self.nblock = nblock - - def encode(self, buf): - return imagecodecs.rcomp_encode(buf, nblock=self.nblock) - - def decode(self, buf, out=None): - return imagecodecs.rcomp_decode( - buf, - shape=self.shape, - dtype=self.dtype, - nblock=self.nblock, - out=out, - ) - - -class Snappy(Codec): - """Snappy codec for numcodecs.""" - - codec_id = 'imagecodecs_snappy' - - def encode(self, buf): - return imagecodecs.snappy_encode(buf) - - def decode(self, buf, out=None): - return imagecodecs.snappy_decode(buf, out=_flat(out)) - - -class Spng(Codec): - """SPNG codec for numcodecs.""" - - codec_id = 'imagecodecs_spng' - - def __init__(self, level=None): - self.level = level - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.spng_encode(buf, level=self.level) - - def decode(self, buf, out=None): - return imagecodecs.spng_decode(buf, out=out) - - -class Tiff(Codec): - """TIFF codec for numcodecs.""" - - codec_id = 'imagecodecs_tiff' - - def __init__(self, index=None, asrgb=None, verbose=0): - self.index = index - self.asrgb = bool(asrgb) - self.verbose = verbose - - def encode(self, buf): - # TODO: not implemented - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.tiff_encode(buf) - - def decode(self, buf, out=None): - return imagecodecs.tiff_decode( - buf, - index=self.index, - asrgb=self.asrgb, - verbose=self.verbose, - out=out, - ) - - -class Webp(Codec): - """WebP codec for numcodecs.""" - - codec_id = 'imagecodecs_webp' - - def __init__(self, level=None, lossless=None, method=None, hasalpha=None): - self.level = level - self.hasalpha = bool(hasalpha) - self.method = method - self.lossless = lossless - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - return imagecodecs.webp_encode( - buf, level=self.level, lossless=self.lossless, method=self.method - ) - - def decode(self, buf, out=None): - return imagecodecs.webp_decode(buf, hasalpha=self.hasalpha, out=out) - - -class Xor(Codec): - """XOR codec for numcodecs.""" - - codec_id = 'imagecodecs_xor' - - def __init__(self, shape=None, dtype=None, axis=-1): - self.shape = None if shape is None else tuple(shape) - self.dtype = None if dtype is None else numpy.dtype(dtype).str - self.axis = axis - - def encode(self, buf): - if self.shape is not None or self.dtype is not None: - buf = protective_squeeze(numpy.asarray(buf)) - assert buf.shape == self.shape - assert buf.dtype == self.dtype - return imagecodecs.xor_encode(buf, axis=self.axis).tobytes() - - def decode(self, buf, out=None): - if self.shape is not None or self.dtype is not None: - buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape) - return imagecodecs.xor_decode(buf, axis=self.axis, out=_flat(out)) - - -class Zfp(Codec): - """ZFP codec for numcodecs.""" - - codec_id = 'imagecodecs_zfp' - - def __init__( - self, - shape=None, - dtype=None, - strides=None, - level=None, - mode=None, - execution=None, - numthreads=None, - chunksize=None, - header=True, - ): - if header: - self.shape = None - self.dtype = None - self.strides = None - elif shape is None or dtype is None: - raise ValueError('invalid shape or dtype') - else: - self.shape = tuple(shape) - self.dtype = numpy.dtype(dtype).str - self.strides = None if strides is None else tuple(strides) - self.level = level - self.mode = mode - self.execution = execution - self.numthreads = numthreads - self.chunksize = chunksize - self.header = bool(header) - - def encode(self, buf): - buf = protective_squeeze(numpy.asarray(buf)) - if not self.header: - assert buf.shape == self.shape - assert buf.dtype == self.dtype - return imagecodecs.zfp_encode( - buf, - level=self.level, - mode=self.mode, - execution=self.execution, - header=self.header, - numthreads=self.numthreads, - chunksize=self.chunksize, - ) - - def decode(self, buf, out=None): - if self.header: - return imagecodecs.zfp_decode(buf, out=out) - return imagecodecs.zfp_decode( - buf, - shape=self.shape, - dtype=numpy.dtype(self.dtype), - strides=self.strides, - numthreads=self.numthreads, - out=out, - ) - - -class Zlib(Codec): - """Zlib codec for numcodecs.""" - - codec_id = 'imagecodecs_zlib' - - def __init__(self, level=None): - self.level = level - - def encode(self, buf): - return imagecodecs.zlib_encode(buf, level=self.level) - - def decode(self, buf, out=None): - return imagecodecs.zlib_decode(buf, out=_flat(out)) - - -class Zlibng(Codec): - """Zlibng codec for numcodecs.""" - - codec_id = 'imagecodecs_zlibng' - - def __init__(self, level=None): - self.level = level - - def encode(self, buf): - return imagecodecs.zlibng_encode(buf, level=self.level) - - def decode(self, buf, out=None): - return imagecodecs.zlibng_decode(buf, out=_flat(out)) - - -class Zopfli(Codec): - """Zopfli codec for numcodecs.""" - - codec_id = 'imagecodecs_zopfli' - - def encode(self, buf): - return imagecodecs.zopfli_encode(buf) - - def decode(self, buf, out=None): - return imagecodecs.zopfli_decode(buf, out=_flat(out)) - - -class Zstd(Codec): - """ZStandard codec for numcodecs.""" - - codec_id = 'imagecodecs_zstd' - - def __init__(self, level=None): - self.level = level - - def encode(self, buf): - return imagecodecs.zstd_encode(buf, level=self.level) - - def decode(self, buf, out=None): - return imagecodecs.zstd_decode(buf, out=_flat(out)) - - -def _flat(out): - """Return numpy array as contiguous view of bytes if possible.""" - if out is None: - return None - view = memoryview(out) - if view.readonly or not view.contiguous: - return None - return view.cast('B') - - -def register_codecs(codecs=None, force=False, verbose=True): - """Register codecs in this module with numcodecs.""" - for name, cls in globals().items(): - if not hasattr(cls, 'codec_id') or name == 'Codec': - continue - if codecs is not None and cls.codec_id not in codecs: - continue - try: - try: - get_codec({'id': cls.codec_id}) - except TypeError: - # registered, but failed - pass - except ValueError: - # not registered yet - pass - else: - if not force: - if verbose: - log_warning( - f'numcodec {cls.codec_id!r} already registered' - ) - continue - if verbose: - log_warning(f'replacing registered numcodec {cls.codec_id!r}') - register_codec(cls) - - -def log_warning(msg, *args, **kwargs): - """Log message with level WARNING.""" - import logging - - logging.getLogger(__name__).warning(msg, *args, **kwargs) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/checkpoint_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/checkpoint_util.py deleted file mode 100644 index fbb083713..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/checkpoint_util.py +++ /dev/null @@ -1,59 +0,0 @@ -from typing import Optional, Dict -import os - -class TopKCheckpointManager: - def __init__(self, - save_dir, - monitor_key: str, - mode='min', - k=1, - format_str='epoch={epoch:03d}-train_loss={train_loss:.3f}.ckpt' - ): - assert mode in ['max', 'min'] - assert k >= 0 - - self.save_dir = save_dir - self.monitor_key = monitor_key - self.mode = mode - self.k = k - self.format_str = format_str - self.path_value_map = dict() - - def get_ckpt_path(self, data: Dict[str, float]) -> Optional[str]: - if self.k == 0: - return None - - value = data[self.monitor_key] - ckpt_path = os.path.join( - self.save_dir, self.format_str.format(**data)) - - if len(self.path_value_map) < self.k: - # under-capacity - self.path_value_map[ckpt_path] = value - return ckpt_path - - # at capacity - sorted_map = sorted(self.path_value_map.items(), key=lambda x: x[1]) - min_path, min_value = sorted_map[0] - max_path, max_value = sorted_map[-1] - - delete_path = None - if self.mode == 'max': - if value > min_value: - delete_path = min_path - else: - if value < max_value: - delete_path = max_path - - if delete_path is None: - return None - else: - del self.path_value_map[delete_path] - self.path_value_map[ckpt_path] = value - - if not os.path.exists(self.save_dir): - os.mkdir(self.save_dir) - - if os.path.exists(delete_path): - os.remove(delete_path) - return ckpt_path diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/cv2_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/cv2_util.py deleted file mode 100644 index 1d749bcec..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/cv2_util.py +++ /dev/null @@ -1,150 +0,0 @@ -from typing import Tuple -import math -import cv2 -import numpy as np - -def draw_reticle(img, u, v, label_color): - """ - Draws a reticle (cross-hair) on the image at the given position on top of - the original image. - @param img (In/Out) uint8 3 channel image - @param u X coordinate (width) - @param v Y coordinate (height) - @param label_color tuple of 3 ints for RGB color used for drawing. - """ - # Cast to int. - u = int(u) - v = int(v) - - white = (255, 255, 255) - cv2.circle(img, (u, v), 10, label_color, 1) - cv2.circle(img, (u, v), 11, white, 1) - cv2.circle(img, (u, v), 12, label_color, 1) - cv2.line(img, (u, v + 1), (u, v + 3), white, 1) - cv2.line(img, (u + 1, v), (u + 3, v), white, 1) - cv2.line(img, (u, v - 1), (u, v - 3), white, 1) - cv2.line(img, (u - 1, v), (u - 3, v), white, 1) - - -def draw_text( - img, - *, - text, - uv_top_left, - color=(255, 255, 255), - fontScale=0.5, - thickness=1, - fontFace=cv2.FONT_HERSHEY_SIMPLEX, - outline_color=(0, 0, 0), - line_spacing=1.5, -): - """ - Draws multiline with an outline. - """ - assert isinstance(text, str) - - uv_top_left = np.array(uv_top_left, dtype=float) - assert uv_top_left.shape == (2,) - - for line in text.splitlines(): - (w, h), _ = cv2.getTextSize( - text=line, - fontFace=fontFace, - fontScale=fontScale, - thickness=thickness, - ) - uv_bottom_left_i = uv_top_left + [0, h] - org = tuple(uv_bottom_left_i.astype(int)) - - if outline_color is not None: - cv2.putText( - img, - text=line, - org=org, - fontFace=fontFace, - fontScale=fontScale, - color=outline_color, - thickness=thickness * 3, - lineType=cv2.LINE_AA, - ) - cv2.putText( - img, - text=line, - org=org, - fontFace=fontFace, - fontScale=fontScale, - color=color, - thickness=thickness, - lineType=cv2.LINE_AA, - ) - - uv_top_left += [0, h * line_spacing] - - -def get_image_transform( - input_res: Tuple[int,int]=(1280,720), - output_res: Tuple[int,int]=(640,480), - bgr_to_rgb: bool=False): - - iw, ih = input_res - ow, oh = output_res - rw, rh = None, None - interp_method = cv2.INTER_AREA - - if (iw/ih) >= (ow/oh): - # input is wider - rh = oh - rw = math.ceil(rh / ih * iw) - if oh > ih: - interp_method = cv2.INTER_LINEAR - else: - rw = ow - rh = math.ceil(rw / iw * ih) - if ow > iw: - interp_method = cv2.INTER_LINEAR - - w_slice_start = (rw - ow) // 2 - w_slice = slice(w_slice_start, w_slice_start + ow) - h_slice_start = (rh - oh) // 2 - h_slice = slice(h_slice_start, h_slice_start + oh) - c_slice = slice(None) - if bgr_to_rgb: - c_slice = slice(None, None, -1) - - def transform(img: np.ndarray): - assert img.shape == ((ih,iw,3)) - # resize - img = cv2.resize(img, (rw, rh), interpolation=interp_method) - # crop - img = img[h_slice, w_slice, c_slice] - return img - return transform - -def optimal_row_cols( - n_cameras, - in_wh_ratio, - max_resolution=(1920, 1080) - ): - out_w, out_h = max_resolution - out_wh_ratio = out_w / out_h - - n_rows = np.arange(n_cameras,dtype=np.int64) + 1 - n_cols = np.ceil(n_cameras / n_rows).astype(np.int64) - cat_wh_ratio = in_wh_ratio * (n_cols / n_rows) - ratio_diff = np.abs(out_wh_ratio - cat_wh_ratio) - best_idx = np.argmin(ratio_diff) - best_n_row = n_rows[best_idx] - best_n_col = n_cols[best_idx] - best_cat_wh_ratio = cat_wh_ratio[best_idx] - - rw, rh = None, None - if best_cat_wh_ratio >= out_wh_ratio: - # cat is wider - rw = math.floor(out_w / best_n_col) - rh = math.floor(rw / in_wh_ratio) - else: - rh = math.floor(out_h / best_n_row) - rw = math.floor(rh * in_wh_ratio) - - # crop_resolution = (rw, rh) - return rw, rh, best_n_col, best_n_row diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/env_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/env_util.py deleted file mode 100644 index a380f5a84..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/env_util.py +++ /dev/null @@ -1,23 +0,0 @@ -import cv2 -import numpy as np - - -def render_env_video(env, states, actions=None): - observations = states - imgs = list() - for i in range(len(observations)): - state = observations[i] - env.set_state(state) - if i == 0: - env.set_state(state) - img = env.render() - # draw action - if actions is not None: - action = actions[i] - coord = (action / 512 * 96).astype(np.int32) - cv2.drawMarker(img, coord, - color=(255,0,0), markerType=cv2.MARKER_CROSS, - markerSize=8, thickness=1) - imgs.append(img) - imgs = np.array(imgs) - return imgs diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/json_logger.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/json_logger.py deleted file mode 100644 index ee9fb4af0..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/json_logger.py +++ /dev/null @@ -1,117 +0,0 @@ -from typing import Optional, Callable, Any, Sequence -import os -import copy -import json -import numbers -import pandas as pd - - -def read_json_log(path: str, - required_keys: Sequence[str]=tuple(), - **kwargs) -> pd.DataFrame: - """ - Read json-per-line file, with potentially incomplete lines. - kwargs passed to pd.read_json - """ - lines = list() - with open(path, 'r') as f: - while True: - # one json per line - line = f.readline() - if len(line) == 0: - # EOF - break - elif not line.endswith('\n'): - # incomplete line - break - is_relevant = False - for k in required_keys: - if k in line: - is_relevant = True - break - if is_relevant: - lines.append(line) - if len(lines) < 1: - return pd.DataFrame() - json_buf = f'[{",".join([line for line in (line.strip() for line in lines) if line])}]' - df = pd.read_json(json_buf, **kwargs) - return df - -class JsonLogger: - def __init__(self, path: str, - filter_fn: Optional[Callable[[str,Any],bool]]=None): - if filter_fn is None: - filter_fn = lambda k,v: isinstance(v, numbers.Number) - - # default to append mode - self.path = path - self.filter_fn = filter_fn - self.file = None - self.last_log = None - - def start(self): - # use line buffering - try: - self.file = file = open(self.path, 'r+', buffering=1) - except FileNotFoundError: - self.file = file = open(self.path, 'w+', buffering=1) - - # Move the pointer (similar to a cursor in a text editor) to the end of the file - pos = file.seek(0, os.SEEK_END) - - # Read each character in the file one at a time from the last - # character going backwards, searching for a newline character - # If we find a new line, exit the search - while pos > 0 and file.read(1) != "\n": - pos -= 1 - file.seek(pos, os.SEEK_SET) - # now the file pointer is at one past the last '\n' - # and pos is at the last '\n'. - last_line_end = file.tell() - - # find the start of second last line - pos = max(0, pos-1) - file.seek(pos, os.SEEK_SET) - while pos > 0 and file.read(1) != "\n": - pos -= 1 - file.seek(pos, os.SEEK_SET) - # now the file pointer is at one past the second last '\n' - last_line_start = file.tell() - - if last_line_start < last_line_end: - # has last line of json - last_line = file.readline() - self.last_log = json.loads(last_line) - - # remove the last incomplete line - file.seek(last_line_end) - file.truncate() - - def stop(self): - self.file.close() - self.file = None - - def __enter__(self): - self.start() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() - - def log(self, data: dict): - filtered_data = dict( - filter(lambda x: self.filter_fn(*x), data.items())) - # save current as last log - self.last_log = filtered_data - for k, v in filtered_data.items(): - if isinstance(v, numbers.Integral): - filtered_data[k] = int(v) - elif isinstance(v, numbers.Number): - filtered_data[k] = float(v) - buf = json.dumps(filtered_data) - # ensure one line per json - buf = buf.replace('\n','') + '\n' - self.file.write(buf) - - def get_last_log(self): - return copy.deepcopy(self.last_log) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/nested_dict_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/nested_dict_util.py deleted file mode 100644 index b186f707f..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/nested_dict_util.py +++ /dev/null @@ -1,32 +0,0 @@ -import functools - -def nested_dict_map(f, x): - """ - Map f over all leaf of nested dict x - """ - - if not isinstance(x, dict): - return f(x) - y = dict() - for key, value in x.items(): - y[key] = nested_dict_map(f, value) - return y - -def nested_dict_reduce(f, x): - """ - Map f over all values of nested dict x, and reduce to a single value - """ - if not isinstance(x, dict): - return x - - reduced_values = list() - for value in x.values(): - reduced_values.append(nested_dict_reduce(f, value)) - y = functools.reduce(f, reduced_values) - return y - - -def nested_dict_check(f, x): - bool_dict = nested_dict_map(f, x) - result = nested_dict_reduce(lambda x, y: x and y, bool_dict) - return result diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/normalize_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/normalize_util.py deleted file mode 100644 index d6fc7b668..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/normalize_util.py +++ /dev/null @@ -1,223 +0,0 @@ -from diffusion_policy.model.common.normalizer import SingleFieldLinearNormalizer -from diffusion_policy.common.pytorch_util import dict_apply, dict_apply_reduce, dict_apply_split -import numpy as np - - -def get_range_normalizer_from_stat(stat, output_max=1, output_min=-1, range_eps=1e-7): - # -1, 1 normalization - input_max = stat['max'] - input_min = stat['min'] - input_range = input_max - input_min - ignore_dim = input_range < range_eps - input_range[ignore_dim] = output_max - output_min - scale = (output_max - output_min) / input_range - offset = output_min - scale * input_min - offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] - - return SingleFieldLinearNormalizer.create_manual( - scale=scale, - offset=offset, - input_stats_dict=stat - ) - -def get_image_range_normalizer(): - scale = np.array([2], dtype=np.float32) - offset = np.array([-1], dtype=np.float32) - stat = { - 'min': np.array([0], dtype=np.float32), - 'max': np.array([1], dtype=np.float32), - 'mean': np.array([0.5], dtype=np.float32), - 'std': np.array([np.sqrt(1/12)], dtype=np.float32) - } - return SingleFieldLinearNormalizer.create_manual( - scale=scale, - offset=offset, - input_stats_dict=stat - ) - -def get_identity_normalizer_from_stat(stat): - scale = np.ones_like(stat['min']) - offset = np.zeros_like(stat['min']) - return SingleFieldLinearNormalizer.create_manual( - scale=scale, - offset=offset, - input_stats_dict=stat - ) - -def robomimic_abs_action_normalizer_from_stat(stat, rotation_transformer): - result = dict_apply_split( - stat, lambda x: { - 'pos': x[...,:3], - 'rot': x[...,3:6], - 'gripper': x[...,6:] - }) - - def get_pos_param_info(stat, output_max=1, output_min=-1, range_eps=1e-7): - # -1, 1 normalization - input_max = stat['max'] - input_min = stat['min'] - input_range = input_max - input_min - ignore_dim = input_range < range_eps - input_range[ignore_dim] = output_max - output_min - scale = (output_max - output_min) / input_range - offset = output_min - scale * input_min - offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] - - return {'scale': scale, 'offset': offset}, stat - - def get_rot_param_info(stat): - example = rotation_transformer.forward(stat['mean']) - scale = np.ones_like(example) - offset = np.zeros_like(example) - info = { - 'max': np.ones_like(example), - 'min': np.full_like(example, -1), - 'mean': np.zeros_like(example), - 'std': np.ones_like(example) - } - return {'scale': scale, 'offset': offset}, info - - def get_gripper_param_info(stat): - example = stat['max'] - scale = np.ones_like(example) - offset = np.zeros_like(example) - info = { - 'max': np.ones_like(example), - 'min': np.full_like(example, -1), - 'mean': np.zeros_like(example), - 'std': np.ones_like(example) - } - return {'scale': scale, 'offset': offset}, info - - pos_param, pos_info = get_pos_param_info(result['pos']) - rot_param, rot_info = get_rot_param_info(result['rot']) - gripper_param, gripper_info = get_gripper_param_info(result['gripper']) - - param = dict_apply_reduce( - [pos_param, rot_param, gripper_param], - lambda x: np.concatenate(x,axis=-1)) - info = dict_apply_reduce( - [pos_info, rot_info, gripper_info], - lambda x: np.concatenate(x,axis=-1)) - - return SingleFieldLinearNormalizer.create_manual( - scale=param['scale'], - offset=param['offset'], - input_stats_dict=info - ) - - -def robomimic_abs_action_only_normalizer_from_stat(stat): - result = dict_apply_split( - stat, lambda x: { - 'pos': x[...,:3], - 'other': x[...,3:] - }) - - def get_pos_param_info(stat, output_max=1, output_min=-1, range_eps=1e-7): - # -1, 1 normalization - input_max = stat['max'] - input_min = stat['min'] - input_range = input_max - input_min - ignore_dim = input_range < range_eps - input_range[ignore_dim] = output_max - output_min - scale = (output_max - output_min) / input_range - offset = output_min - scale * input_min - offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] - - return {'scale': scale, 'offset': offset}, stat - - - def get_other_param_info(stat): - example = stat['max'] - scale = np.ones_like(example) - offset = np.zeros_like(example) - info = { - 'max': np.ones_like(example), - 'min': np.full_like(example, -1), - 'mean': np.zeros_like(example), - 'std': np.ones_like(example) - } - return {'scale': scale, 'offset': offset}, info - - pos_param, pos_info = get_pos_param_info(result['pos']) - other_param, other_info = get_other_param_info(result['other']) - - param = dict_apply_reduce( - [pos_param, other_param], - lambda x: np.concatenate(x,axis=-1)) - info = dict_apply_reduce( - [pos_info, other_info], - lambda x: np.concatenate(x,axis=-1)) - - return SingleFieldLinearNormalizer.create_manual( - scale=param['scale'], - offset=param['offset'], - input_stats_dict=info - ) - - -def robomimic_abs_action_only_dual_arm_normalizer_from_stat(stat): - Da = stat['max'].shape[-1] - Dah = Da // 2 - result = dict_apply_split( - stat, lambda x: { - 'pos0': x[...,:3], - 'other0': x[...,3:Dah], - 'pos1': x[...,Dah:Dah+3], - 'other1': x[...,Dah+3:] - }) - - def get_pos_param_info(stat, output_max=1, output_min=-1, range_eps=1e-7): - # -1, 1 normalization - input_max = stat['max'] - input_min = stat['min'] - input_range = input_max - input_min - ignore_dim = input_range < range_eps - input_range[ignore_dim] = output_max - output_min - scale = (output_max - output_min) / input_range - offset = output_min - scale * input_min - offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] - - return {'scale': scale, 'offset': offset}, stat - - - def get_other_param_info(stat): - example = stat['max'] - scale = np.ones_like(example) - offset = np.zeros_like(example) - info = { - 'max': np.ones_like(example), - 'min': np.full_like(example, -1), - 'mean': np.zeros_like(example), - 'std': np.ones_like(example) - } - return {'scale': scale, 'offset': offset}, info - - pos0_param, pos0_info = get_pos_param_info(result['pos0']) - pos1_param, pos1_info = get_pos_param_info(result['pos1']) - other0_param, other0_info = get_other_param_info(result['other0']) - other1_param, other1_info = get_other_param_info(result['other1']) - - param = dict_apply_reduce( - [pos0_param, other0_param, pos1_param, other1_param], - lambda x: np.concatenate(x,axis=-1)) - info = dict_apply_reduce( - [pos0_info, other0_info, pos1_info, other1_info], - lambda x: np.concatenate(x,axis=-1)) - - return SingleFieldLinearNormalizer.create_manual( - scale=param['scale'], - offset=param['offset'], - input_stats_dict=info - ) - - -def array_to_stats(arr: np.ndarray): - stat = { - 'min': np.min(arr, axis=0), - 'max': np.max(arr, axis=0), - 'mean': np.mean(arr, axis=0), - 'std': np.std(arr, axis=0) - } - return stat diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pose_trajectory_interpolator.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pose_trajectory_interpolator.py deleted file mode 100644 index 3e37397e2..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pose_trajectory_interpolator.py +++ /dev/null @@ -1,208 +0,0 @@ -from typing import Union -import numbers -import numpy as np -import scipy.interpolate as si -import scipy.spatial.transform as st - -def rotation_distance(a: st.Rotation, b: st.Rotation) -> float: - return (b * a.inv()).magnitude() - -def pose_distance(start_pose, end_pose): - start_pose = np.array(start_pose) - end_pose = np.array(end_pose) - start_pos = start_pose[:3] - end_pos = end_pose[:3] - start_rot = st.Rotation.from_rotvec(start_pose[3:]) - end_rot = st.Rotation.from_rotvec(end_pose[3:]) - pos_dist = np.linalg.norm(end_pos - start_pos) - rot_dist = rotation_distance(start_rot, end_rot) - return pos_dist, rot_dist - -class PoseTrajectoryInterpolator: - def __init__(self, times: np.ndarray, poses: np.ndarray): - assert len(times) >= 1 - assert len(poses) == len(times) - if not isinstance(times, np.ndarray): - times = np.array(times) - if not isinstance(poses, np.ndarray): - poses = np.array(poses) - - if len(times) == 1: - # special treatment for single step interpolation - self.single_step = True - self._times = times - self._poses = poses - else: - self.single_step = False - assert np.all(times[1:] >= times[:-1]) - - pos = poses[:,:3] - rot = st.Rotation.from_rotvec(poses[:,3:]) - - self.pos_interp = si.interp1d(times, pos, - axis=0, assume_sorted=True) - self.rot_interp = st.Slerp(times, rot) - - @property - def times(self) -> np.ndarray: - if self.single_step: - return self._times - else: - return self.pos_interp.x - - @property - def poses(self) -> np.ndarray: - if self.single_step: - return self._poses - else: - n = len(self.times) - poses = np.zeros((n, 6)) - poses[:,:3] = self.pos_interp.y - poses[:,3:] = self.rot_interp(self.times).as_rotvec() - return poses - - def trim(self, - start_t: float, end_t: float - ) -> "PoseTrajectoryInterpolator": - assert start_t <= end_t - times = self.times - should_keep = (start_t < times) & (times < end_t) - keep_times = times[should_keep] - all_times = np.concatenate([[start_t], keep_times, [end_t]]) - # remove duplicates, Slerp requires strictly increasing x - all_times = np.unique(all_times) - # interpolate - all_poses = self(all_times) - return PoseTrajectoryInterpolator(times=all_times, poses=all_poses) - - def drive_to_waypoint(self, - pose, time, curr_time, - max_pos_speed=np.inf, - max_rot_speed=np.inf - ) -> "PoseTrajectoryInterpolator": - assert(max_pos_speed > 0) - assert(max_rot_speed > 0) - time = max(time, curr_time) - - curr_pose = self(curr_time) - pos_dist, rot_dist = pose_distance(curr_pose, pose) - pos_min_duration = pos_dist / max_pos_speed - rot_min_duration = rot_dist / max_rot_speed - duration = time - curr_time - duration = max(duration, max(pos_min_duration, rot_min_duration)) - assert duration >= 0 - last_waypoint_time = curr_time + duration - - # insert new pose - trimmed_interp = self.trim(curr_time, curr_time) - times = np.append(trimmed_interp.times, [last_waypoint_time], axis=0) - poses = np.append(trimmed_interp.poses, [pose], axis=0) - - # create new interpolator - final_interp = PoseTrajectoryInterpolator(times, poses) - return final_interp - - def schedule_waypoint(self, - pose, time, - max_pos_speed=np.inf, - max_rot_speed=np.inf, - curr_time=None, - last_waypoint_time=None - ) -> "PoseTrajectoryInterpolator": - assert(max_pos_speed > 0) - assert(max_rot_speed > 0) - if last_waypoint_time is not None: - assert curr_time is not None - - # trim current interpolator to between curr_time and last_waypoint_time - start_time = self.times[0] - end_time = self.times[-1] - assert start_time <= end_time - - if curr_time is not None: - if time <= curr_time: - # if insert time is earlier than current time - # no effect should be done to the interpolator - return self - # now, curr_time < time - start_time = max(curr_time, start_time) - - if last_waypoint_time is not None: - # if last_waypoint_time is earlier than start_time - # use start_time - if time <= last_waypoint_time: - end_time = curr_time - else: - end_time = max(last_waypoint_time, curr_time) - else: - end_time = curr_time - - end_time = min(end_time, time) - start_time = min(start_time, end_time) - # end time should be the latest of all times except time - # after this we can assume order (proven by zhenjia, due to the 2 min operations) - - # Constraints: - # start_time <= end_time <= time (proven by zhenjia) - # curr_time <= start_time (proven by zhenjia) - # curr_time <= time (proven by zhenjia) - - # time can't change - # last_waypoint_time can't change - # curr_time can't change - assert start_time <= end_time - assert end_time <= time - if last_waypoint_time is not None: - if time <= last_waypoint_time: - assert end_time == curr_time - else: - assert end_time == max(last_waypoint_time, curr_time) - - if curr_time is not None: - assert curr_time <= start_time - assert curr_time <= time - - trimmed_interp = self.trim(start_time, end_time) - # after this, all waypoints in trimmed_interp is within start_time and end_time - # and is earlier than time - - # determine speed - duration = time - end_time - end_pose = trimmed_interp(end_time) - pos_dist, rot_dist = pose_distance(pose, end_pose) - pos_min_duration = pos_dist / max_pos_speed - rot_min_duration = rot_dist / max_rot_speed - duration = max(duration, max(pos_min_duration, rot_min_duration)) - assert duration >= 0 - last_waypoint_time = end_time + duration - - # insert new pose - times = np.append(trimmed_interp.times, [last_waypoint_time], axis=0) - poses = np.append(trimmed_interp.poses, [pose], axis=0) - - # create new interpolator - final_interp = PoseTrajectoryInterpolator(times, poses) - return final_interp - - - def __call__(self, t: Union[numbers.Number, np.ndarray]) -> np.ndarray: - is_single = False - if isinstance(t, numbers.Number): - is_single = True - t = np.array([t]) - - pose = np.zeros((len(t), 6)) - if self.single_step: - pose[:] = self._poses[0] - else: - start_time = self.times[0] - end_time = self.times[-1] - t = np.clip(t, start_time, end_time) - - pose = np.zeros((len(t), 6)) - pose[:,:3] = self.pos_interp(t) - pose[:,3:] = self.rot_interp(t).as_rotvec() - - if is_single: - pose = pose[0] - return pose diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/precise_sleep.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/precise_sleep.py deleted file mode 100644 index 83ef62c5d..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/precise_sleep.py +++ /dev/null @@ -1,25 +0,0 @@ -import time - -def precise_sleep(dt: float, slack_time: float=0.001, time_func=time.monotonic): - """ - Use hybrid of time.sleep and spinning to minimize jitter. - Sleep dt - slack_time seconds first, then spin for the rest. - """ - t_start = time_func() - if dt > slack_time: - time.sleep(dt - slack_time) - t_end = t_start + dt - while time_func() < t_end: - pass - return - -def precise_wait(t_end: float, slack_time: float=0.001, time_func=time.monotonic): - t_start = time_func() - t_wait = t_end - t_start - if t_wait > 0: - t_sleep = t_wait - slack_time - if t_sleep > 0: - time.sleep(t_sleep) - while time_func() < t_end: - pass - return diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pymunk_override.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pymunk_override.py deleted file mode 100644 index 2439020a1..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pymunk_override.py +++ /dev/null @@ -1,248 +0,0 @@ -# ---------------------------------------------------------------------------- -# pymunk -# Copyright (c) 2007-2016 Victor Blomqvist -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# ---------------------------------------------------------------------------- - -"""This submodule contains helper functions to help with quick prototyping -using pymunk together with pygame. - -Intended to help with debugging and prototyping, not for actual production use -in a full application. The methods contained in this module is opinionated -about your coordinate system and not in any way optimized. -""" - -__docformat__ = "reStructuredText" - -__all__ = [ - "DrawOptions", - "get_mouse_pos", - "to_pygame", - "from_pygame", - "lighten", - "positive_y_is_up", -] - -from typing import List, Sequence, Tuple - -import pygame - -import numpy as np - -import pymunk -from pymunk.space_debug_draw_options import SpaceDebugColor -from pymunk.vec2d import Vec2d - -positive_y_is_up: bool = False -"""Make increasing values of y point upwards. - -When True:: - - y - ^ - | . (3, 3) - | - | . (2, 2) - | - +------ > x - -When False:: - - +------ > x - | - | . (2, 2) - | - | . (3, 3) - v - y - -""" - - -class DrawOptions(pymunk.SpaceDebugDrawOptions): - def __init__(self, surface: pygame.Surface) -> None: - """Draw a pymunk.Space on a pygame.Surface object. - - Typical usage:: - - >>> import pymunk - >>> surface = pygame.Surface((10,10)) - >>> space = pymunk.Space() - >>> options = pymunk.pygame_util.DrawOptions(surface) - >>> space.debug_draw(options) - - You can control the color of a shape by setting shape.color to the color - you want it drawn in:: - - >>> c = pymunk.Circle(None, 10) - >>> c.color = pygame.Color("pink") - - See pygame_util.demo.py for a full example - - Since pygame uses a coordinate system where y points down (in contrast - to many other cases), you either have to make the physics simulation - with Pymunk also behave in that way, or flip everything when you draw. - - The easiest is probably to just make the simulation behave the same - way as Pygame does. In that way all coordinates used are in the same - orientation and easy to reason about:: - - >>> space = pymunk.Space() - >>> space.gravity = (0, -1000) - >>> body = pymunk.Body() - >>> body.position = (0, 0) # will be positioned in the top left corner - >>> space.debug_draw(options) - - To flip the drawing its possible to set the module property - :py:data:`positive_y_is_up` to True. Then the pygame drawing will flip - the simulation upside down before drawing:: - - >>> positive_y_is_up = True - >>> body = pymunk.Body() - >>> body.position = (0, 0) - >>> # Body will be position in bottom left corner - - :Parameters: - surface : pygame.Surface - Surface that the objects will be drawn on - """ - self.surface = surface - super(DrawOptions, self).__init__() - - def draw_circle( - self, - pos: Vec2d, - angle: float, - radius: float, - outline_color: SpaceDebugColor, - fill_color: SpaceDebugColor, - ) -> None: - p = to_pygame(pos, self.surface) - - pygame.draw.circle(self.surface, fill_color.as_int(), p, round(radius), 0) - pygame.draw.circle(self.surface, light_color(fill_color).as_int(), p, round(radius-4), 0) - - circle_edge = pos + Vec2d(radius, 0).rotated(angle) - p2 = to_pygame(circle_edge, self.surface) - line_r = 2 if radius > 20 else 1 - # pygame.draw.lines(self.surface, outline_color.as_int(), False, [p, p2], line_r) - - def draw_segment(self, a: Vec2d, b: Vec2d, color: SpaceDebugColor) -> None: - p1 = to_pygame(a, self.surface) - p2 = to_pygame(b, self.surface) - - pygame.draw.aalines(self.surface, color.as_int(), False, [p1, p2]) - - def draw_fat_segment( - self, - a: Tuple[float, float], - b: Tuple[float, float], - radius: float, - outline_color: SpaceDebugColor, - fill_color: SpaceDebugColor, - ) -> None: - p1 = to_pygame(a, self.surface) - p2 = to_pygame(b, self.surface) - - r = round(max(1, radius * 2)) - pygame.draw.lines(self.surface, fill_color.as_int(), False, [p1, p2], r) - if r > 2: - orthog = [abs(p2[1] - p1[1]), abs(p2[0] - p1[0])] - if orthog[0] == 0 and orthog[1] == 0: - return - scale = radius / (orthog[0] * orthog[0] + orthog[1] * orthog[1]) ** 0.5 - orthog[0] = round(orthog[0] * scale) - orthog[1] = round(orthog[1] * scale) - points = [ - (p1[0] - orthog[0], p1[1] - orthog[1]), - (p1[0] + orthog[0], p1[1] + orthog[1]), - (p2[0] + orthog[0], p2[1] + orthog[1]), - (p2[0] - orthog[0], p2[1] - orthog[1]), - ] - pygame.draw.polygon(self.surface, fill_color.as_int(), points) - pygame.draw.circle( - self.surface, - fill_color.as_int(), - (round(p1[0]), round(p1[1])), - round(radius), - ) - pygame.draw.circle( - self.surface, - fill_color.as_int(), - (round(p2[0]), round(p2[1])), - round(radius), - ) - - def draw_polygon( - self, - verts: Sequence[Tuple[float, float]], - radius: float, - outline_color: SpaceDebugColor, - fill_color: SpaceDebugColor, - ) -> None: - ps = [to_pygame(v, self.surface) for v in verts] - ps += [ps[0]] - - radius = 2 - pygame.draw.polygon(self.surface, light_color(fill_color).as_int(), ps) - - if radius > 0: - for i in range(len(verts)): - a = verts[i] - b = verts[(i + 1) % len(verts)] - self.draw_fat_segment(a, b, radius, fill_color, fill_color) - - def draw_dot( - self, size: float, pos: Tuple[float, float], color: SpaceDebugColor - ) -> None: - p = to_pygame(pos, self.surface) - pygame.draw.circle(self.surface, color.as_int(), p, round(size), 0) - - -def get_mouse_pos(surface: pygame.Surface) -> Tuple[int, int]: - """Get position of the mouse pointer in pymunk coordinates.""" - p = pygame.mouse.get_pos() - return from_pygame(p, surface) - - -def to_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]: - """Convenience method to convert pymunk coordinates to pygame surface - local coordinates. - - Note that in case positive_y_is_up is False, this function won't actually do - anything except converting the point to integers. - """ - if positive_y_is_up: - return round(p[0]), surface.get_height() - round(p[1]) - else: - return round(p[0]), round(p[1]) - - -def from_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]: - """Convenience method to convert pygame surface local coordinates to - pymunk coordinates - """ - return to_pygame(p, surface) - - -def light_color(color: SpaceDebugColor): - color = np.minimum(1.2 * np.float32([color.r, color.g, color.b, color.a]), np.float32([255])) - color = SpaceDebugColor(r=color[0], g=color[1], b=color[2], a=color[3]) - return color diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pymunk_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pymunk_util.py deleted file mode 100644 index 424599485..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pymunk_util.py +++ /dev/null @@ -1,52 +0,0 @@ -import pygame -import pymunk -import pymunk.pygame_util -import numpy as np - -COLLTYPE_DEFAULT = 0 -COLLTYPE_MOUSE = 1 -COLLTYPE_BALL = 2 - -def get_body_type(static=False): - body_type = pymunk.Body.DYNAMIC - if static: - body_type = pymunk.Body.STATIC - return body_type - - -def create_rectangle(space, - pos_x,pos_y,width,height, - density=3,static=False): - body = pymunk.Body(body_type=get_body_type(static)) - body.position = (pos_x,pos_y) - shape = pymunk.Poly.create_box(body,(width,height)) - shape.density = density - space.add(body,shape) - return body, shape - - -def create_rectangle_bb(space, - left, bottom, right, top, - **kwargs): - pos_x = (left + right) / 2 - pos_y = (top + bottom) / 2 - height = top - bottom - width = right - left - return create_rectangle(space, pos_x, pos_y, width, height, **kwargs) - -def create_circle(space, pos_x, pos_y, radius, density=3, static=False): - body = pymunk.Body(body_type=get_body_type(static)) - body.position = (pos_x, pos_y) - shape = pymunk.Circle(body, radius=radius) - shape.density = density - shape.collision_type = COLLTYPE_BALL - space.add(body, shape) - return body, shape - -def get_body_state(body): - state = np.zeros(6, dtype=np.float32) - state[:2] = body.position - state[2] = body.angle - state[3:5] = body.velocity - state[5] = body.angular_velocity - return state diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pytorch_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pytorch_util.py deleted file mode 100644 index a9262d485..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/pytorch_util.py +++ /dev/null @@ -1,82 +0,0 @@ -from typing import Dict, Callable, List -import collections -import torch -import torch.nn as nn - -def dict_apply( - x: Dict[str, torch.Tensor], - func: Callable[[torch.Tensor], torch.Tensor] - ) -> Dict[str, torch.Tensor]: - result = dict() - for key, value in x.items(): - if isinstance(value, dict): - result[key] = dict_apply(value, func) - else: - result[key] = func(value) - return result - -def pad_remaining_dims(x, target): - assert x.shape == target.shape[:len(x.shape)] - return x.reshape(x.shape + (1,)*(len(target.shape) - len(x.shape))) - -def dict_apply_split( - x: Dict[str, torch.Tensor], - split_func: Callable[[torch.Tensor], Dict[str, torch.Tensor]] - ) -> Dict[str, torch.Tensor]: - results = collections.defaultdict(dict) - for key, value in x.items(): - result = split_func(value) - for k, v in result.items(): - results[k][key] = v - return results - -def dict_apply_reduce( - x: List[Dict[str, torch.Tensor]], - reduce_func: Callable[[List[torch.Tensor]], torch.Tensor] - ) -> Dict[str, torch.Tensor]: - result = dict() - for key in x[0].keys(): - result[key] = reduce_func([x_[key] for x_ in x]) - return result - - -def replace_submodules( - root_module: nn.Module, - predicate: Callable[[nn.Module], bool], - func: Callable[[nn.Module], nn.Module]) -> nn.Module: - """ - predicate: Return true if the module is to be replaced. - func: Return new module to use. - """ - if predicate(root_module): - return func(root_module) - - bn_list = [k.split('.') for k, m - in root_module.named_modules(remove_duplicate=True) - if predicate(m)] - for *parent, k in bn_list: - parent_module = root_module - if len(parent) > 0: - parent_module = root_module.get_submodule('.'.join(parent)) - if isinstance(parent_module, nn.Sequential): - src_module = parent_module[int(k)] - else: - src_module = getattr(parent_module, k) - tgt_module = func(src_module) - if isinstance(parent_module, nn.Sequential): - parent_module[int(k)] = tgt_module - else: - setattr(parent_module, k, tgt_module) - # verify that all BN are replaced - bn_list = [k.split('.') for k, m - in root_module.named_modules(remove_duplicate=True) - if predicate(m)] - assert len(bn_list) == 0 - return root_module - -def optimizer_to(optimizer, device): - for state in optimizer.state.values(): - for k, v in state.items(): - if isinstance(v, torch.Tensor): - state[k] = v.to(device=device) - return optimizer diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/replay_buffer.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/replay_buffer.py deleted file mode 100644 index 022a704ed..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/replay_buffer.py +++ /dev/null @@ -1,588 +0,0 @@ -from typing import Union, Dict, Optional -import os -import math -import numbers -import zarr -import numcodecs -import numpy as np -from functools import cached_property - -def check_chunks_compatible(chunks: tuple, shape: tuple): - assert len(shape) == len(chunks) - for c in chunks: - assert isinstance(c, numbers.Integral) - assert c > 0 - -def rechunk_recompress_array(group, name, - chunks=None, chunk_length=None, - compressor=None, tmp_key='_temp'): - old_arr = group[name] - if chunks is None: - if chunk_length is not None: - chunks = (chunk_length,) + old_arr.chunks[1:] - else: - chunks = old_arr.chunks - check_chunks_compatible(chunks, old_arr.shape) - - if compressor is None: - compressor = old_arr.compressor - - if (chunks == old_arr.chunks) and (compressor == old_arr.compressor): - # no change - return old_arr - - # rechunk recompress - group.move(name, tmp_key) - old_arr = group[tmp_key] - n_copied, n_skipped, n_bytes_copied = zarr.copy( - source=old_arr, - dest=group, - name=name, - chunks=chunks, - compressor=compressor, - ) - del group[tmp_key] - arr = group[name] - return arr - -def get_optimal_chunks(shape, dtype, - target_chunk_bytes=2e6, - max_chunk_length=None): - """ - Common shapes - T,D - T,N,D - T,H,W,C - T,N,H,W,C - """ - itemsize = np.dtype(dtype).itemsize - # reversed - rshape = list(shape[::-1]) - if max_chunk_length is not None: - rshape[-1] = int(max_chunk_length) - split_idx = len(shape)-1 - for i in range(len(shape)-1): - this_chunk_bytes = itemsize * np.prod(rshape[:i]) - next_chunk_bytes = itemsize * np.prod(rshape[:i+1]) - if this_chunk_bytes <= target_chunk_bytes \ - and next_chunk_bytes > target_chunk_bytes: - split_idx = i - - rchunks = rshape[:split_idx] - item_chunk_bytes = itemsize * np.prod(rshape[:split_idx]) - this_max_chunk_length = rshape[split_idx] - next_chunk_length = min(this_max_chunk_length, math.ceil( - target_chunk_bytes / item_chunk_bytes)) - rchunks.append(next_chunk_length) - len_diff = len(shape) - len(rchunks) - rchunks.extend([1] * len_diff) - chunks = tuple(rchunks[::-1]) - # print(np.prod(chunks) * itemsize / target_chunk_bytes) - return chunks - - -class ReplayBuffer: - """ - Zarr-based temporal datastructure. - Assumes first dimension to be time. Only chunk in time dimension. - """ - def __init__(self, - root: Union[zarr.Group, - Dict[str,dict]]): - """ - Dummy constructor. Use copy_from* and create_from* class methods instead. - """ - assert('data' in root) - assert('meta' in root) - assert('episode_ends' in root['meta']) - for key, value in root['data'].items(): - assert(value.shape[0] == root['meta']['episode_ends'][-1]) - self.root = root - - # ============= create constructors =============== - @classmethod - def create_empty_zarr(cls, storage=None, root=None): - if root is None: - if storage is None: - storage = zarr.MemoryStore() - root = zarr.group(store=storage) - data = root.require_group('data', overwrite=False) - meta = root.require_group('meta', overwrite=False) - if 'episode_ends' not in meta: - episode_ends = meta.zeros('episode_ends', shape=(0,), dtype=np.int64, - compressor=None, overwrite=False) - return cls(root=root) - - @classmethod - def create_empty_numpy(cls): - root = { - 'data': dict(), - 'meta': { - 'episode_ends': np.zeros((0,), dtype=np.int64) - } - } - return cls(root=root) - - @classmethod - def create_from_group(cls, group, **kwargs): - if 'data' not in group: - # create from stratch - buffer = cls.create_empty_zarr(root=group, **kwargs) - else: - # already exist - buffer = cls(root=group, **kwargs) - return buffer - - @classmethod - def create_from_path(cls, zarr_path, mode='r', **kwargs): - """ - Open a on-disk zarr directly (for dataset larger than memory). - Slower. - """ - group = zarr.open(os.path.expanduser(zarr_path), mode) - return cls.create_from_group(group, **kwargs) - - # ============= copy constructors =============== - @classmethod - def copy_from_store(cls, src_store, store=None, keys=None, - chunks: Dict[str,tuple]=dict(), - compressors: Union[dict, str, numcodecs.abc.Codec]=dict(), - if_exists='replace', - **kwargs): - """ - Load to memory. - """ - src_root = zarr.group(src_store) - root = None - if store is None: - # numpy backend - meta = dict() - for key, value in src_root['meta'].items(): - if len(value.shape) == 0: - meta[key] = np.array(value) - else: - meta[key] = value[:] - - if keys is None: - keys = src_root['data'].keys() - data = dict() - for key in keys: - arr = src_root['data'][key] - data[key] = arr[:] - - root = { - 'meta': meta, - 'data': data - } - else: - root = zarr.group(store=store) - # copy without recompression - n_copied, n_skipped, n_bytes_copied = zarr.copy_store(source=src_store, dest=store, - source_path='/meta', dest_path='/meta', if_exists=if_exists) - data_group = root.create_group('data', overwrite=True) - if keys is None: - keys = src_root['data'].keys() - for key in keys: - value = src_root['data'][key] - cks = cls._resolve_array_chunks( - chunks=chunks, key=key, array=value) - cpr = cls._resolve_array_compressor( - compressors=compressors, key=key, array=value) - if cks == value.chunks and cpr == value.compressor: - # copy without recompression - this_path = '/data/' + key - n_copied, n_skipped, n_bytes_copied = zarr.copy_store( - source=src_store, dest=store, - source_path=this_path, dest_path=this_path, - if_exists=if_exists - ) - else: - # copy with recompression - n_copied, n_skipped, n_bytes_copied = zarr.copy( - source=value, dest=data_group, name=key, - chunks=cks, compressor=cpr, if_exists=if_exists - ) - buffer = cls(root=root) - return buffer - - @classmethod - def copy_from_path(cls, zarr_path, backend=None, store=None, keys=None, - chunks: Dict[str,tuple]=dict(), - compressors: Union[dict, str, numcodecs.abc.Codec]=dict(), - if_exists='replace', - **kwargs): - """ - Copy a on-disk zarr to in-memory compressed. - Recommended - """ - if backend == 'numpy': - print('backend argument is deprecated!') - store = None - group = zarr.open(os.path.expanduser(zarr_path), 'r') - return cls.copy_from_store(src_store=group.store, store=store, - keys=keys, chunks=chunks, compressors=compressors, - if_exists=if_exists, **kwargs) - - # ============= save methods =============== - def save_to_store(self, store, - chunks: Optional[Dict[str,tuple]]=dict(), - compressors: Union[str, numcodecs.abc.Codec, dict]=dict(), - if_exists='replace', - **kwargs): - - root = zarr.group(store) - if self.backend == 'zarr': - # recompression free copy - n_copied, n_skipped, n_bytes_copied = zarr.copy_store( - source=self.root.store, dest=store, - source_path='/meta', dest_path='/meta', if_exists=if_exists) - else: - meta_group = root.create_group('meta', overwrite=True) - # save meta, no chunking - for key, value in self.root['meta'].items(): - _ = meta_group.array( - name=key, - data=value, - shape=value.shape, - chunks=value.shape) - - # save data, chunk - data_group = root.create_group('data', overwrite=True) - for key, value in self.root['data'].items(): - cks = self._resolve_array_chunks( - chunks=chunks, key=key, array=value) - cpr = self._resolve_array_compressor( - compressors=compressors, key=key, array=value) - if isinstance(value, zarr.Array): - if cks == value.chunks and cpr == value.compressor: - # copy without recompression - this_path = '/data/' + key - n_copied, n_skipped, n_bytes_copied = zarr.copy_store( - source=self.root.store, dest=store, - source_path=this_path, dest_path=this_path, if_exists=if_exists) - else: - # copy with recompression - n_copied, n_skipped, n_bytes_copied = zarr.copy( - source=value, dest=data_group, name=key, - chunks=cks, compressor=cpr, if_exists=if_exists - ) - else: - # numpy - _ = data_group.array( - name=key, - data=value, - chunks=cks, - compressor=cpr - ) - return store - - def save_to_path(self, zarr_path, - chunks: Optional[Dict[str,tuple]]=dict(), - compressors: Union[str, numcodecs.abc.Codec, dict]=dict(), - if_exists='replace', - **kwargs): - store = zarr.DirectoryStore(os.path.expanduser(zarr_path)) - return self.save_to_store(store, chunks=chunks, - compressors=compressors, if_exists=if_exists, **kwargs) - - @staticmethod - def resolve_compressor(compressor='default'): - if compressor == 'default': - compressor = numcodecs.Blosc(cname='lz4', clevel=5, - shuffle=numcodecs.Blosc.NOSHUFFLE) - elif compressor == 'disk': - compressor = numcodecs.Blosc('zstd', clevel=5, - shuffle=numcodecs.Blosc.BITSHUFFLE) - return compressor - - @classmethod - def _resolve_array_compressor(cls, - compressors: Union[dict, str, numcodecs.abc.Codec], key, array): - # allows compressor to be explicitly set to None - cpr = 'nil' - if isinstance(compressors, dict): - if key in compressors: - cpr = cls.resolve_compressor(compressors[key]) - elif isinstance(array, zarr.Array): - cpr = array.compressor - else: - cpr = cls.resolve_compressor(compressors) - # backup default - if cpr == 'nil': - cpr = cls.resolve_compressor('default') - return cpr - - @classmethod - def _resolve_array_chunks(cls, - chunks: Union[dict, tuple], key, array): - cks = None - if isinstance(chunks, dict): - if key in chunks: - cks = chunks[key] - elif isinstance(array, zarr.Array): - cks = array.chunks - elif isinstance(chunks, tuple): - cks = chunks - else: - raise TypeError(f"Unsupported chunks type {type(chunks)}") - # backup default - if cks is None: - cks = get_optimal_chunks(shape=array.shape, dtype=array.dtype) - # check - check_chunks_compatible(chunks=cks, shape=array.shape) - return cks - - # ============= properties ================= - @cached_property - def data(self): - return self.root['data'] - - @cached_property - def meta(self): - return self.root['meta'] - - def update_meta(self, data): - # sanitize data - np_data = dict() - for key, value in data.items(): - if isinstance(value, np.ndarray): - np_data[key] = value - else: - arr = np.array(value) - if arr.dtype == object: - raise TypeError(f"Invalid value type {type(value)}") - np_data[key] = arr - - meta_group = self.meta - if self.backend == 'zarr': - for key, value in np_data.items(): - _ = meta_group.array( - name=key, - data=value, - shape=value.shape, - chunks=value.shape, - overwrite=True) - else: - meta_group.update(np_data) - - return meta_group - - @property - def episode_ends(self): - return self.meta['episode_ends'] - - def get_episode_idxs(self): - import numba - numba.jit(nopython=True) - def _get_episode_idxs(episode_ends): - result = np.zeros((episode_ends[-1],), dtype=np.int64) - for i in range(len(episode_ends)): - start = 0 - if i > 0: - start = episode_ends[i-1] - end = episode_ends[i] - for idx in range(start, end): - result[idx] = i - return result - return _get_episode_idxs(self.episode_ends) - - - @property - def backend(self): - backend = 'numpy' - if isinstance(self.root, zarr.Group): - backend = 'zarr' - return backend - - # =========== dict-like API ============== - def __repr__(self) -> str: - if self.backend == 'zarr': - return str(self.root.tree()) - else: - return super().__repr__() - - def keys(self): - return self.data.keys() - - def values(self): - return self.data.values() - - def items(self): - return self.data.items() - - def __getitem__(self, key): - return self.data[key] - - def __contains__(self, key): - return key in self.data - - # =========== our API ============== - @property - def n_steps(self): - if len(self.episode_ends) == 0: - return 0 - return self.episode_ends[-1] - - @property - def n_episodes(self): - return len(self.episode_ends) - - @property - def chunk_size(self): - if self.backend == 'zarr': - return next(iter(self.data.arrays()))[-1].chunks[0] - return None - - @property - def episode_lengths(self): - ends = self.episode_ends[:] - ends = np.insert(ends, 0, 0) - lengths = np.diff(ends) - return lengths - - def add_episode(self, - data: Dict[str, np.ndarray], - chunks: Optional[Dict[str,tuple]]=dict(), - compressors: Union[str, numcodecs.abc.Codec, dict]=dict()): - assert(len(data) > 0) - is_zarr = (self.backend == 'zarr') - - curr_len = self.n_steps - episode_length = None - for key, value in data.items(): - assert(len(value.shape) >= 1) - if episode_length is None: - episode_length = len(value) - else: - assert(episode_length == len(value)) - new_len = curr_len + episode_length - - for key, value in data.items(): - new_shape = (new_len,) + value.shape[1:] - # create array - if key not in self.data: - if is_zarr: - cks = self._resolve_array_chunks( - chunks=chunks, key=key, array=value) - cpr = self._resolve_array_compressor( - compressors=compressors, key=key, array=value) - arr = self.data.zeros(name=key, - shape=new_shape, - chunks=cks, - dtype=value.dtype, - compressor=cpr) - else: - # copy data to prevent modify - arr = np.zeros(shape=new_shape, dtype=value.dtype) - self.data[key] = arr - else: - arr = self.data[key] - assert(value.shape[1:] == arr.shape[1:]) - # same method for both zarr and numpy - if is_zarr: - arr.resize(new_shape) - else: - arr.resize(new_shape, refcheck=False) - # copy data - arr[-value.shape[0]:] = value - - # append to episode ends - episode_ends = self.episode_ends - if is_zarr: - episode_ends.resize(episode_ends.shape[0] + 1) - else: - episode_ends.resize(episode_ends.shape[0] + 1, refcheck=False) - episode_ends[-1] = new_len - - # rechunk - if is_zarr: - if episode_ends.chunks[0] < episode_ends.shape[0]: - rechunk_recompress_array(self.meta, 'episode_ends', - chunk_length=int(episode_ends.shape[0] * 1.5)) - - def drop_episode(self): - is_zarr = (self.backend == 'zarr') - episode_ends = self.episode_ends[:].copy() - assert(len(episode_ends) > 0) - start_idx = 0 - if len(episode_ends) > 1: - start_idx = episode_ends[-2] - for key, value in self.data.items(): - new_shape = (start_idx,) + value.shape[1:] - if is_zarr: - value.resize(new_shape) - else: - value.resize(new_shape, refcheck=False) - if is_zarr: - self.episode_ends.resize(len(episode_ends)-1) - else: - self.episode_ends.resize(len(episode_ends)-1, refcheck=False) - - def pop_episode(self): - assert(self.n_episodes > 0) - episode = self.get_episode(self.n_episodes-1, copy=True) - self.drop_episode() - return episode - - def extend(self, data): - self.add_episode(data) - - def get_episode(self, idx, copy=False): - idx = list(range(len(self.episode_ends)))[idx] - start_idx = 0 - if idx > 0: - start_idx = self.episode_ends[idx-1] - end_idx = self.episode_ends[idx] - result = self.get_steps_slice(start_idx, end_idx, copy=copy) - return result - - def get_episode_slice(self, idx): - start_idx = 0 - if idx > 0: - start_idx = self.episode_ends[idx-1] - end_idx = self.episode_ends[idx] - return slice(start_idx, end_idx) - - def get_steps_slice(self, start, stop, step=None, copy=False): - _slice = slice(start, stop, step) - - result = dict() - for key, value in self.data.items(): - x = value[_slice] - if copy and isinstance(value, np.ndarray): - x = x.copy() - result[key] = x - return result - - # =========== chunking ============= - def get_chunks(self) -> dict: - assert self.backend == 'zarr' - chunks = dict() - for key, value in self.data.items(): - chunks[key] = value.chunks - return chunks - - def set_chunks(self, chunks: dict): - assert self.backend == 'zarr' - for key, value in chunks.items(): - if key in self.data: - arr = self.data[key] - if value != arr.chunks: - check_chunks_compatible(chunks=value, shape=arr.shape) - rechunk_recompress_array(self.data, key, chunks=value) - - def get_compressors(self) -> dict: - assert self.backend == 'zarr' - compressors = dict() - for key, value in self.data.items(): - compressors[key] = value.compressor - return compressors - - def set_compressors(self, compressors: dict): - assert self.backend == 'zarr' - for key, value in compressors.items(): - if key in self.data: - arr = self.data[key] - compressor = self.resolve_compressor(value) - if compressor != arr.compressor: - rechunk_recompress_array(self.data, key, compressor=compressor) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/robomimic_config_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/robomimic_config_util.py deleted file mode 100644 index 85768bbbf..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/robomimic_config_util.py +++ /dev/null @@ -1,47 +0,0 @@ -from omegaconf import OmegaConf -from robomimic.config import config_factory -import robomimic.scripts.generate_paper_configs as gpc -from robomimic.scripts.generate_paper_configs import ( - modify_config_for_default_image_exp, - modify_config_for_default_low_dim_exp, - modify_config_for_dataset, -) - -def get_robomimic_config( - algo_name='bc_rnn', - hdf5_type='low_dim', - task_name='square', - dataset_type='ph' - ): - base_dataset_dir = '/tmp/null' - filter_key = None - - # decide whether to use low-dim or image training defaults - modifier_for_obs = modify_config_for_default_image_exp - if hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"]: - modifier_for_obs = modify_config_for_default_low_dim_exp - - algo_config_name = "bc" if algo_name == "bc_rnn" else algo_name - config = config_factory(algo_name=algo_config_name) - # turn into default config for observation modalities (e.g.: low-dim or rgb) - config = modifier_for_obs(config) - # add in config based on the dataset - config = modify_config_for_dataset( - config=config, - task_name=task_name, - dataset_type=dataset_type, - hdf5_type=hdf5_type, - base_dataset_dir=base_dataset_dir, - filter_key=filter_key, - ) - # add in algo hypers based on dataset - algo_config_modifier = getattr(gpc, f'modify_{algo_name}_config_for_dataset') - config = algo_config_modifier( - config=config, - task_name=task_name, - dataset_type=dataset_type, - hdf5_type=hdf5_type, - ) - return config - - diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/robomimic_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/robomimic_util.py deleted file mode 100644 index 24f25fbc5..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/robomimic_util.py +++ /dev/null @@ -1,177 +0,0 @@ -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 - - -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 __len__(self): - return len(self.file['data']) - - def convert_actions(self, - states: np.ndarray, - actions: np.ndarray) -> np.ndarray: - """ - Given state and delta action sequence - generate equivalent goal position and orientation for each step - keep the original gripper action intact. - """ - # 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,7) - - env = self.env - # 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_gripper = stacked_actions[...,[-1]] - for i in range(len(states)): - _ = 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_gripper - ], axis=-1) - abs_actions = stacked_abs_actions.reshape(actions.shape) - return abs_actions - - def convert_idx(self, idx): - file = self.file - demo = file[f'data/demo_{idx}'] - # input - states = demo['states'][:] - actions = demo['actions'][:] - - # generate abs actions - abs_actions = self.convert_actions(states, actions) - return abs_actions - - def convert_and_eval_idx(self, idx): - 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[f'data/demo_{idx}'] - # 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 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/sampler.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/sampler.py deleted file mode 100644 index 4accaa093..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/sampler.py +++ /dev/null @@ -1,153 +0,0 @@ -from typing import Optional -import numpy as np -import numba -from diffusion_policy.common.replay_buffer import ReplayBuffer - - -@numba.jit(nopython=True) -def create_indices( - episode_ends:np.ndarray, sequence_length:int, - episode_mask: np.ndarray, - pad_before: int=0, pad_after: int=0, - debug:bool=True) -> np.ndarray: - episode_mask.shape == episode_ends.shape - pad_before = min(max(pad_before, 0), sequence_length-1) - pad_after = min(max(pad_after, 0), sequence_length-1) - - indices = list() - for i in range(len(episode_ends)): - if not episode_mask[i]: - # skip episode - continue - start_idx = 0 - if i > 0: - start_idx = episode_ends[i-1] - end_idx = episode_ends[i] - episode_length = end_idx - start_idx - - min_start = -pad_before - max_start = episode_length - sequence_length + pad_after - - # range stops one idx before end - for idx in range(min_start, max_start+1): - buffer_start_idx = max(idx, 0) + start_idx - buffer_end_idx = min(idx+sequence_length, episode_length) + start_idx - start_offset = buffer_start_idx - (idx+start_idx) - end_offset = (idx+sequence_length+start_idx) - buffer_end_idx - sample_start_idx = 0 + start_offset - sample_end_idx = sequence_length - end_offset - if debug: - assert(start_offset >= 0) - assert(end_offset >= 0) - assert (sample_end_idx - sample_start_idx) == (buffer_end_idx - buffer_start_idx) - indices.append([ - buffer_start_idx, buffer_end_idx, - sample_start_idx, sample_end_idx]) - indices = np.array(indices) - return indices - - -def get_val_mask(n_episodes, val_ratio, seed=0): - val_mask = np.zeros(n_episodes, dtype=bool) - if val_ratio <= 0: - return val_mask - - # have at least 1 episode for validation, and at least 1 episode for train - n_val = min(max(1, round(n_episodes * val_ratio)), n_episodes-1) - rng = np.random.default_rng(seed=seed) - val_idxs = rng.choice(n_episodes, size=n_val, replace=False) - val_mask[val_idxs] = True - return val_mask - - -def downsample_mask(mask, max_n, seed=0): - # subsample training data - train_mask = mask - if (max_n is not None) and (np.sum(train_mask) > max_n): - n_train = int(max_n) - curr_train_idxs = np.nonzero(train_mask)[0] - rng = np.random.default_rng(seed=seed) - train_idxs_idx = rng.choice(len(curr_train_idxs), size=n_train, replace=False) - train_idxs = curr_train_idxs[train_idxs_idx] - train_mask = np.zeros_like(train_mask) - train_mask[train_idxs] = True - assert np.sum(train_mask) == n_train - return train_mask - -class SequenceSampler: - def __init__(self, - replay_buffer: ReplayBuffer, - sequence_length:int, - pad_before:int=0, - pad_after:int=0, - keys=None, - key_first_k=dict(), - episode_mask: Optional[np.ndarray]=None, - ): - """ - key_first_k: dict str: int - Only take first k data from these keys (to improve perf) - """ - - super().__init__() - assert(sequence_length >= 1) - if keys is None: - keys = list(replay_buffer.keys()) - - episode_ends = replay_buffer.episode_ends[:] - if episode_mask is None: - episode_mask = np.ones(episode_ends.shape, dtype=bool) - - if np.any(episode_mask): - indices = create_indices(episode_ends, - sequence_length=sequence_length, - pad_before=pad_before, - pad_after=pad_after, - episode_mask=episode_mask - ) - else: - indices = np.zeros((0,4), dtype=np.int64) - - # (buffer_start_idx, buffer_end_idx, sample_start_idx, sample_end_idx) - self.indices = indices - self.keys = list(keys) # prevent OmegaConf list performance problem - self.sequence_length = sequence_length - self.replay_buffer = replay_buffer - self.key_first_k = key_first_k - - def __len__(self): - return len(self.indices) - - def sample_sequence(self, idx): - buffer_start_idx, buffer_end_idx, sample_start_idx, sample_end_idx \ - = self.indices[idx] - result = dict() - for key in self.keys: - input_arr = self.replay_buffer[key] - # performance optimization, avoid small allocation if possible - if key not in self.key_first_k: - sample = input_arr[buffer_start_idx:buffer_end_idx] - else: - # performance optimization, only load used obs steps - n_data = buffer_end_idx - buffer_start_idx - k_data = min(self.key_first_k[key], n_data) - # fill value with Nan to catch bugs - # the non-loaded region should never be used - sample = np.full((n_data,) + input_arr.shape[1:], - fill_value=np.nan, dtype=input_arr.dtype) - try: - sample[:k_data] = input_arr[buffer_start_idx:buffer_start_idx+k_data] - except Exception as e: - import pdb; pdb.set_trace() - data = sample - if (sample_start_idx > 0) or (sample_end_idx < self.sequence_length): - data = np.zeros( - shape=(self.sequence_length,) + input_arr.shape[1:], - dtype=input_arr.dtype) - if sample_start_idx > 0: - data[:sample_start_idx] = sample[0] - if sample_end_idx < self.sequence_length: - data[sample_end_idx:] = sample[-1] - data[sample_start_idx:sample_end_idx] = sample - result[key] = data - return result diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/timestamp_accumulator.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/timestamp_accumulator.py deleted file mode 100644 index 2e7267288..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/common/timestamp_accumulator.py +++ /dev/null @@ -1,222 +0,0 @@ -from typing import List, Tuple, Optional, Dict -import math -import numpy as np - - -def get_accumulate_timestamp_idxs( - timestamps: List[float], - start_time: float, - dt: float, - eps:float=1e-5, - next_global_idx: Optional[int]=0, - allow_negative=False - ) -> Tuple[List[int], List[int], int]: - """ - For each dt window, choose the first timestamp in the window. - Assumes timestamps sorted. One timestamp might be chosen multiple times due to dropped frames. - next_global_idx should start at 0 normally, and then use the returned next_global_idx. - However, when overwiting previous values are desired, set last_global_idx to None. - - Returns: - local_idxs: which index in the given timestamps array to chose from - global_idxs: the global index of each chosen timestamp - next_global_idx: used for next call. - """ - local_idxs = list() - global_idxs = list() - for local_idx, ts in enumerate(timestamps): - # add eps * dt to timestamps so that when ts == start_time + k * dt - # is always recorded as kth element (avoiding floating point errors) - global_idx = math.floor((ts - start_time) / dt + eps) - if (not allow_negative) and (global_idx < 0): - continue - if next_global_idx is None: - next_global_idx = global_idx - - n_repeats = max(0, global_idx - next_global_idx + 1) - for i in range(n_repeats): - local_idxs.append(local_idx) - global_idxs.append(next_global_idx + i) - next_global_idx += n_repeats - return local_idxs, global_idxs, next_global_idx - - -def align_timestamps( - timestamps: List[float], - target_global_idxs: List[int], - start_time: float, - dt: float, - eps:float=1e-5): - if isinstance(target_global_idxs, np.ndarray): - target_global_idxs = target_global_idxs.tolist() - assert len(target_global_idxs) > 0 - - local_idxs, global_idxs, _ = get_accumulate_timestamp_idxs( - timestamps=timestamps, - start_time=start_time, - dt=dt, - eps=eps, - next_global_idx=target_global_idxs[0], - allow_negative=True - ) - if len(global_idxs) > len(target_global_idxs): - # if more steps available, truncate - global_idxs = global_idxs[:len(target_global_idxs)] - local_idxs = local_idxs[:len(target_global_idxs)] - - if len(global_idxs) == 0: - import pdb; pdb.set_trace() - - for i in range(len(target_global_idxs) - len(global_idxs)): - # if missing, repeat - local_idxs.append(len(timestamps)-1) - global_idxs.append(global_idxs[-1] + 1) - assert global_idxs == target_global_idxs - assert len(local_idxs) == len(global_idxs) - return local_idxs - - -class TimestampObsAccumulator: - def __init__(self, - start_time: float, - dt: float, - eps: float=1e-5): - self.start_time = start_time - self.dt = dt - self.eps = eps - self.obs_buffer = dict() - self.timestamp_buffer = None - self.next_global_idx = 0 - - def __len__(self): - return self.next_global_idx - - @property - def data(self): - if self.timestamp_buffer is None: - return dict() - result = dict() - for key, value in self.obs_buffer.items(): - result[key] = value[:len(self)] - return result - - @property - def actual_timestamps(self): - if self.timestamp_buffer is None: - return np.array([]) - return self.timestamp_buffer[:len(self)] - - @property - def timestamps(self): - if self.timestamp_buffer is None: - return np.array([]) - return self.start_time + np.arange(len(self)) * self.dt - - def put(self, data: Dict[str, np.ndarray], timestamps: np.ndarray): - """ - data: - key: T,* - """ - - local_idxs, global_idxs, self.next_global_idx = get_accumulate_timestamp_idxs( - timestamps=timestamps, - start_time=self.start_time, - dt=self.dt, - eps=self.eps, - next_global_idx=self.next_global_idx - ) - - if len(global_idxs) > 0: - if self.timestamp_buffer is None: - # first allocation - self.obs_buffer = dict() - for key, value in data.items(): - self.obs_buffer[key] = np.zeros_like(value) - self.timestamp_buffer = np.zeros( - (len(timestamps),), dtype=np.float64) - - this_max_size = global_idxs[-1] + 1 - if this_max_size > len(self.timestamp_buffer): - # reallocate - new_size = max(this_max_size, len(self.timestamp_buffer) * 2) - for key in list(self.obs_buffer.keys()): - new_shape = (new_size,) + self.obs_buffer[key].shape[1:] - self.obs_buffer[key] = np.resize(self.obs_buffer[key], new_shape) - self.timestamp_buffer = np.resize(self.timestamp_buffer, (new_size)) - - # write data - for key, value in self.obs_buffer.items(): - value[global_idxs] = data[key][local_idxs] - self.timestamp_buffer[global_idxs] = timestamps[local_idxs] - - -class TimestampActionAccumulator: - def __init__(self, - start_time: float, - dt: float, - eps: float=1e-5): - """ - Different from Obs accumulator, the action accumulator - allows overwriting previous values. - """ - self.start_time = start_time - self.dt = dt - self.eps = eps - self.action_buffer = None - self.timestamp_buffer = None - self.size = 0 - - def __len__(self): - return self.size - - @property - def actions(self): - if self.action_buffer is None: - return np.array([]) - return self.action_buffer[:len(self)] - - @property - def actual_timestamps(self): - if self.timestamp_buffer is None: - return np.array([]) - return self.timestamp_buffer[:len(self)] - - @property - def timestamps(self): - if self.timestamp_buffer is None: - return np.array([]) - return self.start_time + np.arange(len(self)) * self.dt - - def put(self, actions: np.ndarray, timestamps: np.ndarray): - """ - Note: timestamps is the time when the action will be issued, - not when the action will be completed (target_timestamp) - """ - - local_idxs, global_idxs, _ = get_accumulate_timestamp_idxs( - timestamps=timestamps, - start_time=self.start_time, - dt=self.dt, - eps=self.eps, - # allows overwriting previous actions - next_global_idx=None - ) - - if len(global_idxs) > 0: - if self.timestamp_buffer is None: - # first allocation - self.action_buffer = np.zeros_like(actions) - self.timestamp_buffer = np.zeros((len(actions),), dtype=np.float64) - - this_max_size = global_idxs[-1] + 1 - if this_max_size > len(self.timestamp_buffer): - # reallocate - new_size = max(this_max_size, len(self.timestamp_buffer) * 2) - new_shape = (new_size,) + self.action_buffer.shape[1:] - self.action_buffer = np.resize(self.action_buffer, new_shape) - self.timestamp_buffer = np.resize(self.timestamp_buffer, (new_size,)) - - # potentially rewrite old data (as expected) - self.action_buffer[global_idxs] = actions[local_idxs] - self.timestamp_buffer[global_idxs] = timestamps[local_idxs] - self.size = max(self.size, this_max_size) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/blockpush_lowdim_seed.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/blockpush_lowdim_seed.yaml deleted file mode 100644 index aa1e636c4..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/blockpush_lowdim_seed.yaml +++ /dev/null @@ -1,34 +0,0 @@ -name: blockpush_lowdim_seed - -obs_dim: 16 -action_dim: 2 -keypoint_dim: 2 -obs_eef_target: True - -env_runner: - _target_: diffusion_policy.env_runner.blockpush_lowdim_runner.BlockPushLowdimRunner - n_train: 6 - n_train_vis: 2 - train_start_seed: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - max_steps: 350 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - fps: 5 - past_action: ${past_action_visible} - abs_action: False - obs_eef_target: ${task.obs_eef_target} - n_envs: null - -dataset: - _target_: diffusion_policy.dataset.blockpush_lowdim_dataset.BlockPushLowdimDataset - zarr_path: data/block_pushing/multimodal_push_seed.zarr - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_eef_target: ${task.obs_eef_target} - use_manual_normalizer: False - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/blockpush_lowdim_seed_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/blockpush_lowdim_seed_abs.yaml deleted file mode 100644 index 0649121ac..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/blockpush_lowdim_seed_abs.yaml +++ /dev/null @@ -1,34 +0,0 @@ -name: blockpush_lowdim_seed_abs - -obs_dim: 16 -action_dim: 2 -keypoint_dim: 2 -obs_eef_target: True - -env_runner: - _target_: diffusion_policy.env_runner.blockpush_lowdim_runner.BlockPushLowdimRunner - n_train: 6 - n_train_vis: 2 - train_start_seed: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - max_steps: 350 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - fps: 5 - past_action: ${past_action_visible} - abs_action: True - obs_eef_target: ${task.obs_eef_target} - n_envs: null - -dataset: - _target_: diffusion_policy.dataset.blockpush_lowdim_dataset.BlockPushLowdimDataset - zarr_path: data/block_pushing/multimodal_push_seed_abs.zarr - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_eef_target: ${task.obs_eef_target} - use_manual_normalizer: False - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_image.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_image.yaml deleted file mode 100644 index 10158bbd0..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_image.yaml +++ /dev/null @@ -1,64 +0,0 @@ -name: can_image - -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - agentview_image: - shape: [3, 84, 84] - type: rgb - robot0_eye_in_hand_image: - shape: [3, 84, 84] - type: rgb - robot0_eef_pos: - shape: [3] - # type default: low_dim - robot0_eef_quat: - shape: [4] - robot0_gripper_qpos: - shape: [2] - action: - shape: [7] - -task_name: &task_name can -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image.hdf5 -abs_action: &abs_action False - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner - dataset_path: *dataset_path - shape_meta: *shape_meta - # costs 1GB per env - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_obs_key: 'agentview_image' - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset - shape_meta: *shape_meta - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - n_obs_steps: ${dataset_obs_steps} - abs_action: *abs_action - rotation_rep: 'rotation_6d' - use_legacy_normalizer: False - use_cache: True - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_image_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_image_abs.yaml deleted file mode 100644 index eee34dc63..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_image_abs.yaml +++ /dev/null @@ -1,64 +0,0 @@ -name: can_image - -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - agentview_image: - shape: [3, 84, 84] - type: rgb - robot0_eye_in_hand_image: - shape: [3, 84, 84] - type: rgb - robot0_eef_pos: - shape: [3] - # type default: low_dim - robot0_eef_quat: - shape: [4] - robot0_gripper_qpos: - shape: [2] - action: - shape: [10] - -task_name: &task_name can -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5 -abs_action: &abs_action True - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner - dataset_path: *dataset_path - shape_meta: *shape_meta - # costs 1GB per env - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_obs_key: 'agentview_image' - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset - shape_meta: *shape_meta - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - n_obs_steps: ${dataset_obs_steps} - abs_action: *abs_action - rotation_rep: 'rotation_6d' - use_legacy_normalizer: False - use_cache: True - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_lowdim.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_lowdim.yaml deleted file mode 100644 index ba8d20f69..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_lowdim.yaml +++ /dev/null @@ -1,45 +0,0 @@ -name: can_lowdim - -obs_dim: 23 -action_dim: 7 -keypoint_dim: 3 - -obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos'] -task_name: &task_name can -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim.hdf5 -abs_action: &abs_action False - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner - dataset_path: *dataset_path - obs_keys: *obs_keys - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - n_latency_steps: ${n_latency_steps} - render_hw: [128,128] - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - n_envs: 28 - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_keys: *obs_keys - abs_action: *abs_action - use_legacy_normalizer: False - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_lowdim_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_lowdim_abs.yaml deleted file mode 100644 index a9f723bd9..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/can_lowdim_abs.yaml +++ /dev/null @@ -1,46 +0,0 @@ -name: can_lowdim - -obs_dim: 23 -action_dim: 10 -keypoint_dim: 3 - -obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos'] -task_name: &task_name can -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim_abs.hdf5 -abs_action: &abs_action True - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner - dataset_path: *dataset_path - obs_keys: *obs_keys - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - n_latency_steps: ${n_latency_steps} - render_hw: [128,128] - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - n_envs: 28 - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_keys: *obs_keys - abs_action: *abs_action - use_legacy_normalizer: False - rotation_rep: rotation_6d - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/kitchen_lowdim.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/kitchen_lowdim.yaml deleted file mode 100644 index a57d5d833..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/kitchen_lowdim.yaml +++ /dev/null @@ -1,33 +0,0 @@ -name: kitchen_lowdim - -obs_dim: 60 -action_dim: 9 -keypoint_dim: 3 - -dataset_dir: &dataset_dir data/kitchen - -env_runner: - _target_: diffusion_policy.env_runner.kitchen_lowdim_runner.KitchenLowdimRunner - dataset_dir: *dataset_dir - n_train: 6 - n_train_vis: 2 - train_start_seed: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - max_steps: 280 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_hw: [240, 360] - fps: 12.5 - past_action: ${past_action_visible} - n_envs: null - -dataset: - _target_: diffusion_policy.dataset.kitchen_lowdim_dataset.KitchenLowdimDataset - dataset_dir: *dataset_dir - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1'} - pad_after: ${eval:'${n_action_steps}-1'} - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/kitchen_lowdim_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/kitchen_lowdim_abs.yaml deleted file mode 100644 index ecd0691b9..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/kitchen_lowdim_abs.yaml +++ /dev/null @@ -1,38 +0,0 @@ -name: kitchen_lowdim - -obs_dim: 60 -action_dim: 9 -keypoint_dim: 3 - -abs_action: True -robot_noise_ratio: 0.1 - -env_runner: - _target_: diffusion_policy.env_runner.kitchen_lowdim_runner.KitchenLowdimRunner - dataset_dir: data/kitchen - n_train: 6 - n_train_vis: 2 - train_start_seed: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - max_steps: 280 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_hw: [240, 360] - fps: 12.5 - past_action: ${past_action_visible} - abs_action: ${task.abs_action} - robot_noise_ratio: ${task.robot_noise_ratio} - n_envs: null - -dataset: - _target_: diffusion_policy.dataset.kitchen_mjl_lowdim_dataset.KitchenMjlLowdimDataset - dataset_dir: data/kitchen/kitchen_demos_multitask - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1'} - pad_after: ${eval:'${n_action_steps}-1'} - abs_action: ${task.abs_action} - robot_noise_ratio: ${task.robot_noise_ratio} - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_image.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_image.yaml deleted file mode 100644 index 8ddde4567..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_image.yaml +++ /dev/null @@ -1,64 +0,0 @@ -name: lift_image - -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - agentview_image: - shape: [3, 84, 84] - type: rgb - robot0_eye_in_hand_image: - shape: [3, 84, 84] - type: rgb - robot0_eef_pos: - shape: [3] - # type default: low_dim - robot0_eef_quat: - shape: [4] - robot0_gripper_qpos: - shape: [2] - action: - shape: [7] - -task_name: &task_name lift -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image.hdf5 -abs_action: &abs_action False - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner - dataset_path: *dataset_path - shape_meta: *shape_meta - # costs 1GB per env - n_train: 6 - n_train_vis: 1 - train_start_idx: 0 - n_test: 50 - n_test_vis: 3 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_obs_key: 'agentview_image' - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset - shape_meta: *shape_meta - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - n_obs_steps: ${dataset_obs_steps} - abs_action: *abs_action - rotation_rep: 'rotation_6d' - use_legacy_normalizer: False - use_cache: True - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_image_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_image_abs.yaml deleted file mode 100644 index b25002f38..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_image_abs.yaml +++ /dev/null @@ -1,63 +0,0 @@ -name: lift_image - -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - agentview_image: - shape: [3, 84, 84] - type: rgb - robot0_eye_in_hand_image: - shape: [3, 84, 84] - type: rgb - robot0_eef_pos: - shape: [3] - # type default: low_dim - robot0_eef_quat: - shape: [4] - robot0_gripper_qpos: - shape: [2] - action: - shape: [10] - -task_name: &task_name lift -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5 -abs_action: &abs_action True - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner - dataset_path: *dataset_path - shape_meta: *shape_meta - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_obs_key: 'agentview_image' - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset - shape_meta: *shape_meta - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - n_obs_steps: ${dataset_obs_steps} - abs_action: *abs_action - rotation_rep: 'rotation_6d' - use_legacy_normalizer: False - use_cache: True - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_lowdim.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_lowdim.yaml deleted file mode 100644 index 5b1b2d25a..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_lowdim.yaml +++ /dev/null @@ -1,46 +0,0 @@ -name: lift_lowdim - -obs_dim: 19 -action_dim: 7 -keypoint_dim: 3 - -obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos'] -task_name: &task_name lift -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim.hdf5 -abs_action: &abs_action False - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner - dataset_path: *dataset_path - obs_keys: *obs_keys - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - n_latency_steps: ${n_latency_steps} - render_hw: [128,128] - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_keys: *obs_keys - abs_action: *abs_action - use_legacy_normalizer: False - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_lowdim_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_lowdim_abs.yaml deleted file mode 100644 index fe9d17b5b..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/lift_lowdim_abs.yaml +++ /dev/null @@ -1,47 +0,0 @@ -name: lift_lowdim - -obs_dim: 19 -action_dim: 10 -keypoint_dim: 3 - -obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos'] -task_name: &task_name lift -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim_abs.hdf5 -abs_action: &abs_action True - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner - dataset_path: *dataset_path - obs_keys: *obs_keys - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 3 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - n_latency_steps: ${n_latency_steps} - render_hw: [128,128] - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_keys: *obs_keys - abs_action: *abs_action - use_legacy_normalizer: False - rotation_rep: rotation_6d - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/pusht_image.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/pusht_image.yaml deleted file mode 100644 index 36b5dd312..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/pusht_image.yaml +++ /dev/null @@ -1,40 +0,0 @@ -name: pusht_image - -image_shape: &image_shape [3, 96, 96] -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - image: - shape: *image_shape - type: rgb - agent_pos: - shape: [2] - type: low_dim - action: - shape: [2] - -env_runner: - _target_: diffusion_policy.env_runner.pusht_image_runner.PushTImageRunner - n_train: 6 - n_train_vis: 2 - train_start_seed: 0 - n_test: 50 - n_test_vis: 4 - legacy_test: True - test_start_seed: 100000 - max_steps: 300 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - fps: 10 - past_action: ${past_action_visible} - n_envs: null - -dataset: - _target_: diffusion_policy.dataset.pusht_image_dataset.PushTImageDataset - zarr_path: data/pusht/pusht_cchi_v7_replay.zarr - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1'} - pad_after: ${eval:'${n_action_steps}-1'} - seed: 42 - val_ratio: 0.02 - max_train_episodes: 90 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/pusht_lowdim.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/pusht_lowdim.yaml deleted file mode 100644 index 3ec7eb564..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/pusht_lowdim.yaml +++ /dev/null @@ -1,34 +0,0 @@ -name: pusht_lowdim - -obs_dim: 20 # 9*2 keypoints + 2 state -action_dim: 2 -keypoint_dim: 2 - -env_runner: - _target_: diffusion_policy.env_runner.pusht_keypoints_runner.PushTKeypointsRunner - keypoint_visible_rate: ${keypoint_visible_rate} - n_train: 6 - n_train_vis: 2 - train_start_seed: 0 - n_test: 50 - n_test_vis: 4 - legacy_test: True - test_start_seed: 100000 - max_steps: 300 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - n_latency_steps: ${n_latency_steps} - fps: 10 - agent_keypoints: False - past_action: ${past_action_visible} - n_envs: null - -dataset: - _target_: diffusion_policy.dataset.pusht_dataset.PushTLowdimDataset - zarr_path: data/pusht/pusht_cchi_v7_replay.zarr - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - seed: 42 - val_ratio: 0.02 - max_train_episodes: 90 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/real_pusht_image.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/real_pusht_image.yaml deleted file mode 100644 index a3f7c3f08..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/real_pusht_image.yaml +++ /dev/null @@ -1,47 +0,0 @@ -name: real_image - -image_shape: [3, 240, 320] -dataset_path: data/pusht_real/real_pusht_20230105 - -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - # camera_0: - # shape: ${task.image_shape} - # type: rgb - camera_1: - shape: ${task.image_shape} - type: rgb - # camera_2: - # shape: ${task.image_shape} - # type: rgb - camera_3: - shape: ${task.image_shape} - type: rgb - # camera_4: - # shape: ${task.image_shape} - # type: rgb - robot_eef_pose: - shape: [2] - type: low_dim - action: - shape: [2] - -env_runner: - _target_: diffusion_policy.env_runner.real_pusht_image_runner.RealPushTImageRunner - -dataset: - _target_: diffusion_policy.dataset.real_pusht_image_dataset.RealPushTImageDataset - shape_meta: *shape_meta - dataset_path: ${task.dataset_path} - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - n_obs_steps: ${dataset_obs_steps} - n_latency_steps: ${n_latency_steps} - use_cache: True - seed: 42 - val_ratio: 0.00 - max_train_episodes: null - delta_action: False - diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_image.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_image.yaml deleted file mode 100644 index 827151be6..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_image.yaml +++ /dev/null @@ -1,64 +0,0 @@ -name: square_image - -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - agentview_image: - shape: [3, 84, 84] - type: rgb - robot0_eye_in_hand_image: - shape: [3, 84, 84] - type: rgb - robot0_eef_pos: - shape: [3] - # type default: low_dim - robot0_eef_quat: - shape: [4] - robot0_gripper_qpos: - shape: [2] - action: - shape: [7] - -task_name: &task_name square -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image.hdf5 -abs_action: &abs_action False - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner - dataset_path: *dataset_path - shape_meta: *shape_meta - # costs 1GB per env - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_obs_key: 'agentview_image' - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset - shape_meta: *shape_meta - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - n_obs_steps: ${dataset_obs_steps} - abs_action: *abs_action - rotation_rep: 'rotation_6d' - use_legacy_normalizer: False - use_cache: True - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_image_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_image_abs.yaml deleted file mode 100644 index d27a916af..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_image_abs.yaml +++ /dev/null @@ -1,64 +0,0 @@ -name: square_image - -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - agentview_image: - shape: [3, 84, 84] - type: rgb - robot0_eye_in_hand_image: - shape: [3, 84, 84] - type: rgb - robot0_eef_pos: - shape: [3] - # type default: low_dim - robot0_eef_quat: - shape: [4] - robot0_gripper_qpos: - shape: [2] - action: - shape: [10] - -task_name: &task_name square -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5 -abs_action: &abs_action True - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner - dataset_path: *dataset_path - shape_meta: *shape_meta - # costs 1GB per env - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_obs_key: 'agentview_image' - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset - shape_meta: *shape_meta - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - n_obs_steps: ${dataset_obs_steps} - abs_action: *abs_action - rotation_rep: 'rotation_6d' - use_legacy_normalizer: False - use_cache: True - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_lowdim.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_lowdim.yaml deleted file mode 100644 index f93e062cd..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_lowdim.yaml +++ /dev/null @@ -1,46 +0,0 @@ -name: square_lowdim - -obs_dim: 23 -action_dim: 7 -keypoint_dim: 3 - -obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos'] -task_name: &task_name square -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim.hdf5 -abs_action: &abs_action False - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner - dataset_path: *dataset_path - obs_keys: *obs_keys - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - n_latency_steps: ${n_latency_steps} - render_hw: [128,128] - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - n_envs: 28 - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_keys: *obs_keys - abs_action: *abs_action - use_legacy_normalizer: False - seed: 42 - val_ratio: 0.02 - max_train_episodes: null diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_lowdim_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_lowdim_abs.yaml deleted file mode 100644 index f69b3c8d8..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/square_lowdim_abs.yaml +++ /dev/null @@ -1,47 +0,0 @@ -name: square_lowdim - -obs_dim: 23 -action_dim: 10 -keypoint_dim: 3 - -obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos'] -task_name: &task_name square -dataset_type: &dataset_type ph -abs_action: &abs_action True -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim_abs.hdf5 - - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner - dataset_path: *dataset_path - obs_keys: *obs_keys - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - # use python's eval function as resolver, single-quoted string as argument - max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - n_latency_steps: ${n_latency_steps} - render_hw: [128,128] - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - n_envs: 28 - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_keys: *obs_keys - abs_action: *abs_action - use_legacy_normalizer: False - seed: 42 - val_ratio: 0.02 - max_train_episodes: null diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_image.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_image.yaml deleted file mode 100644 index 8a8b298f4..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_image.yaml +++ /dev/null @@ -1,63 +0,0 @@ -name: tool_hang_image - -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - sideview_image: - shape: [3, 240, 240] - type: rgb - robot0_eye_in_hand_image: - shape: [3, 240, 240] - type: rgb - robot0_eef_pos: - shape: [3] - # type default: low_dim - robot0_eef_quat: - shape: [4] - robot0_gripper_qpos: - shape: [2] - action: - shape: [7] - -task_name: &task_name tool_hang -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image.hdf5 -abs_action: &abs_action False - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner - dataset_path: *dataset_path - shape_meta: *shape_meta - # costs 1GB per env - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - max_steps: 700 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_obs_key: 'sideview_image' - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset - shape_meta: *shape_meta - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - n_obs_steps: ${dataset_obs_steps} - abs_action: *abs_action - rotation_rep: 'rotation_6d' - use_legacy_normalizer: False - use_cache: True - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_image_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_image_abs.yaml deleted file mode 100644 index 068bdc93a..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_image_abs.yaml +++ /dev/null @@ -1,63 +0,0 @@ -name: tool_hang_image_abs - -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - sideview_image: - shape: [3, 240, 240] - type: rgb - robot0_eye_in_hand_image: - shape: [3, 240, 240] - type: rgb - robot0_eef_pos: - shape: [3] - # type default: low_dim - robot0_eef_quat: - shape: [4] - robot0_gripper_qpos: - shape: [2] - action: - shape: [10] - -task_name: &task_name tool_hang -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5 -abs_action: &abs_action True - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner - dataset_path: *dataset_path - shape_meta: *shape_meta - # costs 1GB per env - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - max_steps: 700 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_obs_key: 'sideview_image' - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset - shape_meta: *shape_meta - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - n_obs_steps: ${dataset_obs_steps} - abs_action: *abs_action - rotation_rep: 'rotation_6d' - use_legacy_normalizer: False - use_cache: True - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_lowdim.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_lowdim.yaml deleted file mode 100644 index 38264391f..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_lowdim.yaml +++ /dev/null @@ -1,45 +0,0 @@ -name: tool_hang_lowdim - -obs_dim: 53 -action_dim: 7 -keypoint_dim: 3 - -obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos'] -task_name: &task_name tool_hang -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim.hdf5 -abs_action: &abs_action False - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner - dataset_path: *dataset_path - obs_keys: *obs_keys - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - max_steps: 700 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - n_latency_steps: ${n_latency_steps} - render_hw: [128,128] - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - n_envs: 28 -# seed 42 will crash MuJoCo for some reason. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_keys: *obs_keys - abs_action: *abs_action - use_legacy_normalizer: False - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_lowdim_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_lowdim_abs.yaml deleted file mode 100644 index bfdb34daf..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/tool_hang_lowdim_abs.yaml +++ /dev/null @@ -1,46 +0,0 @@ -name: tool_hang_lowdim - -obs_dim: 53 -action_dim: 10 -keypoint_dim: 3 - -obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos'] -task_name: &task_name tool_hang -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim_abs.hdf5 -abs_action: &abs_action True - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner - dataset_path: *dataset_path - obs_keys: *obs_keys - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - max_steps: 700 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - n_latency_steps: ${n_latency_steps} - render_hw: [128,128] - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - n_envs: 28 -# seed 42 will crash MuJoCo for some reason. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_keys: *obs_keys - abs_action: *abs_action - use_legacy_normalizer: False - rotation_rep: rotation_6d - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_image.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_image.yaml deleted file mode 100644 index c9c24cc05..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_image.yaml +++ /dev/null @@ -1,75 +0,0 @@ -name: transport_image - -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - shouldercamera0_image: - shape: [3, 84, 84] - type: rgb - robot0_eye_in_hand_image: - shape: [3, 84, 84] - type: rgb - robot0_eef_pos: - shape: [3] - # type default: low_dim - robot0_eef_quat: - shape: [4] - robot0_gripper_qpos: - shape: [2] - shouldercamera1_image: - shape: [3, 84, 84] - type: rgb - robot1_eye_in_hand_image: - shape: [3, 84, 84] - type: rgb - robot1_eef_pos: - shape: [3] - # type default: low_dim - robot1_eef_quat: - shape: [4] - robot1_gripper_qpos: - shape: [2] - action: - shape: [14] - -task_name: &task_name transport -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image.hdf5 -abs_action: &abs_action False - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner - dataset_path: *dataset_path - shape_meta: *shape_meta - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - max_steps: 700 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_obs_key: 'shouldercamera0_image' - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset - shape_meta: *shape_meta - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - n_obs_steps: ${dataset_obs_steps} - abs_action: *abs_action - rotation_rep: 'rotation_6d' - use_legacy_normalizer: False - use_cache: True - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_image_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_image_abs.yaml deleted file mode 100644 index 74e76c602..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_image_abs.yaml +++ /dev/null @@ -1,75 +0,0 @@ -name: transport_image - -shape_meta: &shape_meta - # acceptable types: rgb, low_dim - obs: - shouldercamera0_image: - shape: [3, 84, 84] - type: rgb - robot0_eye_in_hand_image: - shape: [3, 84, 84] - type: rgb - robot0_eef_pos: - shape: [3] - # type default: low_dim - robot0_eef_quat: - shape: [4] - robot0_gripper_qpos: - shape: [2] - shouldercamera1_image: - shape: [3, 84, 84] - type: rgb - robot1_eye_in_hand_image: - shape: [3, 84, 84] - type: rgb - robot1_eef_pos: - shape: [3] - # type default: low_dim - robot1_eef_quat: - shape: [4] - robot1_gripper_qpos: - shape: [2] - action: - shape: [20] - -task_name: &task_name transport -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5 -abs_action: &abs_action True - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner - dataset_path: *dataset_path - shape_meta: *shape_meta - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - max_steps: 700 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - render_obs_key: 'shouldercamera0_image' - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - tqdm_interval_sec: 1.0 - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset - shape_meta: *shape_meta - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - n_obs_steps: ${dataset_obs_steps} - abs_action: *abs_action - rotation_rep: 'rotation_6d' - use_legacy_normalizer: False - use_cache: True - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_lowdim.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_lowdim.yaml deleted file mode 100644 index 99e279cc8..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_lowdim.yaml +++ /dev/null @@ -1,49 +0,0 @@ -name: transport_lowdim - -obs_dim: 59 # 41+(3+4+2)*2 -action_dim: 14 # 7*2 -keypoint_dim: 3 - -obs_keys: &obs_keys [ - 'object', - 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos', - 'robot1_eef_pos', 'robot1_eef_quat', 'robot1_gripper_qpos' -] -task_name: &task_name transport -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim.hdf5 -abs_action: &abs_action False - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner - dataset_path: *dataset_path - obs_keys: *obs_keys - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 5 - test_start_seed: 100000 - max_steps: 700 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - n_latency_steps: ${n_latency_steps} - render_hw: [128,128] - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_keys: *obs_keys - abs_action: *abs_action - use_legacy_normalizer: False - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_lowdim_abs.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_lowdim_abs.yaml deleted file mode 100644 index 7cc30149d..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/task/transport_lowdim_abs.yaml +++ /dev/null @@ -1,49 +0,0 @@ -name: transport_lowdim - -obs_dim: 59 # 41+(3+4+2)*2 -action_dim: 20 # 10*2 -keypoint_dim: 3 - -obs_keys: &obs_keys [ - 'object', - 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos', - 'robot1_eef_pos', 'robot1_eef_quat', 'robot1_gripper_qpos' -] -task_name: &task_name transport -dataset_type: &dataset_type ph -dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim_abs.hdf5 -abs_action: &abs_action True - -env_runner: - _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner - dataset_path: *dataset_path - obs_keys: *obs_keys - n_train: 6 - n_train_vis: 2 - train_start_idx: 0 - n_test: 50 - n_test_vis: 4 - test_start_seed: 100000 - max_steps: 700 - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - n_latency_steps: ${n_latency_steps} - render_hw: [128,128] - fps: 10 - crf: 22 - past_action: ${past_action_visible} - abs_action: *abs_action - n_envs: 28 -# evaluation at this config requires a 16 core 64GB instance. - -dataset: - _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset - dataset_path: *dataset_path - horizon: ${horizon} - pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'} - pad_after: ${eval:'${n_action_steps}-1'} - obs_keys: *obs_keys - abs_action: *abs_action - use_legacy_normalizer: False - seed: 42 - val_ratio: 0.02 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_bet_lowdim_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_bet_lowdim_workspace.yaml deleted file mode 100644 index 63a3436d8..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_bet_lowdim_workspace.yaml +++ /dev/null @@ -1,132 +0,0 @@ -defaults: - - _self_ - - task: blockpush_lowdim_seed - -name: train_bet_lowdim -_target_: diffusion_policy.workspace.train_bet_lowdim_workspace.TrainBETLowdimWorkspace - -obs_dim: ${task.obs_dim} -action_dim: ${task.action_dim} -keypoint_dim: ${task.keypoint_dim} -task_name: ${task.name} -exp_name: "default" - -horizon: 3 -n_obs_steps: 3 -n_action_steps: 1 -n_latency_steps: 0 -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_local_cond: False -obs_as_global_cond: False -pred_action_steps_only: False - -policy: - _target_: diffusion_policy.policy.bet_lowdim_policy.BETLowdimPolicy - - action_ae: - _target_: diffusion_policy.model.bet.action_ae.discretizers.k_means.KMeansDiscretizer - num_bins: 24 - action_dim: ${action_dim} - predict_offsets: True - - obs_encoding_net: - _target_: torch.nn.Identity - output_dim: ${obs_dim} - - state_prior: - _target_: diffusion_policy.model.bet.latent_generators.mingpt.MinGPT - - discrete_input: false - input_dim: ${obs_dim} - - vocab_size: ${policy.action_ae.num_bins} - - # Architecture details - n_layer: 4 - n_head: 4 - n_embd: 72 - - block_size: ${horizon} # Length of history/context - predict_offsets: True - offset_loss_scale: 1000.0 # actions are very small - focal_loss_gamma: 2.0 - action_dim: ${action_dim} - - horizon: ${horizon} - n_obs_steps: ${n_obs_steps} - n_action_steps: ${n_action_steps} - -dataloader: - batch_size: 256 - num_workers: 1 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 256 - num_workers: 1 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - learning_rate: 0.0001 # 1e-4 - weight_decay: 0.1 - betas: [0.9, 0.95] - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 5000 - gradient_accumulate_every: 1 - grad_norm_clip: 1.0 - enable_normalizer: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_hybrid_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_hybrid_workspace.yaml deleted file mode 100644 index 0a900395c..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_hybrid_workspace.yaml +++ /dev/null @@ -1,141 +0,0 @@ -defaults: - - _self_ - - task: lift_image_abs - -name: train_diffusion_transformer_hybrid -_target_: diffusion_policy.workspace.train_diffusion_transformer_hybrid_workspace.TrainDiffusionTransformerHybridWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 10 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -dataset_obs_steps: ${n_obs_steps} -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_cond: True - -policy: - _target_: diffusion_policy.policy.diffusion_transformer_hybrid_image_policy.DiffusionTransformerHybridImagePolicy - - shape_meta: ${shape_meta} - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - beta_schedule: squaredcos_cap_v2 - variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan - clip_sample: True # required when predict_epsilon=False - prediction_type: epsilon # or sample - - horizon: ${horizon} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 100 - - crop_shape: [76, 76] - obs_encoder_group_norm: True - eval_fixed_crop: True - - n_layer: 8 - n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP - n_head: 4 - n_emb: 256 - p_drop_emb: 0.0 - p_drop_attn: 0.3 - causal_attn: True - time_as_cond: True # if false, use BERT like encoder only arch, time as input - obs_as_cond: ${obs_as_cond} - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 64 - num_workers: 8 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 64 - num_workers: 8 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - transformer_weight_decay: 1.0e-3 - obs_encoder_weight_decay: 1.0e-6 - learning_rate: 1.0e-4 - betas: [0.9, 0.95] - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - # Transformer needs LR warmup - lr_warmup_steps: 1000 - num_epochs: 3050 - gradient_accumulate_every: 1 - # EMA destroys performance when used with BatchNorm - # replace BatchNorm with GroupNorm. - use_ema: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_lowdim_kitchen_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_lowdim_kitchen_workspace.yaml deleted file mode 100644 index 937f2301c..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_lowdim_kitchen_workspace.yaml +++ /dev/null @@ -1,146 +0,0 @@ -defaults: - - _self_ - - task: kitchen_lowdim_abs - -name: train_diffusion_transformer_lowdim -_target_: diffusion_policy.workspace.train_diffusion_transformer_lowdim_workspace.TrainDiffusionTransformerLowdimWorkspace - -obs_dim: ${task.obs_dim} -action_dim: ${task.action_dim} -task_name: ${task.name} -exp_name: "default" - -horizon: 16 -n_obs_steps: 4 -n_action_steps: 8 -n_latency_steps: 0 -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_cond: True -pred_action_steps_only: False - -policy: - _target_: diffusion_policy.policy.diffusion_transformer_lowdim_policy.DiffusionTransformerLowdimPolicy - - model: - _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion - input_dim: ${eval:'${action_dim} if ${obs_as_cond} else ${obs_dim} + ${action_dim}'} - output_dim: ${policy.model.input_dim} - horizon: ${horizon} - n_obs_steps: ${n_obs_steps} - cond_dim: ${eval:'${obs_dim} if ${obs_as_cond} else 0'} - - n_layer: 8 - n_head: 4 - n_emb: 768 - p_drop_emb: 0.0 - p_drop_attn: 0.1 - - causal_attn: True - time_as_cond: True # if false, use BERT like encoder only arch, time as input - obs_as_cond: ${obs_as_cond} - n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - beta_schedule: squaredcos_cap_v2 - variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan - clip_sample: True # required when predict_epsilon=False - prediction_type: epsilon # or sample - - horizon: ${horizon} - obs_dim: ${obs_dim} - action_dim: ${action_dim} - n_action_steps: ${n_action_steps} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 100 - obs_as_cond: ${obs_as_cond} - pred_action_steps_only: ${pred_action_steps_only} - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 256 - num_workers: 1 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 256 - num_workers: 1 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - learning_rate: 1.0e-4 - weight_decay: 1.0e-3 - betas: [0.9, 0.95] - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - # Transformer needs LR warmup - lr_warmup_steps: 1000 - num_epochs: 5000 - gradient_accumulate_every: 1 - use_ema: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_lowdim_pusht_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_lowdim_pusht_workspace.yaml deleted file mode 100644 index ffedd8635..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_lowdim_pusht_workspace.yaml +++ /dev/null @@ -1,146 +0,0 @@ -defaults: - - _self_ - - task: pusht_lowdim - -name: train_diffusion_transformer_lowdim -_target_: diffusion_policy.workspace.train_diffusion_transformer_lowdim_workspace.TrainDiffusionTransformerLowdimWorkspace - -obs_dim: ${task.obs_dim} -action_dim: ${task.action_dim} -task_name: ${task.name} -exp_name: "default" - -horizon: 16 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_cond: True -pred_action_steps_only: False - -policy: - _target_: diffusion_policy.policy.diffusion_transformer_lowdim_policy.DiffusionTransformerLowdimPolicy - - model: - _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion - input_dim: ${eval:'${action_dim} if ${obs_as_cond} else ${obs_dim} + ${action_dim}'} - output_dim: ${policy.model.input_dim} - horizon: ${horizon} - n_obs_steps: ${n_obs_steps} - cond_dim: ${eval:'${obs_dim} if ${obs_as_cond} else 0'} - - n_layer: 8 - n_head: 4 - n_emb: 256 - p_drop_emb: 0.0 - p_drop_attn: 0.01 - - causal_attn: True - time_as_cond: True # if false, use BERT like encoder only arch, time as input - obs_as_cond: ${obs_as_cond} - n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - beta_schedule: squaredcos_cap_v2 - variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan - clip_sample: True # required when predict_epsilon=False - prediction_type: epsilon # or sample - - horizon: ${horizon} - obs_dim: ${obs_dim} - action_dim: ${action_dim} - n_action_steps: ${n_action_steps} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 100 - obs_as_cond: ${obs_as_cond} - pred_action_steps_only: ${pred_action_steps_only} - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 256 - num_workers: 1 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 256 - num_workers: 1 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - learning_rate: 1.0e-4 - weight_decay: 1.0e-1 - betas: [0.9, 0.95] - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - # Transformer needs LR warmup - lr_warmup_steps: 1000 - num_epochs: 8000 - gradient_accumulate_every: 1 - use_ema: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_lowdim_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_lowdim_workspace.yaml deleted file mode 100644 index 5f91d370a..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_lowdim_workspace.yaml +++ /dev/null @@ -1,146 +0,0 @@ -defaults: - - _self_ - - task: blockpush_lowdim_seed - -name: train_diffusion_transformer_lowdim -_target_: diffusion_policy.workspace.train_diffusion_transformer_lowdim_workspace.TrainDiffusionTransformerLowdimWorkspace - -obs_dim: ${task.obs_dim} -action_dim: ${task.action_dim} -task_name: ${task.name} -exp_name: "default" - -horizon: 5 -n_obs_steps: 3 -n_action_steps: 1 -n_latency_steps: 0 -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_cond: True -pred_action_steps_only: False - -policy: - _target_: diffusion_policy.policy.diffusion_transformer_lowdim_policy.DiffusionTransformerLowdimPolicy - - model: - _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion - input_dim: ${eval:'${action_dim} if ${obs_as_cond} else ${obs_dim} + ${action_dim}'} - output_dim: ${policy.model.input_dim} - horizon: ${horizon} - n_obs_steps: ${n_obs_steps} - cond_dim: ${eval:'${obs_dim} if ${obs_as_cond} else 0'} - - n_layer: 8 - n_head: 4 - n_emb: 256 - p_drop_emb: 0.0 - p_drop_attn: 0.3 - - causal_attn: True - time_as_cond: True # if false, use BERT like encoder only arch, time as input - obs_as_cond: ${obs_as_cond} - n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - beta_schedule: squaredcos_cap_v2 - variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan - clip_sample: True # required when predict_epsilon=False - prediction_type: epsilon # or sample - - horizon: ${horizon} - obs_dim: ${obs_dim} - action_dim: ${action_dim} - n_action_steps: ${n_action_steps} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 100 - obs_as_cond: ${obs_as_cond} - pred_action_steps_only: ${pred_action_steps_only} - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 256 - num_workers: 1 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 256 - num_workers: 1 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - learning_rate: 1.0e-4 - weight_decay: 1.0e-3 - betas: [0.9, 0.95] - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - # Transformer needs LR warmup - lr_warmup_steps: 1000 - num_epochs: 5000 - gradient_accumulate_every: 1 - use_ema: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_real_hybrid_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_real_hybrid_workspace.yaml deleted file mode 100644 index 69bfad0f3..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_transformer_real_hybrid_workspace.yaml +++ /dev/null @@ -1,144 +0,0 @@ -defaults: - - _self_ - - task: real_pusht_image - -name: train_diffusion_transformer_hybrid -_target_: diffusion_policy.workspace.train_diffusion_transformer_hybrid_workspace.TrainDiffusionTransformerHybridWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 16 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -dataset_obs_steps: ${n_obs_steps} -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_cond: True - -policy: - _target_: diffusion_policy.policy.diffusion_transformer_hybrid_image_policy.DiffusionTransformerHybridImagePolicy - - shape_meta: ${shape_meta} - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddim.DDIMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - # beta_schedule is important - # this is the best we found - beta_schedule: squaredcos_cap_v2 - clip_sample: True - set_alpha_to_one: True - steps_offset: 0 - prediction_type: epsilon # or sample - - horizon: ${horizon} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 8 - - crop_shape: [216, 288] # ch, cw 320x240 90% - obs_encoder_group_norm: True - eval_fixed_crop: True - - n_layer: 8 - n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP - n_head: 4 - n_emb: 256 - p_drop_emb: 0.0 - p_drop_attn: 0.3 - causal_attn: True - time_as_cond: True # if false, use BERT like encoder only arch, time as input - obs_as_cond: ${obs_as_cond} - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 64 - num_workers: 8 - shuffle: True - pin_memory: True - persistent_workers: True - -val_dataloader: - batch_size: 64 - num_workers: 8 - shuffle: False - pin_memory: True - persistent_workers: True - -optimizer: - transformer_weight_decay: 1.0e-3 - obs_encoder_weight_decay: 1.0e-6 - learning_rate: 1.0e-4 - betas: [0.9, 0.95] - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - # Transformer needs LR warmup - lr_warmup_steps: 500 - num_epochs: 600 - gradient_accumulate_every: 1 - # EMA destroys performance when used with BatchNorm - # replace BatchNorm with GroupNorm. - use_ema: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: train_loss - mode: min - k: 5 - format_str: 'epoch={epoch:04d}-train_loss={train_loss:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_ddim_hybrid_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_ddim_hybrid_workspace.yaml deleted file mode 100644 index 414cbbaef..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_ddim_hybrid_workspace.yaml +++ /dev/null @@ -1,140 +0,0 @@ -defaults: - - _self_ - - task: lift_image_abs - -name: train_diffusion_unet_hybrid -_target_: diffusion_policy.workspace.train_diffusion_unet_hybrid_workspace.TrainDiffusionUnetHybridWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 16 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -dataset_obs_steps: ${n_obs_steps} -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_global_cond: True - -policy: - _target_: diffusion_policy.policy.diffusion_unet_hybrid_image_policy.DiffusionUnetHybridImagePolicy - - shape_meta: ${shape_meta} - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddim.DDIMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - # beta_schedule is important - # this is the best we found - beta_schedule: squaredcos_cap_v2 - clip_sample: True - set_alpha_to_one: True - steps_offset: 0 - prediction_type: epsilon # or sample - - horizon: ${horizon} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 8 - obs_as_global_cond: ${obs_as_global_cond} - # crop_shape: [76, 76] - crop_shape: null - diffusion_step_embed_dim: 128 - down_dims: [256,512,1024] - kernel_size: 5 - n_groups: 8 - cond_predict_scale: True - obs_encoder_group_norm: True - eval_fixed_crop: True - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 64 - num_workers: 8 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 64 - num_workers: 8 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 3000 - gradient_accumulate_every: 1 - # EMA destroys performance when used with BatchNorm - # replace BatchNorm with GroupNorm. - use_ema: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_ddim_lowdim_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_ddim_lowdim_workspace.yaml deleted file mode 100644 index 90cd2f850..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_ddim_lowdim_workspace.yaml +++ /dev/null @@ -1,146 +0,0 @@ -defaults: - - _self_ - - task: pusht_lowdim - -name: train_diffusion_unet_lowdim -_target_: diffusion_policy.workspace.train_diffusion_unet_lowdim_workspace.TrainDiffusionUnetLowdimWorkspace - -obs_dim: ${task.obs_dim} -action_dim: ${task.action_dim} -keypoint_dim: ${task.keypoint_dim} -task_name: ${task.name} -exp_name: "default" - -horizon: 16 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_local_cond: False -obs_as_global_cond: False -pred_action_steps_only: False - -policy: - _target_: diffusion_policy.policy.diffusion_unet_lowdim_policy.DiffusionUnetLowdimPolicy - - model: - _target_: diffusion_policy.model.diffusion.conditional_unet1d.ConditionalUnet1D - input_dim: "${eval: ${task.action_dim} if ${obs_as_local_cond} or ${obs_as_global_cond} else ${task.obs_dim} + ${task.action_dim}}" - local_cond_dim: "${eval: ${task.obs_dim} if ${obs_as_local_cond} else None}" - global_cond_dim: "${eval: ${task.obs_dim}*${n_obs_steps} if ${obs_as_global_cond} else None}" - diffusion_step_embed_dim: 256 - down_dims: [256, 512, 1024] - kernel_size: 5 - n_groups: 8 - cond_predict_scale: True - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddim.DDIMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - # beta_schedule is important - # this is the best we found - beta_schedule: squaredcos_cap_v2 - clip_sample: True - set_alpha_to_one: True - steps_offset: 0 - prediction_type: epsilon # or sample - - horizon: ${horizon} - obs_dim: ${obs_dim} - action_dim: ${action_dim} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 8 - obs_as_local_cond: ${obs_as_local_cond} - obs_as_global_cond: ${obs_as_global_cond} - pred_action_steps_only: ${pred_action_steps_only} - oa_step_convention: True - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 256 - num_workers: 1 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 256 - num_workers: 1 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 8000 - gradient_accumulate_every: 1 - use_ema: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_hybrid_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_hybrid_workspace.yaml deleted file mode 100644 index 9f9b1a6f8..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_hybrid_workspace.yaml +++ /dev/null @@ -1,137 +0,0 @@ -defaults: - - _self_ - - task: lift_image_abs - -name: train_diffusion_unet_hybrid -_target_: diffusion_policy.workspace.train_diffusion_unet_hybrid_workspace.TrainDiffusionUnetHybridWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 16 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -dataset_obs_steps: ${n_obs_steps} -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_global_cond: True - -policy: - _target_: diffusion_policy.policy.diffusion_unet_hybrid_image_policy.DiffusionUnetHybridImagePolicy - - shape_meta: ${shape_meta} - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - beta_schedule: squaredcos_cap_v2 - variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan - clip_sample: True # required when predict_epsilon=False - prediction_type: epsilon # or sample - - horizon: ${horizon} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 100 - obs_as_global_cond: ${obs_as_global_cond} - crop_shape: [76, 76] - # crop_shape: null - diffusion_step_embed_dim: 128 - down_dims: [512, 1024, 2048] - kernel_size: 5 - n_groups: 8 - cond_predict_scale: True - obs_encoder_group_norm: True - eval_fixed_crop: True - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 64 - num_workers: 8 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 64 - num_workers: 8 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 3050 - gradient_accumulate_every: 1 - # EMA destroys performance when used with BatchNorm - # replace BatchNorm with GroupNorm. - use_ema: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_image_pretrained_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_image_pretrained_workspace.yaml deleted file mode 100644 index 122bea335..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_image_pretrained_workspace.yaml +++ /dev/null @@ -1,149 +0,0 @@ -defaults: - - _self_ - - task: lift_image_abs - -name: train_diffusion_unet_image -_target_: diffusion_policy.workspace.train_diffusion_unet_image_workspace.TrainDiffusionUnetImageWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 16 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -dataset_obs_steps: ${n_obs_steps} -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_global_cond: True - -policy: - _target_: diffusion_policy.policy.diffusion_unet_image_policy.DiffusionUnetImagePolicy - - shape_meta: ${shape_meta} - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - beta_schedule: squaredcos_cap_v2 - variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan - clip_sample: True # required when predict_epsilon=False - prediction_type: epsilon # or sample - - obs_encoder: - _target_: diffusion_policy.model.vision.multi_image_obs_encoder.MultiImageObsEncoder - shape_meta: ${shape_meta} - rgb_model: - _target_: diffusion_policy.model.vision.model_getter.get_resnet - name: resnet18 - weights: IMAGENET1K_V1 # or r3m - resize_shape: [256, 256] - crop_shape: [224, 224] - random_crop: False - use_group_norm: False - share_rgb_model: True - imagenet_norm: True - - horizon: ${horizon} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 100 - obs_as_global_cond: ${obs_as_global_cond} - # crop_shape: null - diffusion_step_embed_dim: 128 - down_dims: [512, 1024, 2048] - kernel_size: 5 - n_groups: 8 - cond_predict_scale: True - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 64 - num_workers: 4 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 64 - num_workers: 4 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 8000 - gradient_accumulate_every: 1 - # EMA destroys performance when used with BatchNorm - # replace BatchNorm with GroupNorm. - use_ema: True - freeze_encoder: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_image_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_image_workspace.yaml deleted file mode 100644 index f5636e392..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_image_workspace.yaml +++ /dev/null @@ -1,150 +0,0 @@ -defaults: - - _self_ - - task: lift_image_abs - -name: train_diffusion_unet_image -_target_: diffusion_policy.workspace.train_diffusion_unet_image_workspace.TrainDiffusionUnetImageWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 16 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -dataset_obs_steps: ${n_obs_steps} -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_global_cond: True - -policy: - _target_: diffusion_policy.policy.diffusion_unet_image_policy.DiffusionUnetImagePolicy - - shape_meta: ${shape_meta} - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - beta_schedule: squaredcos_cap_v2 - variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan - clip_sample: True # required when predict_epsilon=False - prediction_type: epsilon # or sample - - obs_encoder: - _target_: diffusion_policy.model.vision.multi_image_obs_encoder.MultiImageObsEncoder - shape_meta: ${shape_meta} - rgb_model: - _target_: diffusion_policy.model.vision.model_getter.get_resnet - name: resnet18 - weights: null - resize_shape: null - crop_shape: [76, 76] - # constant center crop - random_crop: True - use_group_norm: True - share_rgb_model: False - imagenet_norm: True - - horizon: ${horizon} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 100 - obs_as_global_cond: ${obs_as_global_cond} - # crop_shape: null - diffusion_step_embed_dim: 128 - down_dims: [512, 1024, 2048] - kernel_size: 5 - n_groups: 8 - cond_predict_scale: True - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 64 - num_workers: 4 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 64 - num_workers: 4 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 8000 - gradient_accumulate_every: 1 - # EMA destroys performance when used with BatchNorm - # replace BatchNorm with GroupNorm. - use_ema: True - freeze_encoder: False - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_lowdim_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_lowdim_workspace.yaml deleted file mode 100644 index 7b4d35f48..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_lowdim_workspace.yaml +++ /dev/null @@ -1,143 +0,0 @@ -defaults: - - _self_ - - task: pusht_lowdim - -name: train_diffusion_unet_lowdim -_target_: diffusion_policy.workspace.train_diffusion_unet_lowdim_workspace.TrainDiffusionUnetLowdimWorkspace - -obs_dim: ${task.obs_dim} -action_dim: ${task.action_dim} -keypoint_dim: ${task.keypoint_dim} -task_name: ${task.name} -exp_name: "default" - -horizon: 16 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_local_cond: False -obs_as_global_cond: True -pred_action_steps_only: False - -policy: - _target_: diffusion_policy.policy.diffusion_unet_lowdim_policy.DiffusionUnetLowdimPolicy - - model: - _target_: diffusion_policy.model.diffusion.conditional_unet1d.ConditionalUnet1D - input_dim: "${eval: ${task.action_dim} if ${obs_as_local_cond} or ${obs_as_global_cond} else ${task.obs_dim} + ${task.action_dim}}" - local_cond_dim: "${eval: ${task.obs_dim} if ${obs_as_local_cond} else None}" - global_cond_dim: "${eval: ${task.obs_dim}*${n_obs_steps} if ${obs_as_global_cond} else None}" - diffusion_step_embed_dim: 256 - down_dims: [256, 512, 1024] - kernel_size: 5 - n_groups: 8 - cond_predict_scale: True - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - beta_schedule: squaredcos_cap_v2 - variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan - clip_sample: True # required when predict_epsilon=False - prediction_type: epsilon # or sample - - horizon: ${horizon} - obs_dim: ${obs_dim} - action_dim: ${action_dim} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 100 - obs_as_local_cond: ${obs_as_local_cond} - obs_as_global_cond: ${obs_as_global_cond} - pred_action_steps_only: ${pred_action_steps_only} - oa_step_convention: True - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 256 - num_workers: 1 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 256 - num_workers: 1 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 5000 - gradient_accumulate_every: 1 - use_ema: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_real_hybrid_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_real_hybrid_workspace.yaml deleted file mode 100644 index 5b323661a..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_real_hybrid_workspace.yaml +++ /dev/null @@ -1,141 +0,0 @@ -defaults: - - _self_ - - task: real_pusht_image - -name: train_diffusion_unet_hybrid -_target_: diffusion_policy.workspace.train_diffusion_unet_hybrid_workspace.TrainDiffusionUnetHybridWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 16 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -dataset_obs_steps: ${n_obs_steps} -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_global_cond: True - -policy: - _target_: diffusion_policy.policy.diffusion_unet_hybrid_image_policy.DiffusionUnetHybridImagePolicy - - shape_meta: ${shape_meta} - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddim.DDIMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - # beta_schedule is important - # this is the best we found - beta_schedule: squaredcos_cap_v2 - clip_sample: True - set_alpha_to_one: True - steps_offset: 0 - prediction_type: epsilon # or sample - - horizon: ${horizon} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 8 - obs_as_global_cond: ${obs_as_global_cond} - # crop_shape: [76, 76] # 84x84 90% - crop_shape: [216, 288] # ch, cw 320x240 90% - # crop_shape: null - diffusion_step_embed_dim: 128 - down_dims: [256,512,1024] - kernel_size: 5 - n_groups: 8 - cond_predict_scale: True - obs_encoder_group_norm: True - eval_fixed_crop: True - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 64 - num_workers: 8 - shuffle: True - pin_memory: True - persistent_workers: True - -val_dataloader: - batch_size: 64 - num_workers: 8 - shuffle: False - pin_memory: True - persistent_workers: True - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 600 - gradient_accumulate_every: 1 - # EMA destroys performance when used with BatchNorm - # replace BatchNorm with GroupNorm. - use_ema: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: train_loss - mode: min - k: 5 - format_str: 'epoch={epoch:04d}-train_loss={train_loss:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_real_image_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_real_image_workspace.yaml deleted file mode 100644 index f6187ef14..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_real_image_workspace.yaml +++ /dev/null @@ -1,152 +0,0 @@ -defaults: - - _self_ - - task: real_pusht_image - -name: train_diffusion_unet_image -_target_: diffusion_policy.workspace.train_diffusion_unet_image_workspace.TrainDiffusionUnetImageWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 16 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -dataset_obs_steps: ${n_obs_steps} -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_global_cond: True - -policy: - _target_: diffusion_policy.policy.diffusion_unet_image_policy.DiffusionUnetImagePolicy - - shape_meta: ${shape_meta} - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddim.DDIMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - # beta_schedule is important - # this is the best we found - beta_schedule: squaredcos_cap_v2 - clip_sample: True - set_alpha_to_one: True - steps_offset: 0 - prediction_type: epsilon # or sample - - obs_encoder: - _target_: diffusion_policy.model.vision.multi_image_obs_encoder.MultiImageObsEncoder - shape_meta: ${shape_meta} - rgb_model: - _target_: diffusion_policy.model.vision.model_getter.get_resnet - name: resnet18 - weights: null - resize_shape: [240, 320] - crop_shape: [216, 288] # ch, cw 240x320 90% - random_crop: True - use_group_norm: True - share_rgb_model: False - imagenet_norm: True - - horizon: ${horizon} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 100 - obs_as_global_cond: ${obs_as_global_cond} - # crop_shape: null - diffusion_step_embed_dim: 128 - down_dims: [512, 1024, 2048] - kernel_size: 5 - n_groups: 8 - cond_predict_scale: True - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 64 - num_workers: 8 - shuffle: True - pin_memory: True - persistent_workers: True - -val_dataloader: - batch_size: 64 - num_workers: 8 - shuffle: False - pin_memory: True - persistent_workers: True - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 600 - gradient_accumulate_every: 1 - # EMA destroys performance when used with BatchNorm - # replace BatchNorm with GroupNorm. - use_ema: True - freeze_encoder: False - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: train_loss - mode: min - k: 5 - format_str: 'epoch={epoch:04d}-train_loss={train_loss:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_real_pretrained_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_real_pretrained_workspace.yaml deleted file mode 100644 index f129a1edf..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_real_pretrained_workspace.yaml +++ /dev/null @@ -1,152 +0,0 @@ -defaults: - - _self_ - - task: real_pusht_image - -name: train_diffusion_unet_image -_target_: diffusion_policy.workspace.train_diffusion_unet_image_workspace.TrainDiffusionUnetImageWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 16 -n_obs_steps: 2 -n_action_steps: 8 -n_latency_steps: 0 -dataset_obs_steps: ${n_obs_steps} -past_action_visible: False -keypoint_visible_rate: 1.0 -obs_as_global_cond: True - -policy: - _target_: diffusion_policy.policy.diffusion_unet_image_policy.DiffusionUnetImagePolicy - - shape_meta: ${shape_meta} - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddim.DDIMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - # beta_schedule is important - # this is the best we found - beta_schedule: squaredcos_cap_v2 - clip_sample: True - set_alpha_to_one: True - steps_offset: 0 - prediction_type: epsilon # or sample - - obs_encoder: - _target_: diffusion_policy.model.vision.multi_image_obs_encoder.MultiImageObsEncoder - shape_meta: ${shape_meta} - rgb_model: - _target_: diffusion_policy.model.vision.model_getter.get_resnet - name: resnet18 - weights: IMAGENET1K_V1 # or r3m - resize_shape: [224,224] - crop_shape: null - random_crop: False - use_group_norm: False - share_rgb_model: True - imagenet_norm: True - - horizon: ${horizon} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 100 - obs_as_global_cond: ${obs_as_global_cond} - # crop_shape: null - diffusion_step_embed_dim: 128 - down_dims: [512, 1024, 2048] - kernel_size: 5 - n_groups: 8 - cond_predict_scale: True - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 64 - num_workers: 8 - shuffle: True - pin_memory: True - persistent_workers: True - -val_dataloader: - batch_size: 64 - num_workers: 8 - shuffle: False - pin_memory: True - persistent_workers: True - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 600 - gradient_accumulate_every: 1 - # EMA destroys performance when used with BatchNorm - # replace BatchNorm with GroupNorm. - use_ema: True - freeze_encoder: True - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: train_loss - mode: min - k: 5 - format_str: 'epoch={epoch:04d}-train_loss={train_loss:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_video_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_video_workspace.yaml deleted file mode 100644 index 4af5f1fe7..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_diffusion_unet_video_workspace.yaml +++ /dev/null @@ -1,128 +0,0 @@ -defaults: - - _self_ - - task: lift_image_abs - -name: train_diffusion_unet_video -_target_: diffusion_policy.workspace.train_diffusion_unet_video_workspace.TrainDiffusionUnetVideoWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 16 -n_obs_steps: 4 -n_action_steps: 8 -past_action_visible: False -keypoint_visible_rate: 1.0 -lowdim_as_global_cond: True - -policy: - _target_: diffusion_policy.policy.diffusion_unet_video_policy.DiffusionUnetVideoPolicy - - shape_meta: ${shape_meta} - - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler - num_train_timesteps: 100 - beta_start: 0.0001 - beta_end: 0.02 - beta_schedule: squaredcos_cap_v2 - variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan - clip_sample: True # required when predict_epsilon=False - prediction_type: epsilon # or sample - - rgb_net: - _target_: diffusion_policy.model.obs_encoder.video_core.VideoCore - - backbone: - _target_: diffusion_policy.model.obs_encoder.video_core.VideoResNet - - norm_groups: 8 - input_channel: 3 - model_depth: 50 # ResNet 50 (18,34 not yet available) - - pool: - _target_: diffusion_policy.model.ibc.global_avgpool.GlobalAvgpool - - dim: [2,3,4] - - horizon: ${horizon} - n_action_steps: ${n_action_steps} - n_obs_steps: ${n_obs_steps} - num_inference_steps: 100 - lowdim_as_global_cond: ${lowdim_as_global_cond} - diffusion_step_embed_dim: 128 - down_dims: [512, 1024, 2048] - kernel_size: 5 - n_groups: 8 - cond_predict_scale: True - - # TemporalAggregator Parameters - channel_mults: [1,1] - n_blocks_per_level: 1 - ta_kernel_size: 3 - ta_n_groups: 1 - - # scheduler.step params - # predict_epsilon: True - -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - update_after_step: 0 - inv_gamma: 1.0 - power: 0.75 - min_value: 0.0 - max_value: 0.9999 - -dataloader: - batch_size: 32 - num_workers: 1 - shuffle: True - pin_memory: True - -optimizer: - _target_: torch.optim.AdamW - lr: 0.0001 # 1e-4 - betas: [0.95, 0.999] - eps: 0.00000001 # 1e-8 - weight_decay: 0.000001 # 1e-6 - -training: - device: "cuda:0" - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 1500 - gradient_accumulate_every: 1 - eval_every: 5000 - eval_first: False - val_every: 300 - use_ema: False - tqdm_interval_sec: 1.0 - seed: 42 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_score - mode: max - k: 5 - format_str: 'epoch={epoch:03d}-test_score={test_score:.3f}.ckpt' - save_last_ckpt: False - save_last_snapshot: False - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_ibc_dfo_hybrid_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_ibc_dfo_hybrid_workspace.yaml deleted file mode 100644 index b6f5d47c8..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_ibc_dfo_hybrid_workspace.yaml +++ /dev/null @@ -1,111 +0,0 @@ -defaults: - - _self_ - - task: pusht_image - -name: train_ibc_dfo_hybrid -_target_: diffusion_policy.workspace.train_ibc_dfo_hybrid_workspace.TrainIbcDfoHybridWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 2 -n_obs_steps: 2 -n_action_steps: 1 -n_latency_steps: 0 -dataset_obs_steps: ${n_obs_steps} -past_action_visible: False -keypoint_visible_rate: 1.0 - -policy: - _target_: diffusion_policy.policy.ibc_dfo_hybrid_image_policy.IbcDfoHybridImagePolicy - - shape_meta: ${shape_meta} - - horizon: ${horizon} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - dropout: 0.1 - train_n_neg: 1024 - pred_n_iter: 5 - pred_n_samples: 1024 - kevin_inference: False - andy_train: False - obs_encoder_group_norm: True - eval_fixed_crop: True - crop_shape: [84, 84] - -dataloader: - batch_size: 128 - num_workers: 8 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 128 - num_workers: 8 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 3050 - gradient_accumulate_every: 1 - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - sample_max_batch: 128 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_ibc_dfo_lowdim_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_ibc_dfo_lowdim_workspace.yaml deleted file mode 100644 index 9e73a93f8..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_ibc_dfo_lowdim_workspace.yaml +++ /dev/null @@ -1,109 +0,0 @@ -defaults: - - _self_ - - task: pusht_lowdim - -name: train_ibc_dfo_lowdim -_target_: diffusion_policy.workspace.train_ibc_dfo_lowdim_workspace.TrainIbcDfoLowdimWorkspace - -obs_dim: ${task.obs_dim} -action_dim: ${task.action_dim} -keypoint_dim: ${task.keypoint_dim} -task_name: ${task.name} -exp_name: "default" - -horizon: 2 -n_obs_steps: 2 -n_action_steps: 1 -n_latency_steps: 0 -past_action_visible: False -keypoint_visible_rate: 1.0 - -policy: - _target_: diffusion_policy.policy.ibc_dfo_lowdim_policy.IbcDfoLowdimPolicy - - horizon: ${horizon} - obs_dim: ${obs_dim} - action_dim: ${action_dim} - n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'} - n_obs_steps: ${n_obs_steps} - dropout: 0.1 - train_n_neg: 1024 - pred_n_iter: 5 - pred_n_samples: 1024 - kevin_inference: False - andy_train: False - -dataloader: - batch_size: 256 - num_workers: 1 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 256 - num_workers: 1 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 5000 - gradient_accumulate_every: 1 - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - sample_max_batch: 128 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_ibc_dfo_real_hybrid_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_ibc_dfo_real_hybrid_workspace.yaml deleted file mode 100644 index 8255dc1b2..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_ibc_dfo_real_hybrid_workspace.yaml +++ /dev/null @@ -1,111 +0,0 @@ -defaults: - - _self_ - - task: real_pusht_image - -name: train_ibc_dfo_hybrid -_target_: diffusion_policy.workspace.train_ibc_dfo_hybrid_workspace.TrainIbcDfoHybridWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: 2 -n_obs_steps: 2 -n_action_steps: 1 -n_latency_steps: 1 -dataset_obs_steps: ${n_obs_steps} -past_action_visible: False -keypoint_visible_rate: 1.0 - -policy: - _target_: diffusion_policy.policy.ibc_dfo_hybrid_image_policy.IbcDfoHybridImagePolicy - - shape_meta: ${shape_meta} - - horizon: ${horizon} - n_action_steps: ${n_action_steps} - n_obs_steps: ${n_obs_steps} - dropout: 0.1 - train_n_neg: 256 - pred_n_iter: 3 - pred_n_samples: 1024 - kevin_inference: False - andy_train: False - obs_encoder_group_norm: True - eval_fixed_crop: True - crop_shape: [216, 288] # ch, cw 320x240 90% - -dataloader: - batch_size: 128 - num_workers: 8 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 128 - num_workers: 1 - shuffle: False - pin_memory: True - persistent_workers: False - -optimizer: - _target_: torch.optim.AdamW - lr: 1.0e-4 - betas: [0.95, 0.999] - eps: 1.0e-8 - weight_decay: 1.0e-6 - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - lr_scheduler: cosine - lr_warmup_steps: 500 - num_epochs: 1000 - gradient_accumulate_every: 1 - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 5 - val_every: 1 - sample_every: 5 - sample_max_batch: 128 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: train_action_mse_error - mode: min - k: 5 - format_str: 'epoch={epoch:04d}-train_action_mse_error={train_action_mse_error:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_robomimic_image_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_robomimic_image_workspace.yaml deleted file mode 100644 index 4cea6ab05..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_robomimic_image_workspace.yaml +++ /dev/null @@ -1,92 +0,0 @@ -defaults: - - _self_ - - task: lift_image - -name: train_robomimic_image -_target_: diffusion_policy.workspace.train_robomimic_image_workspace.TrainRobomimicImageWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: &horizon 10 -n_obs_steps: 1 -n_action_steps: 1 -n_latency_steps: 0 -dataset_obs_steps: *horizon -past_action_visible: False -keypoint_visible_rate: 1.0 - -policy: - _target_: diffusion_policy.policy.robomimic_image_policy.RobomimicImagePolicy - shape_meta: ${shape_meta} - algo_name: bc_rnn - obs_type: image - # oc.select resolver: key, default - task_name: ${oc.select:task.task_name,lift} - dataset_type: ${oc.select:task.dataset_type,ph} - crop_shape: [76,76] - -dataloader: - batch_size: 64 - num_workers: 16 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 64 - num_workers: 16 - shuffle: False - pin_memory: True - persistent_workers: False - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - num_epochs: 3050 - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_robomimic_lowdim_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_robomimic_lowdim_workspace.yaml deleted file mode 100644 index a657d1ba2..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_robomimic_lowdim_workspace.yaml +++ /dev/null @@ -1,92 +0,0 @@ -defaults: - - _self_ - - task: pusht_lowdim - -name: train_robomimic_lowdim -_target_: diffusion_policy.workspace.train_robomimic_lowdim_workspace.TrainRobomimicLowdimWorkspace - -obs_dim: ${task.obs_dim} -action_dim: ${task.action_dim} -transition_dim: "${eval: ${task.obs_dim} + ${task.action_dim}}" -task_name: ${task.name} -exp_name: "default" - -horizon: 10 -n_obs_steps: 1 -n_action_steps: 1 -n_latency_steps: 0 -past_action_visible: False -keypoint_visible_rate: 1.0 - -policy: - _target_: diffusion_policy.policy.robomimic_lowdim_policy.RobomimicLowdimPolicy - action_dim: ${action_dim} - obs_dim: ${obs_dim} - algo_name: bc_rnn - obs_type: low_dim - # oc.select resolver: key, default - task_name: ${oc.select:task.task_name,lift} - dataset_type: ${oc.select:task.dataset_type,ph} - -dataloader: - batch_size: 256 - num_workers: 1 - shuffle: True - pin_memory: True - persistent_workers: False - -val_dataloader: - batch_size: 256 - num_workers: 1 - shuffle: False - pin_memory: True - persistent_workers: False - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - num_epochs: 5000 - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: test_mean_score - mode: max - k: 5 - format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_robomimic_real_image_workspace.yaml b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_robomimic_real_image_workspace.yaml deleted file mode 100644 index 983976c46..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/config/train_robomimic_real_image_workspace.yaml +++ /dev/null @@ -1,92 +0,0 @@ -defaults: - - _self_ - - task: real_pusht_image - -name: train_robomimic_image -_target_: diffusion_policy.workspace.train_robomimic_image_workspace.TrainRobomimicImageWorkspace - -task_name: ${task.name} -shape_meta: ${task.shape_meta} -exp_name: "default" - -horizon: &horizon 10 -n_obs_steps: 1 -n_action_steps: 1 -n_latency_steps: 1 -dataset_obs_steps: *horizon -past_action_visible: False -keypoint_visible_rate: 1.0 - -policy: - _target_: diffusion_policy.policy.robomimic_image_policy.RobomimicImagePolicy - shape_meta: ${shape_meta} - algo_name: bc_rnn - obs_type: image - # oc.select resolver: key, default - task_name: ${oc.select:task.task_name,tool_hang} - dataset_type: ${oc.select:task.dataset_type,ph} - crop_shape: [216, 288] # ch, cw 320x240 90% - -dataloader: - batch_size: 32 - num_workers: 8 - shuffle: True - pin_memory: True - persistent_workers: True - -val_dataloader: - batch_size: 32 - num_workers: 1 - shuffle: False - pin_memory: True - persistent_workers: False - -training: - device: "cuda:0" - seed: 42 - debug: False - resume: True - # optimization - num_epochs: 1000 - # training loop control - # in epochs - rollout_every: 50 - checkpoint_every: 50 - val_every: 1 - sample_every: 5 - # steps per epoch - max_train_steps: null - max_val_steps: null - # misc - tqdm_interval_sec: 1.0 - -logging: - project: diffusion_policy_debug - resume: True - mode: online - name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - tags: ["${name}", "${task_name}", "${exp_name}"] - id: null - group: null - -checkpoint: - topk: - monitor_key: train_loss - mode: min - k: 5 - format_str: 'epoch={epoch:04d}-train_loss={train_loss:.3f}.ckpt' - save_last_ckpt: True - save_last_snapshot: False - -multi_run: - run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name} - -hydra: - job: - override_dirname: ${name} - run: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - sweep: - dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name} - subdir: ${hydra.job.num} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/base_dataset.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/base_dataset.py deleted file mode 100644 index 72a39e892..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/base_dataset.py +++ /dev/null @@ -1,51 +0,0 @@ -from typing import Dict - -import torch -import torch.nn -from diffusion_policy.model.common.normalizer import LinearNormalizer - -class BaseLowdimDataset(torch.utils.data.Dataset): - def get_validation_dataset(self) -> 'BaseLowdimDataset': - # return an empty dataset by default - return BaseLowdimDataset() - - def get_normalizer(self, **kwargs) -> LinearNormalizer: - raise NotImplementedError() - - def get_all_actions(self) -> torch.Tensor: - raise NotImplementedError() - - def __len__(self) -> int: - return 0 - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - """ - output: - obs: T, Do - action: T, Da - """ - raise NotImplementedError() - - -class BaseImageDataset(torch.utils.data.Dataset): - def get_validation_dataset(self) -> 'BaseLowdimDataset': - # return an empty dataset by default - return BaseImageDataset() - - def get_normalizer(self, **kwargs) -> LinearNormalizer: - raise NotImplementedError() - - def get_all_actions(self) -> torch.Tensor: - raise NotImplementedError() - - def __len__(self) -> int: - return 0 - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - """ - output: - obs: - key: T, * - action: T, Da - """ - raise NotImplementedError() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/blockpush_lowdim_dataset.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/blockpush_lowdim_dataset.py deleted file mode 100644 index 86242c224..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/blockpush_lowdim_dataset.py +++ /dev/null @@ -1,126 +0,0 @@ -from typing import Dict -import torch -import numpy as np -import copy -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.common.replay_buffer import ReplayBuffer -from diffusion_policy.common.sampler import SequenceSampler, get_val_mask -from diffusion_policy.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer -from diffusion_policy.dataset.base_dataset import BaseLowdimDataset - -class BlockPushLowdimDataset(BaseLowdimDataset): - def __init__(self, - zarr_path, - horizon=1, - pad_before=0, - pad_after=0, - obs_key='obs', - action_key='action', - obs_eef_target=True, - use_manual_normalizer=False, - seed=42, - val_ratio=0.0 - ): - super().__init__() - self.replay_buffer = ReplayBuffer.copy_from_path( - zarr_path, keys=[obs_key, action_key]) - - val_mask = get_val_mask( - n_episodes=self.replay_buffer.n_episodes, - val_ratio=val_ratio, - seed=seed) - train_mask = ~val_mask - self.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=horizon, - pad_before=pad_before, - pad_after=pad_after, - episode_mask=train_mask) - self.obs_key = obs_key - self.action_key = action_key - self.obs_eef_target = obs_eef_target - self.use_manual_normalizer = use_manual_normalizer - self.train_mask = train_mask - self.horizon = horizon - self.pad_before = pad_before - self.pad_after = pad_after - - def get_validation_dataset(self): - val_set = copy.copy(self) - val_set.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=self.horizon, - pad_before=self.pad_before, - pad_after=self.pad_after, - episode_mask=~self.train_mask - ) - val_set.train_mask = ~self.train_mask - return val_set - - def get_normalizer(self, mode='limits', **kwargs): - data = self._sample_to_data(self.replay_buffer) - - normalizer = LinearNormalizer() - if not self.use_manual_normalizer: - normalizer.fit(data=data, last_n_dims=1, mode=mode, **kwargs) - else: - x = data['obs'] - stat = { - 'max': np.max(x, axis=0), - 'min': np.min(x, axis=0), - 'mean': np.mean(x, axis=0), - 'std': np.std(x, axis=0) - } - - is_x = np.zeros(stat['max'].shape, dtype=bool) - is_y = np.zeros_like(is_x) - is_x[[0,3,6,8,10,13]] = True - is_y[[1,4,7,9,11,14]] = True - is_rot = ~(is_x|is_y) - - def normalizer_with_masks(stat, masks): - global_scale = np.ones_like(stat['max']) - global_offset = np.zeros_like(stat['max']) - for mask in masks: - output_max = 1 - output_min = -1 - input_max = stat['max'][mask].max() - input_min = stat['min'][mask].min() - input_range = input_max - input_min - scale = (output_max - output_min) / input_range - offset = output_min - scale * input_min - global_scale[mask] = scale - global_offset[mask] = offset - return SingleFieldLinearNormalizer.create_manual( - scale=global_scale, - offset=global_offset, - input_stats_dict=stat - ) - - normalizer['obs'] = normalizer_with_masks(stat, [is_x, is_y, is_rot]) - normalizer['action'] = SingleFieldLinearNormalizer.create_fit( - data['action'], last_n_dims=1, mode=mode, **kwargs) - return normalizer - - def get_all_actions(self) -> torch.Tensor: - return torch.from_numpy(self.replay_buffer['action']) - - def __len__(self) -> int: - return len(self.sampler) - - def _sample_to_data(self, sample): - obs = sample[self.obs_key] # T, D_o - if not self.obs_eef_target: - obs[:,8:10] = 0 - data = { - 'obs': obs, - 'action': sample[self.action_key], # T, D_a - } - return data - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - sample = self.sampler.sample_sequence(idx) - data = self._sample_to_data(sample) - - torch_data = dict_apply(data, torch.from_numpy) - return torch_data diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/kitchen_lowdim_dataset.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/kitchen_lowdim_dataset.py deleted file mode 100644 index 601e21cb1..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/kitchen_lowdim_dataset.py +++ /dev/null @@ -1,91 +0,0 @@ -from typing import Dict -import torch -import numpy as np -import copy -import pathlib -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.common.replay_buffer import ReplayBuffer -from diffusion_policy.common.sampler import SequenceSampler, get_val_mask -from diffusion_policy.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer -from diffusion_policy.dataset.base_dataset import BaseLowdimDataset - -class KitchenLowdimDataset(BaseLowdimDataset): - def __init__(self, - dataset_dir, - horizon=1, - pad_before=0, - pad_after=0, - seed=42, - val_ratio=0.0 - ): - super().__init__() - - data_directory = pathlib.Path(dataset_dir) - observations = np.load(data_directory / "observations_seq.npy") - actions = np.load(data_directory / "actions_seq.npy") - masks = np.load(data_directory / "existence_mask.npy") - - self.replay_buffer = ReplayBuffer.create_empty_numpy() - for i in range(len(masks)): - eps_len = int(masks[i].sum()) - obs = observations[i,:eps_len].astype(np.float32) - action = actions[i,:eps_len].astype(np.float32) - data = { - 'obs': obs, - 'action': action - } - self.replay_buffer.add_episode(data) - - val_mask = get_val_mask( - n_episodes=self.replay_buffer.n_episodes, - val_ratio=val_ratio, - seed=seed) - train_mask = ~val_mask - self.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=horizon, - pad_before=pad_before, - pad_after=pad_after, - episode_mask=train_mask) - - self.train_mask = train_mask - self.horizon = horizon - self.pad_before = pad_before - self.pad_after = pad_after - - def get_validation_dataset(self): - val_set = copy.copy(self) - val_set.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=self.horizon, - pad_before=self.pad_before, - pad_after=self.pad_after, - episode_mask=~self.train_mask - ) - val_set.train_mask = ~self.train_mask - return val_set - - def get_normalizer(self, mode='limits', **kwargs): - data = { - 'obs': self.replay_buffer['obs'], - 'action': self.replay_buffer['action'] - } - if 'range_eps' not in kwargs: - # to prevent blowing up dims that barely change - kwargs['range_eps'] = 5e-2 - normalizer = LinearNormalizer() - normalizer.fit(data=data, last_n_dims=1, mode=mode, **kwargs) - return normalizer - - def get_all_actions(self) -> torch.Tensor: - return torch.from_numpy(self.replay_buffer['action']) - - def __len__(self) -> int: - return len(self.sampler) - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - sample = self.sampler.sample_sequence(idx) - data = sample - - torch_data = dict_apply(data, torch.from_numpy) - return torch_data diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/kitchen_mjl_lowdim_dataset.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/kitchen_mjl_lowdim_dataset.py deleted file mode 100644 index e3173818c..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/kitchen_mjl_lowdim_dataset.py +++ /dev/null @@ -1,112 +0,0 @@ -from typing import Dict -import torch -import numpy as np -import copy -import pathlib -from tqdm import tqdm -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.common.replay_buffer import ReplayBuffer -from diffusion_policy.common.sampler import SequenceSampler, get_val_mask -from diffusion_policy.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer -from diffusion_policy.dataset.base_dataset import BaseLowdimDataset -from diffusion_policy.env.kitchen.kitchen_util import parse_mjl_logs - -class KitchenMjlLowdimDataset(BaseLowdimDataset): - def __init__(self, - dataset_dir, - horizon=1, - pad_before=0, - pad_after=0, - abs_action=True, - robot_noise_ratio=0.0, - seed=42, - val_ratio=0.0 - ): - super().__init__() - - if not abs_action: - raise NotImplementedError() - - robot_pos_noise_amp = np.array([0.1 , 0.1 , 0.1 , 0.1 , 0.1 , 0.1 , 0.1 , 0.1 , - 0.1 , 0.005 , 0.005 , 0.0005, 0.0005, 0.0005, 0.0005, 0.0005, - 0.0005, 0.005 , 0.005 , 0.005 , 0.1 , 0.1 , 0.1 , 0.005 , - 0.005 , 0.005 , 0.1 , 0.1 , 0.1 , 0.005 ], dtype=np.float32) - rng = np.random.default_rng(seed=seed) - - data_directory = pathlib.Path(dataset_dir) - self.replay_buffer = ReplayBuffer.create_empty_numpy() - for i, mjl_path in enumerate(tqdm(list(data_directory.glob('*/*.mjl')))): - try: - data = parse_mjl_logs(str(mjl_path.absolute()), skipamount=40) - qpos = data['qpos'].astype(np.float32) - obs = np.concatenate([ - qpos[:,:9], - qpos[:,-21:], - np.zeros((len(qpos),30),dtype=np.float32) - ], axis=-1) - if robot_noise_ratio > 0: - # add observation noise to match real robot - noise = robot_noise_ratio * robot_pos_noise_amp * rng.uniform( - low=-1., high=1., size=(obs.shape[0], 30)) - obs[:,:30] += noise - episode = { - 'obs': obs, - 'action': data['ctrl'].astype(np.float32) - } - self.replay_buffer.add_episode(episode) - except Exception as e: - print(i, e) - - val_mask = get_val_mask( - n_episodes=self.replay_buffer.n_episodes, - val_ratio=val_ratio, - seed=seed) - train_mask = ~val_mask - self.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=horizon, - pad_before=pad_before, - pad_after=pad_after, - episode_mask=train_mask) - - self.train_mask = train_mask - self.horizon = horizon - self.pad_before = pad_before - self.pad_after = pad_after - - def get_validation_dataset(self): - val_set = copy.copy(self) - val_set.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=self.horizon, - pad_before=self.pad_before, - pad_after=self.pad_after, - episode_mask=~self.train_mask - ) - val_set.train_mask = ~self.train_mask - return val_set - - def get_normalizer(self, mode='limits', **kwargs): - data = { - 'obs': self.replay_buffer['obs'], - 'action': self.replay_buffer['action'] - } - if 'range_eps' not in kwargs: - # to prevent blowing up dims that barely change - kwargs['range_eps'] = 5e-2 - normalizer = LinearNormalizer() - normalizer.fit(data=data, last_n_dims=1, mode=mode, **kwargs) - return normalizer - - def get_all_actions(self) -> torch.Tensor: - return torch.from_numpy(self.replay_buffer['action']) - - def __len__(self) -> int: - return len(self.sampler) - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - sample = self.sampler.sample_sequence(idx) - data = sample - - torch_data = dict_apply(data, torch.from_numpy) - return torch_data diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/pusht_dataset.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/pusht_dataset.py deleted file mode 100644 index dc3ec1c81..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/pusht_dataset.py +++ /dev/null @@ -1,97 +0,0 @@ -from typing import Dict -import torch -import numpy as np -import copy -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.common.replay_buffer import ReplayBuffer -from diffusion_policy.common.sampler import ( - SequenceSampler, get_val_mask, downsample_mask) -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.dataset.base_dataset import BaseLowdimDataset - -class PushTLowdimDataset(BaseLowdimDataset): - def __init__(self, - zarr_path, - horizon=1, - pad_before=0, - pad_after=0, - obs_key='keypoint', - state_key='state', - action_key='action', - seed=42, - val_ratio=0.0, - max_train_episodes=None - ): - super().__init__() - self.replay_buffer = ReplayBuffer.copy_from_path( - zarr_path, keys=[obs_key, state_key, action_key]) - - val_mask = get_val_mask( - n_episodes=self.replay_buffer.n_episodes, - val_ratio=val_ratio, - seed=seed) - train_mask = ~val_mask - train_mask = downsample_mask( - mask=train_mask, - max_n=max_train_episodes, - seed=seed) - - self.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=horizon, - pad_before=pad_before, - pad_after=pad_after, - episode_mask=train_mask - ) - self.obs_key = obs_key - self.state_key = state_key - self.action_key = action_key - self.train_mask = train_mask - self.horizon = horizon - self.pad_before = pad_before - self.pad_after = pad_after - - def get_validation_dataset(self): - val_set = copy.copy(self) - val_set.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=self.horizon, - pad_before=self.pad_before, - pad_after=self.pad_after, - episode_mask=~self.train_mask - ) - val_set.train_mask = ~self.train_mask - return val_set - - def get_normalizer(self, mode='limits', **kwargs): - data = self._sample_to_data(self.replay_buffer) - normalizer = LinearNormalizer() - normalizer.fit(data=data, last_n_dims=1, mode=mode, **kwargs) - return normalizer - - def get_all_actions(self) -> torch.Tensor: - return torch.from_numpy(self.replay_buffer[self.action_key]) - - def __len__(self) -> int: - return len(self.sampler) - - def _sample_to_data(self, sample): - keypoint = sample[self.obs_key] - state = sample[self.state_key] - agent_pos = state[:,:2] - obs = np.concatenate([ - keypoint.reshape(keypoint.shape[0], -1), - agent_pos], axis=-1) - - data = { - 'obs': obs, # T, D_o - 'action': sample[self.action_key], # T, D_a - } - return data - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - sample = self.sampler.sample_sequence(idx) - data = self._sample_to_data(sample) - - torch_data = dict_apply(data, torch.from_numpy) - return torch_data diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/pusht_image_dataset.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/pusht_image_dataset.py deleted file mode 100644 index f096a8f0f..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/pusht_image_dataset.py +++ /dev/null @@ -1,102 +0,0 @@ -from typing import Dict -import torch -import numpy as np -import copy -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.common.replay_buffer import ReplayBuffer -from diffusion_policy.common.sampler import ( - SequenceSampler, get_val_mask, downsample_mask) -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.dataset.base_dataset import BaseImageDataset -from diffusion_policy.common.normalize_util import get_image_range_normalizer - -class PushTImageDataset(BaseImageDataset): - def __init__(self, - zarr_path, - horizon=1, - pad_before=0, - pad_after=0, - seed=42, - val_ratio=0.0, - max_train_episodes=None - ): - - super().__init__() - self.replay_buffer = ReplayBuffer.copy_from_path( - zarr_path, keys=['img', 'state', 'action']) - val_mask = get_val_mask( - n_episodes=self.replay_buffer.n_episodes, - val_ratio=val_ratio, - seed=seed) - train_mask = ~val_mask - train_mask = downsample_mask( - mask=train_mask, - max_n=max_train_episodes, - seed=seed) - - self.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=horizon, - pad_before=pad_before, - pad_after=pad_after, - episode_mask=train_mask) - self.train_mask = train_mask - self.horizon = horizon - self.pad_before = pad_before - self.pad_after = pad_after - - def get_validation_dataset(self): - val_set = copy.copy(self) - val_set.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=self.horizon, - pad_before=self.pad_before, - pad_after=self.pad_after, - episode_mask=~self.train_mask - ) - val_set.train_mask = ~self.train_mask - return val_set - - def get_normalizer(self, mode='limits', **kwargs): - data = { - 'action': self.replay_buffer['action'], - 'agent_pos': self.replay_buffer['state'][...,:2] - } - normalizer = LinearNormalizer() - normalizer.fit(data=data, last_n_dims=1, mode=mode, **kwargs) - normalizer['image'] = get_image_range_normalizer() - return normalizer - - def __len__(self) -> int: - return len(self.sampler) - - def _sample_to_data(self, sample): - agent_pos = sample['state'][:,:2].astype(np.float32) # (agent_posx2, block_posex3) - image = np.moveaxis(sample['img'],-1,1)/255 - - data = { - 'obs': { - 'image': image, # T, 3, 96, 96 - 'agent_pos': agent_pos, # T, 2 - }, - 'action': sample['action'].astype(np.float32) # T, 2 - } - return data - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - sample = self.sampler.sample_sequence(idx) - data = self._sample_to_data(sample) - torch_data = dict_apply(data, torch.from_numpy) - return torch_data - - -def test(): - import os - zarr_path = os.path.expanduser('~/dev/diffusion_policy/data/pusht/pusht_cchi_v7_replay.zarr') - dataset = PushTImageDataset(zarr_path, horizon=16) - - # from matplotlib import pyplot as plt - # normalizer = dataset.get_normalizer() - # nactions = normalizer['action'].normalize(dataset.replay_buffer['action']) - # diff = np.diff(nactions, axis=0) - # dists = np.linalg.norm(np.diff(nactions, axis=0), axis=-1) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/real_pusht_image_dataset.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/real_pusht_image_dataset.py deleted file mode 100644 index c5a238f7f..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/real_pusht_image_dataset.py +++ /dev/null @@ -1,291 +0,0 @@ -from typing import Dict, List -import torch -import numpy as np -import zarr -import os -import shutil -from filelock import FileLock -from threadpoolctl import threadpool_limits -from omegaconf import OmegaConf -import cv2 -import json -import hashlib -import copy -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.dataset.base_dataset import BaseImageDataset -from diffusion_policy.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer -from diffusion_policy.common.replay_buffer import ReplayBuffer -from diffusion_policy.common.sampler import ( - SequenceSampler, get_val_mask, downsample_mask) -from diffusion_policy.real_world.real_data_conversion import real_data_to_replay_buffer -from diffusion_policy.common.normalize_util import ( - get_range_normalizer_from_stat, - get_image_range_normalizer, - get_identity_normalizer_from_stat, - array_to_stats -) - -class RealPushTImageDataset(BaseImageDataset): - def __init__(self, - shape_meta: dict, - dataset_path: str, - horizon=1, - pad_before=0, - pad_after=0, - n_obs_steps=None, - n_latency_steps=0, - use_cache=False, - seed=42, - val_ratio=0.0, - max_train_episodes=None, - delta_action=False, - ): - assert os.path.isdir(dataset_path) - - replay_buffer = None - if use_cache: - # fingerprint shape_meta - shape_meta_json = json.dumps(OmegaConf.to_container(shape_meta), sort_keys=True) - shape_meta_hash = hashlib.md5(shape_meta_json.encode('utf-8')).hexdigest() - cache_zarr_path = os.path.join(dataset_path, shape_meta_hash + '.zarr.zip') - cache_lock_path = cache_zarr_path + '.lock' - print('Acquiring lock on cache.') - with FileLock(cache_lock_path): - if not os.path.exists(cache_zarr_path): - # cache does not exists - try: - print('Cache does not exist. Creating!') - replay_buffer = _get_replay_buffer( - dataset_path=dataset_path, - shape_meta=shape_meta, - store=zarr.MemoryStore() - ) - print('Saving cache to disk.') - with zarr.ZipStore(cache_zarr_path) as zip_store: - replay_buffer.save_to_store( - store=zip_store - ) - except Exception as e: - shutil.rmtree(cache_zarr_path) - raise e - else: - print('Loading cached ReplayBuffer from Disk.') - with zarr.ZipStore(cache_zarr_path, mode='r') as zip_store: - replay_buffer = ReplayBuffer.copy_from_store( - src_store=zip_store, store=zarr.MemoryStore()) - print('Loaded!') - else: - replay_buffer = _get_replay_buffer( - dataset_path=dataset_path, - shape_meta=shape_meta, - store=zarr.MemoryStore() - ) - - if delta_action: - # replace action as relative to previous frame - actions = replay_buffer['action'][:] - # support positions only at this time - assert actions.shape[1] <= 3 - actions_diff = np.zeros_like(actions) - episode_ends = replay_buffer.episode_ends[:] - for i in range(len(episode_ends)): - start = 0 - if i > 0: - start = episode_ends[i-1] - end = episode_ends[i] - # delta action is the difference between previous desired position and the current - # it should be scheduled at the previous timestep for the current timestep - # to ensure consistency with positional mode - actions_diff[start+1:end] = np.diff(actions[start:end], axis=0) - replay_buffer['action'][:] = actions_diff - - rgb_keys = list() - lowdim_keys = list() - obs_shape_meta = shape_meta['obs'] - for key, attr in obs_shape_meta.items(): - type = attr.get('type', 'low_dim') - if type == 'rgb': - rgb_keys.append(key) - elif type == 'low_dim': - lowdim_keys.append(key) - - key_first_k = dict() - if n_obs_steps is not None: - # only take first k obs from images - for key in rgb_keys + lowdim_keys: - key_first_k[key] = n_obs_steps - - val_mask = get_val_mask( - n_episodes=replay_buffer.n_episodes, - val_ratio=val_ratio, - seed=seed) - train_mask = ~val_mask - train_mask = downsample_mask( - mask=train_mask, - max_n=max_train_episodes, - seed=seed) - - sampler = SequenceSampler( - replay_buffer=replay_buffer, - sequence_length=horizon+n_latency_steps, - pad_before=pad_before, - pad_after=pad_after, - episode_mask=train_mask, - key_first_k=key_first_k) - - self.replay_buffer = replay_buffer - self.sampler = sampler - self.shape_meta = shape_meta - self.rgb_keys = rgb_keys - self.lowdim_keys = lowdim_keys - self.n_obs_steps = n_obs_steps - self.val_mask = val_mask - self.horizon = horizon - self.n_latency_steps = n_latency_steps - self.pad_before = pad_before - self.pad_after = pad_after - - def get_validation_dataset(self): - val_set = copy.copy(self) - val_set.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=self.horizon+self.n_latency_steps, - pad_before=self.pad_before, - pad_after=self.pad_after, - episode_mask=self.val_mask - ) - val_set.val_mask = ~self.val_mask - return val_set - - def get_normalizer(self, **kwargs) -> LinearNormalizer: - normalizer = LinearNormalizer() - - # action - normalizer['action'] = SingleFieldLinearNormalizer.create_fit( - self.replay_buffer['action']) - - # obs - for key in self.lowdim_keys: - normalizer[key] = SingleFieldLinearNormalizer.create_fit( - self.replay_buffer[key]) - - # image - for key in self.rgb_keys: - normalizer[key] = get_image_range_normalizer() - return normalizer - - def get_all_actions(self) -> torch.Tensor: - return torch.from_numpy(self.replay_buffer['action']) - - def __len__(self): - return len(self.sampler) - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - threadpool_limits(1) - data = self.sampler.sample_sequence(idx) - - # to save RAM, only return first n_obs_steps of OBS - # since the rest will be discarded anyway. - # when self.n_obs_steps is None - # this slice does nothing (takes all) - T_slice = slice(self.n_obs_steps) - - obs_dict = dict() - for key in self.rgb_keys: - # move channel last to channel first - # T,H,W,C - # convert uint8 image to float32 - obs_dict[key] = np.moveaxis(data[key][T_slice],-1,1 - ).astype(np.float32) / 255. - # T,C,H,W - # save ram - del data[key] - for key in self.lowdim_keys: - obs_dict[key] = data[key][T_slice].astype(np.float32) - # save ram - del data[key] - - action = data['action'].astype(np.float32) - # handle latency by dropping first n_latency_steps action - # observations are already taken care of by T_slice - if self.n_latency_steps > 0: - action = action[self.n_latency_steps:] - - torch_data = { - 'obs': dict_apply(obs_dict, torch.from_numpy), - 'action': torch.from_numpy(action) - } - return torch_data - -def zarr_resize_index_last_dim(zarr_arr, idxs): - actions = zarr_arr[:] - actions = actions[...,idxs] - zarr_arr.resize(zarr_arr.shape[:-1] + (len(idxs),)) - zarr_arr[:] = actions - return zarr_arr - -def _get_replay_buffer(dataset_path, shape_meta, store): - # parse shape meta - rgb_keys = list() - lowdim_keys = list() - out_resolutions = dict() - lowdim_shapes = dict() - obs_shape_meta = shape_meta['obs'] - for key, attr in obs_shape_meta.items(): - type = attr.get('type', 'low_dim') - shape = tuple(attr.get('shape')) - if type == 'rgb': - rgb_keys.append(key) - c,h,w = shape - out_resolutions[key] = (w,h) - elif type == 'low_dim': - lowdim_keys.append(key) - lowdim_shapes[key] = tuple(shape) - if 'pose' in key: - assert tuple(shape) in [(2,),(6,)] - - action_shape = tuple(shape_meta['action']['shape']) - assert action_shape in [(2,),(6,)] - - # load data - cv2.setNumThreads(1) - with threadpool_limits(1): - replay_buffer = real_data_to_replay_buffer( - dataset_path=dataset_path, - out_store=store, - out_resolutions=out_resolutions, - lowdim_keys=lowdim_keys + ['action'], - image_keys=rgb_keys - ) - - # transform lowdim dimensions - if action_shape == (2,): - # 2D action space, only controls X and Y - zarr_arr = replay_buffer['action'] - zarr_resize_index_last_dim(zarr_arr, idxs=[0,1]) - - for key, shape in lowdim_shapes.items(): - if 'pose' in key and shape == (2,): - # only take X and Y - zarr_arr = replay_buffer[key] - zarr_resize_index_last_dim(zarr_arr, idxs=[0,1]) - - return replay_buffer - - -def test(): - import hydra - from omegaconf import OmegaConf - OmegaConf.register_new_resolver("eval", eval, replace=True) - - with hydra.initialize('../diffusion_policy/config'): - cfg = hydra.compose('train_robomimic_real_image_workspace') - OmegaConf.resolve(cfg) - dataset = hydra.utils.instantiate(cfg.task.dataset) - - from matplotlib import pyplot as plt - normalizer = dataset.get_normalizer() - nactions = normalizer['action'].normalize(dataset.replay_buffer['action'][:]) - diff = np.diff(nactions, axis=0) - dists = np.linalg.norm(np.diff(nactions, axis=0), axis=-1) - _ = plt.hist(dists, bins=100); plt.title('real action velocity') diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/robomimic_replay_image_dataset.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/robomimic_replay_image_dataset.py deleted file mode 100644 index 2728e9e9e..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/robomimic_replay_image_dataset.py +++ /dev/null @@ -1,373 +0,0 @@ -from typing import Dict, List -import torch -import numpy as np -import h5py -from tqdm import tqdm -import zarr -import os -import shutil -import copy -import json -import hashlib -from filelock import FileLock -from threadpoolctl import threadpool_limits -import concurrent.futures -import multiprocessing -from omegaconf import OmegaConf -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.dataset.base_dataset import BaseImageDataset, LinearNormalizer -from diffusion_policy.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer -from diffusion_policy.model.common.rotation_transformer import RotationTransformer -from diffusion_policy.codecs.imagecodecs_numcodecs import register_codecs, Jpeg2k -from diffusion_policy.common.replay_buffer import ReplayBuffer -from diffusion_policy.common.sampler import SequenceSampler, get_val_mask -from diffusion_policy.common.normalize_util import ( - robomimic_abs_action_only_normalizer_from_stat, - robomimic_abs_action_only_dual_arm_normalizer_from_stat, - get_range_normalizer_from_stat, - get_image_range_normalizer, - get_identity_normalizer_from_stat, - array_to_stats -) -register_codecs() - -class RobomimicReplayImageDataset(BaseImageDataset): - def __init__(self, - shape_meta: dict, - dataset_path: str, - horizon=1, - pad_before=0, - pad_after=0, - n_obs_steps=None, - abs_action=False, - rotation_rep='rotation_6d', # ignored when abs_action=False - use_legacy_normalizer=False, - use_cache=False, - seed=42, - val_ratio=0.0 - ): - rotation_transformer = RotationTransformer( - from_rep='axis_angle', to_rep=rotation_rep) - - replay_buffer = None - if use_cache: - cache_zarr_path = dataset_path + '.zarr.zip' - cache_lock_path = cache_zarr_path + '.lock' - print('Acquiring lock on cache.') - with FileLock(cache_lock_path): - if not os.path.exists(cache_zarr_path): - # cache does not exists - try: - print('Cache does not exist. Creating!') - # store = zarr.DirectoryStore(cache_zarr_path) - replay_buffer = _convert_robomimic_to_replay( - store=zarr.MemoryStore(), - shape_meta=shape_meta, - dataset_path=dataset_path, - abs_action=abs_action, - rotation_transformer=rotation_transformer) - print('Saving cache to disk.') - with zarr.ZipStore(cache_zarr_path) as zip_store: - replay_buffer.save_to_store( - store=zip_store - ) - except Exception as e: - shutil.rmtree(cache_zarr_path) - raise e - else: - print('Loading cached ReplayBuffer from Disk.') - with zarr.ZipStore(cache_zarr_path, mode='r') as zip_store: - replay_buffer = ReplayBuffer.copy_from_store( - src_store=zip_store, store=zarr.MemoryStore()) - print('Loaded!') - else: - replay_buffer = _convert_robomimic_to_replay( - store=zarr.MemoryStore(), - shape_meta=shape_meta, - dataset_path=dataset_path, - abs_action=abs_action, - rotation_transformer=rotation_transformer) - - rgb_keys = list() - lowdim_keys = list() - obs_shape_meta = shape_meta['obs'] - for key, attr in obs_shape_meta.items(): - type = attr.get('type', 'low_dim') - if type == 'rgb': - rgb_keys.append(key) - elif type == 'low_dim': - lowdim_keys.append(key) - - # for key in rgb_keys: - # replay_buffer[key].compressor.numthreads=1 - - key_first_k = dict() - if n_obs_steps is not None: - # only take first k obs from images - for key in rgb_keys + lowdim_keys: - key_first_k[key] = n_obs_steps - - val_mask = get_val_mask( - n_episodes=replay_buffer.n_episodes, - val_ratio=val_ratio, - seed=seed) - train_mask = ~val_mask - sampler = SequenceSampler( - replay_buffer=replay_buffer, - sequence_length=horizon, - pad_before=pad_before, - pad_after=pad_after, - episode_mask=train_mask, - key_first_k=key_first_k) - - self.replay_buffer = replay_buffer - self.sampler = sampler - self.shape_meta = shape_meta - self.rgb_keys = rgb_keys - self.lowdim_keys = lowdim_keys - self.abs_action = abs_action - self.n_obs_steps = n_obs_steps - self.train_mask = train_mask - self.horizon = horizon - self.pad_before = pad_before - self.pad_after = pad_after - self.use_legacy_normalizer = use_legacy_normalizer - - def get_validation_dataset(self): - val_set = copy.copy(self) - val_set.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=self.horizon, - pad_before=self.pad_before, - pad_after=self.pad_after, - episode_mask=~self.train_mask - ) - val_set.train_mask = ~self.train_mask - return val_set - - def get_normalizer(self, **kwargs) -> LinearNormalizer: - normalizer = LinearNormalizer() - - # action - stat = array_to_stats(self.replay_buffer['action']) - if self.abs_action: - if stat['mean'].shape[-1] > 10: - # dual arm - this_normalizer = robomimic_abs_action_only_dual_arm_normalizer_from_stat(stat) - else: - this_normalizer = robomimic_abs_action_only_normalizer_from_stat(stat) - - if self.use_legacy_normalizer: - this_normalizer = normalizer_from_stat(stat) - else: - # already normalized - this_normalizer = get_identity_normalizer_from_stat(stat) - normalizer['action'] = this_normalizer - - # obs - for key in self.lowdim_keys: - stat = array_to_stats(self.replay_buffer[key]) - - if key.endswith('pos'): - this_normalizer = get_range_normalizer_from_stat(stat) - elif key.endswith('quat'): - # quaternion is in [-1,1] already - this_normalizer = get_identity_normalizer_from_stat(stat) - elif key.endswith('qpos'): - this_normalizer = get_range_normalizer_from_stat(stat) - else: - raise RuntimeError('unsupported') - normalizer[key] = this_normalizer - - # image - for key in self.rgb_keys: - normalizer[key] = get_image_range_normalizer() - return normalizer - - def get_all_actions(self) -> torch.Tensor: - return torch.from_numpy(self.replay_buffer['action']) - - def __len__(self): - return len(self.sampler) - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - threadpool_limits(1) - data = self.sampler.sample_sequence(idx) - - # to save RAM, only return first n_obs_steps of OBS - # since the rest will be discarded anyway. - # when self.n_obs_steps is None - # this slice does nothing (takes all) - T_slice = slice(self.n_obs_steps) - - obs_dict = dict() - for key in self.rgb_keys: - # move channel last to channel first - # T,H,W,C - # convert uint8 image to float32 - obs_dict[key] = np.moveaxis(data[key][T_slice],-1,1 - ).astype(np.float32) / 255. - # T,C,H,W - del data[key] - for key in self.lowdim_keys: - obs_dict[key] = data[key][T_slice].astype(np.float32) - del data[key] - - torch_data = { - 'obs': dict_apply(obs_dict, torch.from_numpy), - 'action': torch.from_numpy(data['action'].astype(np.float32)) - } - return torch_data - - -def _convert_actions(raw_actions, abs_action, rotation_transformer): - actions = raw_actions - if abs_action: - is_dual_arm = False - if raw_actions.shape[-1] == 14: - # dual arm - raw_actions = raw_actions.reshape(-1,2,7) - is_dual_arm = True - - pos = raw_actions[...,:3] - rot = raw_actions[...,3:6] - gripper = raw_actions[...,6:] - rot = rotation_transformer.forward(rot) - raw_actions = np.concatenate([ - pos, rot, gripper - ], axis=-1).astype(np.float32) - - if is_dual_arm: - raw_actions = raw_actions.reshape(-1,20) - actions = raw_actions - return actions - - -def _convert_robomimic_to_replay(store, shape_meta, dataset_path, abs_action, rotation_transformer, - n_workers=None, max_inflight_tasks=None): - if n_workers is None: - n_workers = multiprocessing.cpu_count() - if max_inflight_tasks is None: - max_inflight_tasks = n_workers * 5 - - # parse shape_meta - rgb_keys = list() - lowdim_keys = list() - # construct compressors and chunks - obs_shape_meta = shape_meta['obs'] - for key, attr in obs_shape_meta.items(): - shape = attr['shape'] - type = attr.get('type', 'low_dim') - if type == 'rgb': - rgb_keys.append(key) - elif type == 'low_dim': - lowdim_keys.append(key) - - root = zarr.group(store) - data_group = root.require_group('data', overwrite=True) - meta_group = root.require_group('meta', overwrite=True) - - with h5py.File(dataset_path) as file: - # count total steps - demos = file['data'] - episode_ends = list() - prev_end = 0 - for i in range(len(demos)): - demo = demos[f'demo_{i}'] - episode_length = demo['actions'].shape[0] - episode_end = prev_end + episode_length - prev_end = episode_end - episode_ends.append(episode_end) - n_steps = episode_ends[-1] - episode_starts = [0] + episode_ends[:-1] - _ = meta_group.array('episode_ends', episode_ends, - dtype=np.int64, compressor=None, overwrite=True) - - # save lowdim data - for key in tqdm(lowdim_keys + ['action'], desc="Loading lowdim data"): - data_key = 'obs/' + key - if key == 'action': - data_key = 'actions' - this_data = list() - for i in range(len(demos)): - demo = demos[f'demo_{i}'] - this_data.append(demo[data_key][:].astype(np.float32)) - this_data = np.concatenate(this_data, axis=0) - if key == 'action': - this_data = _convert_actions( - raw_actions=this_data, - abs_action=abs_action, - rotation_transformer=rotation_transformer - ) - assert this_data.shape == (n_steps,) + tuple(shape_meta['action']['shape']) - else: - assert this_data.shape == (n_steps,) + tuple(shape_meta['obs'][key]['shape']) - _ = data_group.array( - name=key, - data=this_data, - shape=this_data.shape, - chunks=this_data.shape, - compressor=None, - dtype=this_data.dtype - ) - - def img_copy(zarr_arr, zarr_idx, hdf5_arr, hdf5_idx): - try: - zarr_arr[zarr_idx] = hdf5_arr[hdf5_idx] - # make sure we can successfully decode - _ = zarr_arr[zarr_idx] - return True - except Exception as e: - return False - - with tqdm(total=n_steps*len(rgb_keys), desc="Loading image data", mininterval=1.0) as pbar: - # one chunk per thread, therefore no synchronization needed - with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as executor: - futures = set() - for key in rgb_keys: - data_key = 'obs/' + key - shape = tuple(shape_meta['obs'][key]['shape']) - c,h,w = shape - this_compressor = Jpeg2k(level=50) - img_arr = data_group.require_dataset( - name=key, - shape=(n_steps,h,w,c), - chunks=(1,h,w,c), - compressor=this_compressor, - dtype=np.uint8 - ) - for episode_idx in range(len(demos)): - demo = demos[f'demo_{episode_idx}'] - hdf5_arr = demo['obs'][key] - for hdf5_idx in range(hdf5_arr.shape[0]): - if len(futures) >= max_inflight_tasks: - # limit number of inflight tasks - completed, futures = concurrent.futures.wait(futures, - return_when=concurrent.futures.FIRST_COMPLETED) - for f in completed: - if not f.result(): - raise RuntimeError('Failed to encode image!') - pbar.update(len(completed)) - - zarr_idx = episode_starts[episode_idx] + hdf5_idx - futures.add( - executor.submit(img_copy, - img_arr, zarr_idx, hdf5_arr, hdf5_idx)) - completed, futures = concurrent.futures.wait(futures) - for f in completed: - if not f.result(): - raise RuntimeError('Failed to encode image!') - pbar.update(len(completed)) - - replay_buffer = ReplayBuffer(root) - return replay_buffer - -def normalizer_from_stat(stat): - max_abs = np.maximum(stat['max'].max(), np.abs(stat['min']).max()) - scale = np.full_like(stat['max'], fill_value=1/max_abs) - offset = np.zeros_like(stat['max']) - return SingleFieldLinearNormalizer.create_manual( - scale=scale, - offset=offset, - input_stats_dict=stat - ) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/robomimic_replay_lowdim_dataset.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/robomimic_replay_lowdim_dataset.py deleted file mode 100644 index c642da9f2..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/dataset/robomimic_replay_lowdim_dataset.py +++ /dev/null @@ -1,168 +0,0 @@ -from typing import Dict, List -import torch -import numpy as np -import h5py -from tqdm import tqdm -import copy -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.dataset.base_dataset import BaseLowdimDataset, LinearNormalizer -from diffusion_policy.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer -from diffusion_policy.model.common.rotation_transformer import RotationTransformer -from diffusion_policy.common.replay_buffer import ReplayBuffer -from diffusion_policy.common.sampler import ( - SequenceSampler, get_val_mask, downsample_mask) -from diffusion_policy.common.normalize_util import ( - robomimic_abs_action_only_normalizer_from_stat, - robomimic_abs_action_only_dual_arm_normalizer_from_stat, - get_identity_normalizer_from_stat, - array_to_stats -) - -class RobomimicReplayLowdimDataset(BaseLowdimDataset): - def __init__(self, - dataset_path: str, - horizon=1, - pad_before=0, - pad_after=0, - obs_keys: List[str]=[ - 'object', - 'robot0_eef_pos', - 'robot0_eef_quat', - 'robot0_gripper_qpos'], - abs_action=False, - rotation_rep='rotation_6d', - use_legacy_normalizer=False, - seed=42, - val_ratio=0.0, - max_train_episodes=None - ): - obs_keys = list(obs_keys) - rotation_transformer = RotationTransformer( - from_rep='axis_angle', to_rep=rotation_rep) - - replay_buffer = ReplayBuffer.create_empty_numpy() - with h5py.File(dataset_path) as file: - demos = file['data'] - for i in tqdm(range(len(demos)), desc="Loading hdf5 to ReplayBuffer"): - demo = demos[f'demo_{i}'] - episode = _data_to_obs( - raw_obs=demo['obs'], - raw_actions=demo['actions'][:].astype(np.float32), - obs_keys=obs_keys, - abs_action=abs_action, - rotation_transformer=rotation_transformer) - replay_buffer.add_episode(episode) - - val_mask = get_val_mask( - n_episodes=replay_buffer.n_episodes, - val_ratio=val_ratio, - seed=seed) - train_mask = ~val_mask - train_mask = downsample_mask( - mask=train_mask, - max_n=max_train_episodes, - seed=seed) - - sampler = SequenceSampler( - replay_buffer=replay_buffer, - sequence_length=horizon, - pad_before=pad_before, - pad_after=pad_after, - episode_mask=train_mask) - - self.replay_buffer = replay_buffer - self.sampler = sampler - self.abs_action = abs_action - self.train_mask = train_mask - self.horizon = horizon - self.pad_before = pad_before - self.pad_after = pad_after - self.use_legacy_normalizer = use_legacy_normalizer - - def get_validation_dataset(self): - val_set = copy.copy(self) - val_set.sampler = SequenceSampler( - replay_buffer=self.replay_buffer, - sequence_length=self.horizon, - pad_before=self.pad_before, - pad_after=self.pad_after, - episode_mask=~self.train_mask - ) - val_set.train_mask = ~self.train_mask - return val_set - - def get_normalizer(self, **kwargs) -> LinearNormalizer: - normalizer = LinearNormalizer() - - # action - stat = array_to_stats(self.replay_buffer['action']) - if self.abs_action: - if stat['mean'].shape[-1] > 10: - # dual arm - this_normalizer = robomimic_abs_action_only_dual_arm_normalizer_from_stat(stat) - else: - this_normalizer = robomimic_abs_action_only_normalizer_from_stat(stat) - - if self.use_legacy_normalizer: - this_normalizer = normalizer_from_stat(stat) - else: - # already normalized - this_normalizer = get_identity_normalizer_from_stat(stat) - normalizer['action'] = this_normalizer - - # aggregate obs stats - obs_stat = array_to_stats(self.replay_buffer['obs']) - - - normalizer['obs'] = normalizer_from_stat(obs_stat) - return normalizer - - def get_all_actions(self) -> torch.Tensor: - return torch.from_numpy(self.replay_buffer['action']) - - def __len__(self): - return len(self.sampler) - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - data = self.sampler.sample_sequence(idx) - torch_data = dict_apply(data, torch.from_numpy) - return torch_data - -def normalizer_from_stat(stat): - max_abs = np.maximum(stat['max'].max(), np.abs(stat['min']).max()) - scale = np.full_like(stat['max'], fill_value=1/max_abs) - offset = np.zeros_like(stat['max']) - return SingleFieldLinearNormalizer.create_manual( - scale=scale, - offset=offset, - input_stats_dict=stat - ) - -def _data_to_obs(raw_obs, raw_actions, obs_keys, abs_action, rotation_transformer): - obs = np.concatenate([ - raw_obs[key] for key in obs_keys - ], axis=-1).astype(np.float32) - - if abs_action: - is_dual_arm = False - if raw_actions.shape[-1] == 14: - # dual arm - raw_actions = raw_actions.reshape(-1,2,7) - is_dual_arm = True - - pos = raw_actions[...,:3] - rot = raw_actions[...,3:6] - gripper = raw_actions[...,6:] - rot = rotation_transformer.forward(rot) - raw_actions = np.concatenate([ - pos, rot, gripper - ], axis=-1).astype(np.float32) - - if is_dual_arm: - raw_actions = raw_actions.reshape(-1,20) - - data = { - 'obs': obs, - 'action': raw_actions - } - return data diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/base_image_runner.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/base_image_runner.py deleted file mode 100644 index 065200716..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/base_image_runner.py +++ /dev/null @@ -1,9 +0,0 @@ -from typing import Dict -from diffusion_policy.policy.base_image_policy import BaseImagePolicy - -class BaseImageRunner: - def __init__(self, output_dir): - self.output_dir = output_dir - - def run(self, policy: BaseImagePolicy) -> Dict: - raise NotImplementedError() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/base_lowdim_runner.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/base_lowdim_runner.py deleted file mode 100644 index 45437ec88..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/base_lowdim_runner.py +++ /dev/null @@ -1,9 +0,0 @@ -from typing import Dict -from diffusion_policy.policy.base_lowdim_policy import BaseLowdimPolicy - -class BaseLowdimRunner: - def __init__(self, output_dir): - self.output_dir = output_dir - - def run(self, policy: BaseLowdimPolicy) -> Dict: - raise NotImplementedError() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/blockpush_lowdim_runner.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/blockpush_lowdim_runner.py deleted file mode 100644 index 5193dafbb..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/blockpush_lowdim_runner.py +++ /dev/null @@ -1,293 +0,0 @@ -import wandb -import numpy as np -import torch -import collections -import pathlib -import tqdm -import dill -import math -import wandb.sdk.data_types.video as wv -from diffusion_policy.env.block_pushing.block_pushing_multimodal import BlockPushMultimodal -from diffusion_policy.gym_util.async_vector_env import AsyncVectorEnv -from diffusion_policy.gym_util.sync_vector_env import SyncVectorEnv -from diffusion_policy.gym_util.multistep_wrapper import MultiStepWrapper -from diffusion_policy.gym_util.video_recording_wrapper import VideoRecordingWrapper, VideoRecorder -from gym.wrappers import FlattenObservation - -from diffusion_policy.policy.base_lowdim_policy import BaseLowdimPolicy -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner - -class BlockPushLowdimRunner(BaseLowdimRunner): - def __init__(self, - output_dir, - n_train=10, - n_train_vis=3, - train_start_seed=0, - n_test=22, - n_test_vis=6, - test_start_seed=10000, - max_steps=200, - n_obs_steps=8, - n_action_steps=8, - fps=5, - crf=22, - past_action=False, - abs_action=False, - obs_eef_target=True, - tqdm_interval_sec=5.0, - n_envs=None - ): - super().__init__(output_dir) - - if n_envs is None: - n_envs = n_train + n_test - - task_fps = 10 - steps_per_render = max(10 // fps, 1) - - def env_fn(): - return MultiStepWrapper( - VideoRecordingWrapper( - FlattenObservation( - BlockPushMultimodal( - control_frequency=task_fps, - shared_memory=False, - seed=seed, - abs_action=abs_action - ) - ), - video_recoder=VideoRecorder.create_h264( - fps=fps, - codec='h264', - input_pix_fmt='rgb24', - crf=crf, - thread_type='FRAME', - thread_count=1 - ), - file_path=None, - steps_per_render=steps_per_render - ), - n_obs_steps=n_obs_steps, - n_action_steps=n_action_steps, - max_episode_steps=max_steps - ) - - env_fns = [env_fn] * n_envs - env_seeds = list() - env_prefixs = list() - env_init_fn_dills = list() - # train - for i in range(n_train): - seed = train_start_seed + i - enable_render = i < n_train_vis - - def init_fn(env, seed=seed, enable_render=enable_render): - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # set seed - assert isinstance(env, MultiStepWrapper) - env.seed(seed) - - env_seeds.append(seed) - env_prefixs.append('train/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - # test - for i in range(n_test): - seed = test_start_seed + i - enable_render = i < n_test_vis - - def init_fn(env, seed=seed, enable_render=enable_render): - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # set seed - assert isinstance(env, MultiStepWrapper) - env.seed(seed) - - env_seeds.append(seed) - env_prefixs.append('test/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - env = AsyncVectorEnv(env_fns) - # env = SyncVectorEnv(env_fns) - - self.env = env - self.env_fns = env_fns - self.env_seeds = env_seeds - self.env_prefixs = env_prefixs - self.env_init_fn_dills = env_init_fn_dills - self.fps = fps - self.crf = crf - self.n_obs_steps = n_obs_steps - self.n_action_steps = n_action_steps - self.past_action = past_action - self.max_steps = max_steps - self.tqdm_interval_sec = tqdm_interval_sec - self.obs_eef_target = obs_eef_target - - - def run(self, policy: BaseLowdimPolicy): - device = policy.device - dtype = policy.dtype - env = self.env - - # plan for rollout - n_envs = len(self.env_fns) - n_inits = len(self.env_init_fn_dills) - n_chunks = math.ceil(n_inits / n_envs) - - # allocate data - all_video_paths = [None] * n_inits - all_rewards = [None] * n_inits - last_info = [None] * n_inits - - for chunk_idx in range(n_chunks): - start = chunk_idx * n_envs - end = min(n_inits, start + n_envs) - this_global_slice = slice(start, end) - this_n_active_envs = end - start - this_local_slice = slice(0,this_n_active_envs) - - this_init_fns = self.env_init_fn_dills[this_global_slice] - n_diff = n_envs - len(this_init_fns) - if n_diff > 0: - this_init_fns.extend([self.env_init_fn_dills[0]]*n_diff) - assert len(this_init_fns) == n_envs - - # init envs - env.call_each('run_dill_function', - args_list=[(x,) for x in this_init_fns]) - - # start rollout - obs = env.reset() - past_action = None - policy.reset() - - pbar = tqdm.tqdm(total=self.max_steps, desc=f"Eval BlockPushLowdimRunner {chunk_idx+1}/{n_chunks}", - leave=False, mininterval=self.tqdm_interval_sec) - done = False - while not done: - # create obs dict - if not self.obs_eef_target: - obs[...,8:10] = 0 - np_obs_dict = { - 'obs': obs.astype(np.float32) - } - if self.past_action and (past_action is not None): - # TODO: not tested - np_obs_dict['past_action'] = past_action[ - :,-(self.n_obs_steps-1):].astype(np.float32) - # device transfer - obs_dict = dict_apply(np_obs_dict, - lambda x: torch.from_numpy(x).to( - device=device)) - - # run policy - with torch.no_grad(): - action_dict = policy.predict_action(obs_dict) - - # device_transfer - np_action_dict = dict_apply(action_dict, - lambda x: x.detach().to('cpu').numpy()) - - action = np_action_dict['action'] - - # step env - obs, reward, done, info = env.step(action) - done = np.all(done) - past_action = action - - # update pbar - pbar.update(action.shape[1]) - pbar.close() - - # collect data for this round - all_video_paths[this_global_slice] = env.render()[this_local_slice] - all_rewards[this_global_slice] = env.call('get_attr', 'reward')[this_local_slice] - last_info[this_global_slice] = [dict((k,v[-1]) for k, v in x.items()) for x in info][this_local_slice] - - # log - total_rewards = collections.defaultdict(list) - total_p1 = collections.defaultdict(list) - total_p2 = collections.defaultdict(list) - prefix_event_counts = collections.defaultdict(lambda :collections.defaultdict(lambda : 0)) - prefix_counts = collections.defaultdict(lambda : 0) - - log_data = dict() - # results reported in the paper are generated using the commented out line below - # which will only report and average metrics from first n_envs initial condition and seeds - # fortunately this won't invalidate our conclusion since - # 1. This bug only affects the variance of metrics, not their mean - # 2. All baseline methods are evaluated using the same code - # to completely reproduce reported numbers, uncomment this line: - # for i in range(len(self.env_fns)): - # and comment out this line - for i in range(n_inits): - seed = self.env_seeds[i] - prefix = self.env_prefixs[i] - this_rewards = all_rewards[i] - total_reward = np.unique(this_rewards).sum() # (0, 0.49, 0.51) - p1 = total_reward > 0.4 - p2 = total_reward > 0.9 - - total_rewards[prefix].append(total_reward) - total_p1[prefix].append(p1) - total_p2[prefix].append(p2) - log_data[prefix+f'sim_max_reward_{seed}'] = total_reward - - # aggregate event counts - prefix_counts[prefix] += 1 - for key, value in last_info[i].items(): - delta_count = 1 if value > 0 else 0 - prefix_event_counts[prefix][key] += delta_count - - # visualize sim - video_path = all_video_paths[i] - if video_path is not None: - sim_video = wandb.Video(video_path) - log_data[prefix+f'sim_video_{seed}'] = sim_video - - # log aggregate metrics - for prefix, value in total_rewards.items(): - name = prefix+'mean_score' - value = np.mean(value) - log_data[name] = value - for prefix, value in total_p1.items(): - name = prefix+'p1' - value = np.mean(value) - log_data[name] = value - for prefix, value in total_p2.items(): - name = prefix+'p2' - value = np.mean(value) - log_data[name] = value - - # summarize probabilities - for prefix, events in prefix_event_counts.items(): - prefix_count = prefix_counts[prefix] - for event, count in events.items(): - prob = count / prefix_count - key = prefix + event - log_data[key] = prob - - return log_data diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/kitchen_lowdim_runner.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/kitchen_lowdim_runner.py deleted file mode 100644 index 39f562333..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/kitchen_lowdim_runner.py +++ /dev/null @@ -1,319 +0,0 @@ -import wandb -import numpy as np -import torch -import collections -import pathlib -import tqdm -import dill -import math -import logging -import wandb.sdk.data_types.video as wv -import gym -import gym.spaces -import multiprocessing as mp -from diffusion_policy.gym_util.async_vector_env import AsyncVectorEnv -from diffusion_policy.gym_util.sync_vector_env import SyncVectorEnv -from diffusion_policy.gym_util.multistep_wrapper import MultiStepWrapper -from diffusion_policy.gym_util.video_recording_wrapper import VideoRecordingWrapper, VideoRecorder - -from diffusion_policy.policy.base_lowdim_policy import BaseLowdimPolicy -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner - -module_logger = logging.getLogger(__name__) - -class KitchenLowdimRunner(BaseLowdimRunner): - def __init__(self, - output_dir, - dataset_dir, - n_train=10, - n_train_vis=3, - train_start_seed=0, - n_test=22, - n_test_vis=6, - test_start_seed=10000, - max_steps=280, - n_obs_steps=2, - n_action_steps=8, - render_hw=(240,360), - fps=12.5, - crf=22, - past_action=False, - tqdm_interval_sec=5.0, - abs_action=False, - robot_noise_ratio=0.1, - n_envs=None - ): - super().__init__(output_dir) - - if n_envs is None: - n_envs = n_train + n_test - - task_fps = 12.5 - steps_per_render = int(max(task_fps // fps, 1)) - - def env_fn(): - from diffusion_policy.env.kitchen.v0 import KitchenAllV0 - from diffusion_policy.env.kitchen.kitchen_lowdim_wrapper import KitchenLowdimWrapper - env = KitchenAllV0(use_abs_action=abs_action) - env.robot_noise_ratio = robot_noise_ratio - return MultiStepWrapper( - VideoRecordingWrapper( - KitchenLowdimWrapper( - env=env, - init_qpos=None, - init_qvel=None, - render_hw=tuple(render_hw) - ), - video_recoder=VideoRecorder.create_h264( - fps=fps, - codec='h264', - input_pix_fmt='rgb24', - crf=crf, - thread_type='FRAME', - thread_count=1 - ), - file_path=None, - steps_per_render=steps_per_render - ), - n_obs_steps=n_obs_steps, - n_action_steps=n_action_steps, - max_episode_steps=max_steps - ) - - all_init_qpos = np.load(pathlib.Path(dataset_dir) / "all_init_qpos.npy") - all_init_qvel = np.load(pathlib.Path(dataset_dir) / "all_init_qvel.npy") - module_logger.info(f'Loaded {len(all_init_qpos)} known initial conditions.') - - env_fns = [env_fn] * n_envs - env_seeds = list() - env_prefixs = list() - env_init_fn_dills = list() - # train - for i in range(n_train): - seed = train_start_seed + i - enable_render = i < n_train_vis - init_qpos = None - init_qvel = None - if i < len(all_init_qpos): - init_qpos = all_init_qpos[i] - init_qvel = all_init_qvel[i] - - def init_fn(env, init_qpos=init_qpos, init_qvel=init_qvel, enable_render=enable_render): - from diffusion_policy.env.kitchen.kitchen_lowdim_wrapper import KitchenLowdimWrapper - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # set initial condition - assert isinstance(env.env.env, KitchenLowdimWrapper) - env.env.env.init_qpos = init_qpos - env.env.env.init_qvel = init_qvel - - env_seeds.append(seed) - env_prefixs.append('train/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - # test - for i in range(n_test): - seed = test_start_seed + i - enable_render = i < n_test_vis - - def init_fn(env, seed=seed, enable_render=enable_render): - from diffusion_policy.env.kitchen.kitchen_lowdim_wrapper import KitchenLowdimWrapper - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # set initial condition - assert isinstance(env.env.env, KitchenLowdimWrapper) - env.env.env.init_qpos = None - env.env.env.init_qvel = None - - # set seed - assert isinstance(env, MultiStepWrapper) - env.seed(seed) - - env_seeds.append(seed) - env_prefixs.append('test/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - def dummy_env_fn(): - # Avoid importing or using env in the main process - # to prevent OpenGL context issue with fork. - # Create a fake env whose sole purpos is to provide - # obs/action spaces and metadata. - env = gym.Env() - env.observation_space = gym.spaces.Box( - -8, 8, shape=(60,), dtype=np.float32) - env.action_space = gym.spaces.Box( - -8, 8, shape=(9,), dtype=np.float32) - env.metadata = { - 'render.modes': ['human', 'rgb_array', 'depth_array'], - 'video.frames_per_second': 12 - } - env = MultiStepWrapper( - env=env, - n_obs_steps=n_obs_steps, - n_action_steps=n_action_steps, - max_episode_steps=max_steps - ) - return env - - env = AsyncVectorEnv(env_fns, dummy_env_fn=dummy_env_fn) - # env = SyncVectorEnv(env_fns) - - self.env = env - self.env_fns = env_fns - self.env_seeds = env_seeds - self.env_prefixs = env_prefixs - self.env_init_fn_dills = env_init_fn_dills - self.fps = fps - self.crf = crf - self.n_obs_steps = n_obs_steps - self.n_action_steps = n_action_steps - self.past_action = past_action - self.max_steps = max_steps - self.tqdm_interval_sec = tqdm_interval_sec - - - def run(self, policy: BaseLowdimPolicy): - device = policy.device - dtype = policy.dtype - env = self.env - - # plan for rollout - n_envs = len(self.env_fns) - n_inits = len(self.env_init_fn_dills) - n_chunks = math.ceil(n_inits / n_envs) - - # allocate data - all_video_paths = [None] * n_inits - all_rewards = [None] * n_inits - last_info = [None] * n_inits - - for chunk_idx in range(n_chunks): - start = chunk_idx * n_envs - end = min(n_inits, start + n_envs) - this_global_slice = slice(start, end) - this_n_active_envs = end - start - this_local_slice = slice(0,this_n_active_envs) - - this_init_fns = self.env_init_fn_dills[this_global_slice] - n_diff = n_envs - len(this_init_fns) - if n_diff > 0: - this_init_fns.extend([self.env_init_fn_dills[0]]*n_diff) - assert len(this_init_fns) == n_envs - - # init envs - env.call_each('run_dill_function', - args_list=[(x,) for x in this_init_fns]) - - # start rollout - obs = env.reset() - past_action = None - policy.reset() - - pbar = tqdm.tqdm(total=self.max_steps, desc=f"Eval BlockPushLowdimRunner {chunk_idx+1}/{n_chunks}", - leave=False, mininterval=self.tqdm_interval_sec) - done = False - while not done: - # create obs dict - np_obs_dict = { - 'obs': obs.astype(np.float32) - } - if self.past_action and (past_action is not None): - # TODO: not tested - np_obs_dict['past_action'] = past_action[ - :,-(self.n_obs_steps-1):].astype(np.float32) - # device transfer - obs_dict = dict_apply(np_obs_dict, - lambda x: torch.from_numpy(x).to( - device=device)) - - # run policy - with torch.no_grad(): - action_dict = policy.predict_action(obs_dict) - - # device_transfer - np_action_dict = dict_apply(action_dict, - lambda x: x.detach().to('cpu').numpy()) - - action = np_action_dict['action'] - - # step env - obs, reward, done, info = env.step(action) - done = np.all(done) - past_action = action - - # update pbar - pbar.update(action.shape[1]) - pbar.close() - - # collect data for this round - all_video_paths[this_global_slice] = env.render()[this_local_slice] - all_rewards[this_global_slice] = env.call('get_attr', 'reward')[this_local_slice] - last_info[this_global_slice] = [dict((k,v[-1]) for k, v in x.items()) for x in info][this_local_slice] - - # reward is number of tasks completed, max 7 - # use info to record the order of task completion? - # also report the probably to completing n tasks (different aggregation of reward). - - # log - log_data = dict() - prefix_total_reward_map = collections.defaultdict(list) - prefix_n_completed_map = collections.defaultdict(list) - # results reported in the paper are generated using the commented out line below - # which will only report and average metrics from first n_envs initial condition and seeds - # fortunately this won't invalidate our conclusion since - # 1. This bug only affects the variance of metrics, not their mean - # 2. All baseline methods are evaluated using the same code - # to completely reproduce reported numbers, uncomment this line: - # for i in range(len(self.env_fns)): - # and comment out this line - for i in range(n_inits): - seed = self.env_seeds[i] - prefix = self.env_prefixs[i] - this_rewards = all_rewards[i] - total_reward = np.sum(this_rewards) / 7 - prefix_total_reward_map[prefix].append(total_reward) - - n_completed_tasks = len(last_info[i]['completed_tasks']) - prefix_n_completed_map[prefix].append(n_completed_tasks) - - # visualize sim - video_path = all_video_paths[i] - if video_path is not None: - sim_video = wandb.Video(video_path) - log_data[prefix+f'sim_video_{seed}'] = sim_video - - # log aggregate metrics - for prefix, value in prefix_total_reward_map.items(): - name = prefix+'mean_score' - value = np.mean(value) - log_data[name] = value - for prefix, value in prefix_n_completed_map.items(): - n_completed = np.array(value) - for i in range(7): - n = i + 1 - p_n = np.mean(n_completed >= n) - name = prefix + f'p_{n}' - log_data[name] = p_n - - return log_data diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/pusht_image_runner.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/pusht_image_runner.py deleted file mode 100644 index f65c06a81..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/pusht_image_runner.py +++ /dev/null @@ -1,251 +0,0 @@ -import wandb -import numpy as np -import torch -import collections -import pathlib -import tqdm -import dill -import math -import wandb.sdk.data_types.video as wv -from diffusion_policy.env.pusht.pusht_image_env import PushTImageEnv -from diffusion_policy.gym_util.async_vector_env import AsyncVectorEnv -# from diffusion_policy.gym_util.sync_vector_env import SyncVectorEnv -from diffusion_policy.gym_util.multistep_wrapper import MultiStepWrapper -from diffusion_policy.gym_util.video_recording_wrapper import VideoRecordingWrapper, VideoRecorder - -from diffusion_policy.policy.base_image_policy import BaseImagePolicy -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.env_runner.base_image_runner import BaseImageRunner - -class PushTImageRunner(BaseImageRunner): - def __init__(self, - output_dir, - n_train=10, - n_train_vis=3, - train_start_seed=0, - n_test=22, - n_test_vis=6, - legacy_test=False, - test_start_seed=10000, - max_steps=200, - n_obs_steps=8, - n_action_steps=8, - fps=10, - crf=22, - render_size=96, - past_action=False, - tqdm_interval_sec=5.0, - n_envs=None - ): - super().__init__(output_dir) - if n_envs is None: - n_envs = n_train + n_test - - steps_per_render = max(10 // fps, 1) - def env_fn(): - return MultiStepWrapper( - VideoRecordingWrapper( - PushTImageEnv( - legacy=legacy_test, - render_size=render_size - ), - video_recoder=VideoRecorder.create_h264( - fps=fps, - codec='h264', - input_pix_fmt='rgb24', - crf=crf, - thread_type='FRAME', - thread_count=1 - ), - file_path=None, - steps_per_render=steps_per_render - ), - n_obs_steps=n_obs_steps, - n_action_steps=n_action_steps, - max_episode_steps=max_steps - ) - - env_fns = [env_fn] * n_envs - env_seeds = list() - env_prefixs = list() - env_init_fn_dills = list() - # train - for i in range(n_train): - seed = train_start_seed + i - enable_render = i < n_train_vis - - def init_fn(env, seed=seed, enable_render=enable_render): - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # set seed - assert isinstance(env, MultiStepWrapper) - env.seed(seed) - - env_seeds.append(seed) - env_prefixs.append('train/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - # test - for i in range(n_test): - seed = test_start_seed + i - enable_render = i < n_test_vis - - def init_fn(env, seed=seed, enable_render=enable_render): - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # set seed - assert isinstance(env, MultiStepWrapper) - env.seed(seed) - - env_seeds.append(seed) - env_prefixs.append('test/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - env = AsyncVectorEnv(env_fns) - - # test env - # env.reset(seed=env_seeds) - # x = env.step(env.action_space.sample()) - # imgs = env.call('render') - # import pdb; pdb.set_trace() - - self.env = env - self.env_fns = env_fns - self.env_seeds = env_seeds - self.env_prefixs = env_prefixs - self.env_init_fn_dills = env_init_fn_dills - self.fps = fps - self.crf = crf - self.n_obs_steps = n_obs_steps - self.n_action_steps = n_action_steps - self.past_action = past_action - self.max_steps = max_steps - self.tqdm_interval_sec = tqdm_interval_sec - - def run(self, policy: BaseImagePolicy): - device = policy.device - dtype = policy.dtype - env = self.env - - # plan for rollout - n_envs = len(self.env_fns) - n_inits = len(self.env_init_fn_dills) - n_chunks = math.ceil(n_inits / n_envs) - - # allocate data - all_video_paths = [None] * n_inits - all_rewards = [None] * n_inits - - for chunk_idx in range(n_chunks): - start = chunk_idx * n_envs - end = min(n_inits, start + n_envs) - this_global_slice = slice(start, end) - this_n_active_envs = end - start - this_local_slice = slice(0,this_n_active_envs) - - this_init_fns = self.env_init_fn_dills[this_global_slice] - n_diff = n_envs - len(this_init_fns) - if n_diff > 0: - this_init_fns.extend([self.env_init_fn_dills[0]]*n_diff) - assert len(this_init_fns) == n_envs - - # init envs - env.call_each('run_dill_function', - args_list=[(x,) for x in this_init_fns]) - - # start rollout - obs = env.reset() - past_action = None - policy.reset() - - pbar = tqdm.tqdm(total=self.max_steps, desc=f"Eval PushtImageRunner {chunk_idx+1}/{n_chunks}", - leave=False, mininterval=self.tqdm_interval_sec) - done = False - while not done: - # create obs dict - np_obs_dict = dict(obs) - if self.past_action and (past_action is not None): - # TODO: not tested - np_obs_dict['past_action'] = past_action[ - :,-(self.n_obs_steps-1):].astype(np.float32) - - # device transfer - obs_dict = dict_apply(np_obs_dict, - lambda x: torch.from_numpy(x).to( - device=device)) - - # run policy - with torch.no_grad(): - action_dict = policy.predict_action(obs_dict) - - # device_transfer - np_action_dict = dict_apply(action_dict, - lambda x: x.detach().to('cpu').numpy()) - - action = np_action_dict['action'] - - # step env - obs, reward, done, info = env.step(action) - done = np.all(done) - past_action = action - - # update pbar - pbar.update(action.shape[1]) - pbar.close() - - all_video_paths[this_global_slice] = env.render()[this_local_slice] - all_rewards[this_global_slice] = env.call('get_attr', 'reward')[this_local_slice] - # clear out video buffer - _ = env.reset() - - # log - max_rewards = collections.defaultdict(list) - log_data = dict() - # results reported in the paper are generated using the commented out line below - # which will only report and average metrics from first n_envs initial condition and seeds - # fortunately this won't invalidate our conclusion since - # 1. This bug only affects the variance of metrics, not their mean - # 2. All baseline methods are evaluated using the same code - # to completely reproduce reported numbers, uncomment this line: - # for i in range(len(self.env_fns)): - # and comment out this line - for i in range(n_inits): - seed = self.env_seeds[i] - prefix = self.env_prefixs[i] - max_reward = np.max(all_rewards[i]) - max_rewards[prefix].append(max_reward) - log_data[prefix+f'sim_max_reward_{seed}'] = max_reward - - # visualize sim - video_path = all_video_paths[i] - if video_path is not None: - sim_video = wandb.Video(video_path) - log_data[prefix+f'sim_video_{seed}'] = sim_video - - # log aggregate metrics - for prefix, value in max_rewards.items(): - name = prefix+'mean_score' - value = np.mean(value) - log_data[name] = value - - return log_data diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/pusht_keypoints_runner.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/pusht_keypoints_runner.py deleted file mode 100644 index a16bd58ed..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/pusht_keypoints_runner.py +++ /dev/null @@ -1,273 +0,0 @@ -import wandb -import numpy as np -import torch -import collections -import pathlib -import tqdm -import dill -import math -import wandb.sdk.data_types.video as wv -from diffusion_policy.env.pusht.pusht_keypoints_env import PushTKeypointsEnv -from diffusion_policy.gym_util.async_vector_env import AsyncVectorEnv -# from diffusion_policy.gym_util.sync_vector_env import SyncVectorEnv -from diffusion_policy.gym_util.multistep_wrapper import MultiStepWrapper -from diffusion_policy.gym_util.video_recording_wrapper import VideoRecordingWrapper, VideoRecorder - -from diffusion_policy.policy.base_lowdim_policy import BaseLowdimPolicy -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner - -class PushTKeypointsRunner(BaseLowdimRunner): - def __init__(self, - output_dir, - keypoint_visible_rate=1.0, - n_train=10, - n_train_vis=3, - train_start_seed=0, - n_test=22, - n_test_vis=6, - legacy_test=False, - test_start_seed=10000, - max_steps=200, - n_obs_steps=8, - n_action_steps=8, - n_latency_steps=0, - fps=10, - crf=22, - agent_keypoints=False, - past_action=False, - tqdm_interval_sec=5.0, - n_envs=None - ): - super().__init__(output_dir) - - if n_envs is None: - n_envs = n_train + n_test - - # handle latency step - # to mimic latency, we request n_latency_steps additional steps - # of past observations, and the discard the last n_latency_steps - env_n_obs_steps = n_obs_steps + n_latency_steps - env_n_action_steps = n_action_steps - - # assert n_obs_steps <= n_action_steps - kp_kwargs = PushTKeypointsEnv.genenerate_keypoint_manager_params() - - def env_fn(): - return MultiStepWrapper( - VideoRecordingWrapper( - PushTKeypointsEnv( - legacy=legacy_test, - keypoint_visible_rate=keypoint_visible_rate, - agent_keypoints=agent_keypoints, - **kp_kwargs - ), - video_recoder=VideoRecorder.create_h264( - fps=fps, - codec='h264', - input_pix_fmt='rgb24', - crf=crf, - thread_type='FRAME', - thread_count=1 - ), - file_path=None, - ), - n_obs_steps=env_n_obs_steps, - n_action_steps=env_n_action_steps, - max_episode_steps=max_steps - ) - - env_fns = [env_fn] * n_envs - env_seeds = list() - env_prefixs = list() - env_init_fn_dills = list() - # train - for i in range(n_train): - seed = train_start_seed + i - enable_render = i < n_train_vis - - def init_fn(env, seed=seed, enable_render=enable_render): - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # set seed - assert isinstance(env, MultiStepWrapper) - env.seed(seed) - - env_seeds.append(seed) - env_prefixs.append('train/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - # test - for i in range(n_test): - seed = test_start_seed + i - enable_render = i < n_test_vis - - def init_fn(env, seed=seed, enable_render=enable_render): - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # set seed - assert isinstance(env, MultiStepWrapper) - env.seed(seed) - - env_seeds.append(seed) - env_prefixs.append('test/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - env = AsyncVectorEnv(env_fns) - - # test env - # env.reset(seed=env_seeds) - # x = env.step(env.action_space.sample()) - # imgs = env.call('render') - # import pdb; pdb.set_trace() - - self.env = env - self.env_fns = env_fns - self.env_seeds = env_seeds - self.env_prefixs = env_prefixs - self.env_init_fn_dills = env_init_fn_dills - self.fps = fps - self.crf = crf - self.agent_keypoints = agent_keypoints - self.n_obs_steps = n_obs_steps - self.n_action_steps = n_action_steps - self.n_latency_steps = n_latency_steps - self.past_action = past_action - self.max_steps = max_steps - self.tqdm_interval_sec = tqdm_interval_sec - - def run(self, policy: BaseLowdimPolicy): - device = policy.device - dtype = policy.dtype - - env = self.env - - # plan for rollout - n_envs = len(self.env_fns) - n_inits = len(self.env_init_fn_dills) - n_chunks = math.ceil(n_inits / n_envs) - - # allocate data - all_video_paths = [None] * n_inits - all_rewards = [None] * n_inits - - for chunk_idx in range(n_chunks): - start = chunk_idx * n_envs - end = min(n_inits, start + n_envs) - this_global_slice = slice(start, end) - this_n_active_envs = end - start - this_local_slice = slice(0,this_n_active_envs) - - this_init_fns = self.env_init_fn_dills[this_global_slice] - n_diff = n_envs - len(this_init_fns) - if n_diff > 0: - this_init_fns.extend([self.env_init_fn_dills[0]]*n_diff) - assert len(this_init_fns) == n_envs - - # init envs - env.call_each('run_dill_function', - args_list=[(x,) for x in this_init_fns]) - - # start rollout - obs = env.reset() - past_action = None - policy.reset() - - pbar = tqdm.tqdm(total=self.max_steps, desc=f"Eval PushtKeypointsRunner {chunk_idx+1}/{n_chunks}", - leave=False, mininterval=self.tqdm_interval_sec) - done = False - while not done: - Do = obs.shape[-1] // 2 - # create obs dict - np_obs_dict = { - # handle n_latency_steps by discarding the last n_latency_steps - 'obs': obs[...,:self.n_obs_steps,:Do].astype(np.float32), - 'obs_mask': obs[...,:self.n_obs_steps,Do:] > 0.5 - } - if self.past_action and (past_action is not None): - # TODO: not tested - np_obs_dict['past_action'] = past_action[ - :,-(self.n_obs_steps-1):].astype(np.float32) - - # device transfer - obs_dict = dict_apply(np_obs_dict, - lambda x: torch.from_numpy(x).to( - device=device)) - - # run policy - with torch.no_grad(): - action_dict = policy.predict_action(obs_dict) - - # device_transfer - np_action_dict = dict_apply(action_dict, - lambda x: x.detach().to('cpu').numpy()) - - # handle latency_steps, we discard the first n_latency_steps actions - # to simulate latency - action = np_action_dict['action'][:,self.n_latency_steps:] - - # step env - obs, reward, done, info = env.step(action) - done = np.all(done) - past_action = action - - # update pbar - pbar.update(action.shape[1]) - pbar.close() - - # collect data for this round - all_video_paths[this_global_slice] = env.render()[this_local_slice] - all_rewards[this_global_slice] = env.call('get_attr', 'reward')[this_local_slice] - # import pdb; pdb.set_trace() - - # log - max_rewards = collections.defaultdict(list) - log_data = dict() - # results reported in the paper are generated using the commented out line below - # which will only report and average metrics from first n_envs initial condition and seeds - # fortunately this won't invalidate our conclusion since - # 1. This bug only affects the variance of metrics, not their mean - # 2. All baseline methods are evaluated using the same code - # to completely reproduce reported numbers, uncomment this line: - # for i in range(len(self.env_fns)): - # and comment out this line - for i in range(n_inits): - seed = self.env_seeds[i] - prefix = self.env_prefixs[i] - max_reward = np.max(all_rewards[i]) - max_rewards[prefix].append(max_reward) - log_data[prefix+f'sim_max_reward_{seed}'] = max_reward - - # visualize sim - video_path = all_video_paths[i] - if video_path is not None: - sim_video = wandb.Video(video_path) - log_data[prefix+f'sim_video_{seed}'] = sim_video - - # log aggregate metrics - for prefix, value in max_rewards.items(): - name = prefix+'mean_score' - value = np.mean(value) - log_data[name] = value - - return log_data diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/real_pusht_image_runner.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/real_pusht_image_runner.py deleted file mode 100644 index 3b58780ce..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/real_pusht_image_runner.py +++ /dev/null @@ -1,10 +0,0 @@ -from diffusion_policy.policy.base_image_policy import BaseImagePolicy -from diffusion_policy.env_runner.base_image_runner import BaseImageRunner - -class RealPushTImageRunner(BaseImageRunner): - def __init__(self, - output_dir): - super().__init__(output_dir) - - def run(self, policy: BaseImagePolicy): - return dict() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/robomimic_image_runner.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/robomimic_image_runner.py deleted file mode 100644 index cbae74e86..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/robomimic_image_runner.py +++ /dev/null @@ -1,375 +0,0 @@ -import os -import wandb -import numpy as np -import torch -import collections -import pathlib -import tqdm -import h5py -import math -import dill -import wandb.sdk.data_types.video as wv -from diffusion_policy.gym_util.async_vector_env import AsyncVectorEnv -from diffusion_policy.gym_util.sync_vector_env import SyncVectorEnv -from diffusion_policy.gym_util.multistep_wrapper import MultiStepWrapper -from diffusion_policy.gym_util.video_recording_wrapper import VideoRecordingWrapper, VideoRecorder -from diffusion_policy.model.common.rotation_transformer import RotationTransformer - -from diffusion_policy.policy.base_image_policy import BaseImagePolicy -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.env_runner.base_image_runner import BaseImageRunner -from diffusion_policy.env.robomimic.robomimic_image_wrapper import RobomimicImageWrapper -import robomimic.utils.file_utils as FileUtils -import robomimic.utils.env_utils as EnvUtils -import robomimic.utils.obs_utils as ObsUtils - - -def create_env(env_meta, shape_meta, enable_render=True): - modality_mapping = collections.defaultdict(list) - for key, attr in shape_meta['obs'].items(): - modality_mapping[attr.get('type', 'low_dim')].append(key) - ObsUtils.initialize_obs_modality_mapping_from_dict(modality_mapping) - - env = EnvUtils.create_env_from_metadata( - env_meta=env_meta, - render=False, - render_offscreen=enable_render, - use_image_obs=enable_render, - ) - return env - - -class RobomimicImageRunner(BaseImageRunner): - """ - Robomimic envs already enforces number of steps. - """ - - def __init__(self, - output_dir, - dataset_path, - shape_meta:dict, - n_train=10, - n_train_vis=3, - train_start_idx=0, - n_test=22, - n_test_vis=6, - test_start_seed=10000, - max_steps=400, - n_obs_steps=2, - n_action_steps=8, - render_obs_key='agentview_image', - fps=10, - crf=22, - past_action=False, - abs_action=False, - tqdm_interval_sec=5.0, - n_envs=None - ): - super().__init__(output_dir) - - if n_envs is None: - n_envs = n_train + n_test - - # assert n_obs_steps <= n_action_steps - dataset_path = os.path.expanduser(dataset_path) - robosuite_fps = 20 - steps_per_render = max(robosuite_fps // fps, 1) - - # read from dataset - env_meta = FileUtils.get_env_metadata_from_dataset( - dataset_path) - # disable object state observation - env_meta['env_kwargs']['use_object_obs'] = False - - rotation_transformer = None - if abs_action: - env_meta['env_kwargs']['controller_configs']['control_delta'] = False - rotation_transformer = RotationTransformer('axis_angle', 'rotation_6d') - - def env_fn(): - robomimic_env = create_env( - env_meta=env_meta, - shape_meta=shape_meta - ) - # Robosuite's hard reset causes excessive memory consumption. - # Disabled to run more envs. - # https://github.com/ARISE-Initiative/robosuite/blob/92abf5595eddb3a845cd1093703e5a3ccd01e77e/robosuite/environments/base.py#L247-L248 - robomimic_env.env.hard_reset = False - return MultiStepWrapper( - VideoRecordingWrapper( - RobomimicImageWrapper( - env=robomimic_env, - shape_meta=shape_meta, - init_state=None, - render_obs_key=render_obs_key - ), - video_recoder=VideoRecorder.create_h264( - fps=fps, - codec='h264', - input_pix_fmt='rgb24', - crf=crf, - thread_type='FRAME', - thread_count=1 - ), - file_path=None, - steps_per_render=steps_per_render - ), - n_obs_steps=n_obs_steps, - n_action_steps=n_action_steps, - max_episode_steps=max_steps - ) - - # For each process the OpenGL context can only be initialized once - # Since AsyncVectorEnv uses fork to create worker process, - # a separate env_fn that does not create OpenGL context (enable_render=False) - # is needed to initialize spaces. - def dummy_env_fn(): - robomimic_env = create_env( - env_meta=env_meta, - shape_meta=shape_meta, - enable_render=False - ) - return MultiStepWrapper( - VideoRecordingWrapper( - RobomimicImageWrapper( - env=robomimic_env, - shape_meta=shape_meta, - init_state=None, - render_obs_key=render_obs_key - ), - video_recoder=VideoRecorder.create_h264( - fps=fps, - codec='h264', - input_pix_fmt='rgb24', - crf=crf, - thread_type='FRAME', - thread_count=1 - ), - file_path=None, - steps_per_render=steps_per_render - ), - n_obs_steps=n_obs_steps, - n_action_steps=n_action_steps, - max_episode_steps=max_steps - ) - - env_fns = [env_fn] * n_envs - env_seeds = list() - env_prefixs = list() - env_init_fn_dills = list() - - # train - with h5py.File(dataset_path, 'r') as f: - for i in range(n_train): - train_idx = train_start_idx + i - enable_render = i < n_train_vis - init_state = f[f'data/demo_{train_idx}/states'][0] - - def init_fn(env, init_state=init_state, - enable_render=enable_render): - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # switch to init_state reset - assert isinstance(env.env.env, RobomimicImageWrapper) - env.env.env.init_state = init_state - - env_seeds.append(train_idx) - env_prefixs.append('train/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - # test - for i in range(n_test): - seed = test_start_seed + i - enable_render = i < n_test_vis - - def init_fn(env, seed=seed, - enable_render=enable_render): - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # switch to seed reset - assert isinstance(env.env.env, RobomimicImageWrapper) - env.env.env.init_state = None - env.seed(seed) - - env_seeds.append(seed) - env_prefixs.append('test/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - env = AsyncVectorEnv(env_fns, dummy_env_fn=dummy_env_fn) - # env = SyncVectorEnv(env_fns) - - - self.env_meta = env_meta - self.env = env - self.env_fns = env_fns - self.env_seeds = env_seeds - self.env_prefixs = env_prefixs - self.env_init_fn_dills = env_init_fn_dills - self.fps = fps - self.crf = crf - self.n_obs_steps = n_obs_steps - self.n_action_steps = n_action_steps - self.past_action = past_action - self.max_steps = max_steps - self.rotation_transformer = rotation_transformer - self.abs_action = abs_action - self.tqdm_interval_sec = tqdm_interval_sec - - def run(self, policy: BaseImagePolicy): - device = policy.device - dtype = policy.dtype - env = self.env - - # plan for rollout - n_envs = len(self.env_fns) - n_inits = len(self.env_init_fn_dills) - n_chunks = math.ceil(n_inits / n_envs) - - # allocate data - all_video_paths = [None] * n_inits - all_rewards = [None] * n_inits - - for chunk_idx in range(n_chunks): - start = chunk_idx * n_envs - end = min(n_inits, start + n_envs) - this_global_slice = slice(start, end) - this_n_active_envs = end - start - this_local_slice = slice(0,this_n_active_envs) - - this_init_fns = self.env_init_fn_dills[this_global_slice] - n_diff = n_envs - len(this_init_fns) - if n_diff > 0: - this_init_fns.extend([self.env_init_fn_dills[0]]*n_diff) - assert len(this_init_fns) == n_envs - - # init envs - env.call_each('run_dill_function', - args_list=[(x,) for x in this_init_fns]) - - # start rollout - obs = env.reset() - past_action = None - policy.reset() - - env_name = self.env_meta['env_name'] - pbar = tqdm.tqdm(total=self.max_steps, desc=f"Eval {env_name}Image {chunk_idx+1}/{n_chunks}", - leave=False, mininterval=self.tqdm_interval_sec) - - done = False - while not done: - # create obs dict - np_obs_dict = dict(obs) - if self.past_action and (past_action is not None): - # TODO: not tested - np_obs_dict['past_action'] = past_action[ - :,-(self.n_obs_steps-1):].astype(np.float32) - - # device transfer - obs_dict = dict_apply(np_obs_dict, - lambda x: torch.from_numpy(x).to( - device=device)) - - # run policy - with torch.no_grad(): - action_dict = policy.predict_action(obs_dict) - - # device_transfer - np_action_dict = dict_apply(action_dict, - lambda x: x.detach().to('cpu').numpy()) - - action = np_action_dict['action'] - if not np.all(np.isfinite(action)): - print(action) - raise RuntimeError("Nan or Inf action") - - # step env - env_action = action - if self.abs_action: - env_action = self.undo_transform_action(action) - - obs, reward, done, info = env.step(env_action) - done = np.all(done) - past_action = action - - # update pbar - pbar.update(action.shape[1]) - pbar.close() - - # collect data for this round - all_video_paths[this_global_slice] = env.render()[this_local_slice] - all_rewards[this_global_slice] = env.call('get_attr', 'reward')[this_local_slice] - # clear out video buffer - _ = env.reset() - - # log - max_rewards = collections.defaultdict(list) - log_data = dict() - # results reported in the paper are generated using the commented out line below - # which will only report and average metrics from first n_envs initial condition and seeds - # fortunately this won't invalidate our conclusion since - # 1. This bug only affects the variance of metrics, not their mean - # 2. All baseline methods are evaluated using the same code - # to completely reproduce reported numbers, uncomment this line: - # for i in range(len(self.env_fns)): - # and comment out this line - for i in range(n_inits): - seed = self.env_seeds[i] - prefix = self.env_prefixs[i] - max_reward = np.max(all_rewards[i]) - max_rewards[prefix].append(max_reward) - log_data[prefix+f'sim_max_reward_{seed}'] = max_reward - - # visualize sim - video_path = all_video_paths[i] - if video_path is not None: - sim_video = wandb.Video(video_path) - log_data[prefix+f'sim_video_{seed}'] = sim_video - - # log aggregate metrics - for prefix, value in max_rewards.items(): - name = prefix+'mean_score' - value = np.mean(value) - log_data[name] = value - - return log_data - - def undo_transform_action(self, action): - raw_shape = action.shape - if raw_shape[-1] == 20: - # dual arm - action = action.reshape(-1,2,10) - - d_rot = action.shape[-1] - 4 - pos = action[...,:3] - rot = action[...,3:3+d_rot] - gripper = action[...,[-1]] - rot = self.rotation_transformer.inverse(rot) - uaction = np.concatenate([ - pos, rot, gripper - ], axis=-1) - - if raw_shape[-1] == 20: - # dual arm - uaction = uaction.reshape(*raw_shape[:-1], 14) - - return uaction diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/robomimic_lowdim_runner.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/robomimic_lowdim_runner.py deleted file mode 100644 index f3ba642e4..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/env_runner/robomimic_lowdim_runner.py +++ /dev/null @@ -1,368 +0,0 @@ -import os -import wandb -import numpy as np -import torch -import collections -import pathlib -import tqdm -import h5py -import dill -import math -import wandb.sdk.data_types.video as wv -from diffusion_policy.gym_util.async_vector_env import AsyncVectorEnv -# from diffusion_policy.gym_util.sync_vector_env import SyncVectorEnv -from diffusion_policy.gym_util.multistep_wrapper import MultiStepWrapper -from diffusion_policy.gym_util.video_recording_wrapper import VideoRecordingWrapper, VideoRecorder -from diffusion_policy.model.common.rotation_transformer import RotationTransformer - -from diffusion_policy.policy.base_lowdim_policy import BaseLowdimPolicy -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner -from diffusion_policy.env.robomimic.robomimic_lowdim_wrapper import RobomimicLowdimWrapper -import robomimic.utils.file_utils as FileUtils -import robomimic.utils.env_utils as EnvUtils -import robomimic.utils.obs_utils as ObsUtils - - -def create_env(env_meta, obs_keys): - ObsUtils.initialize_obs_modality_mapping_from_dict( - {'low_dim': obs_keys}) - env = EnvUtils.create_env_from_metadata( - env_meta=env_meta, - render=False, - # only way to not show collision geometry - # is to enable render_offscreen - # which uses a lot of RAM. - render_offscreen=False, - use_image_obs=False, - ) - return env - - -class RobomimicLowdimRunner(BaseLowdimRunner): - """ - Robomimic envs already enforces number of steps. - """ - - def __init__(self, - output_dir, - dataset_path, - obs_keys, - n_train=10, - n_train_vis=3, - train_start_idx=0, - n_test=22, - n_test_vis=6, - test_start_seed=10000, - max_steps=400, - n_obs_steps=2, - n_action_steps=8, - n_latency_steps=0, - render_hw=(256,256), - render_camera_name='agentview', - fps=10, - crf=22, - past_action=False, - abs_action=False, - tqdm_interval_sec=5.0, - n_envs=None - ): - """ - Assuming: - n_obs_steps=2 - n_latency_steps=3 - n_action_steps=4 - o: obs - i: inference - a: action - Batch t: - |o|o| | | | | | | - | |i|i|i| | | | | - | | | | |a|a|a|a| - Batch t+1 - | | | | |o|o| | | | | | | - | | | | | |i|i|i| | | | | - | | | | | | | | |a|a|a|a| - """ - - super().__init__(output_dir) - - if n_envs is None: - n_envs = n_train + n_test - - # handle latency step - # to mimic latency, we request n_latency_steps additional steps - # of past observations, and the discard the last n_latency_steps - env_n_obs_steps = n_obs_steps + n_latency_steps - env_n_action_steps = n_action_steps - - # assert n_obs_steps <= n_action_steps - dataset_path = os.path.expanduser(dataset_path) - robosuite_fps = 20 - steps_per_render = max(robosuite_fps // fps, 1) - - # read from dataset - env_meta = FileUtils.get_env_metadata_from_dataset( - dataset_path) - rotation_transformer = None - if abs_action: - env_meta['env_kwargs']['controller_configs']['control_delta'] = False - rotation_transformer = RotationTransformer('axis_angle', 'rotation_6d') - - def env_fn(): - robomimic_env = create_env( - env_meta=env_meta, - obs_keys=obs_keys - ) - # hard reset doesn't influence lowdim env - # robomimic_env.env.hard_reset = False - return MultiStepWrapper( - VideoRecordingWrapper( - RobomimicLowdimWrapper( - env=robomimic_env, - obs_keys=obs_keys, - init_state=None, - render_hw=render_hw, - render_camera_name=render_camera_name - ), - video_recoder=VideoRecorder.create_h264( - fps=fps, - codec='h264', - input_pix_fmt='rgb24', - crf=crf, - thread_type='FRAME', - thread_count=1 - ), - file_path=None, - steps_per_render=steps_per_render - ), - n_obs_steps=env_n_obs_steps, - n_action_steps=env_n_action_steps, - max_episode_steps=max_steps - ) - - env_fns = [env_fn] * n_envs - env_seeds = list() - env_prefixs = list() - env_init_fn_dills = list() - - # train - with h5py.File(dataset_path, 'r') as f: - for i in range(n_train): - train_idx = train_start_idx + i - enable_render = i < n_train_vis - init_state = f[f'data/demo_{train_idx}/states'][0] - - def init_fn(env, init_state=init_state, - enable_render=enable_render): - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # switch to init_state reset - assert isinstance(env.env.env, RobomimicLowdimWrapper) - env.env.env.init_state = init_state - - env_seeds.append(train_idx) - env_prefixs.append('train/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - # test - for i in range(n_test): - seed = test_start_seed + i - enable_render = i < n_test_vis - - def init_fn(env, seed=seed, - enable_render=enable_render): - # setup rendering - # video_wrapper - assert isinstance(env.env, VideoRecordingWrapper) - env.env.video_recoder.stop() - env.env.file_path = None - if enable_render: - filename = pathlib.Path(output_dir).joinpath( - 'media', wv.util.generate_id() + ".mp4") - filename.parent.mkdir(parents=False, exist_ok=True) - filename = str(filename) - env.env.file_path = filename - - # switch to seed reset - assert isinstance(env.env.env, RobomimicLowdimWrapper) - env.env.env.init_state = None - env.seed(seed) - - env_seeds.append(seed) - env_prefixs.append('test/') - env_init_fn_dills.append(dill.dumps(init_fn)) - - env = AsyncVectorEnv(env_fns) - # env = SyncVectorEnv(env_fns) - - self.env_meta = env_meta - self.env = env - self.env_fns = env_fns - self.env_seeds = env_seeds - self.env_prefixs = env_prefixs - self.env_init_fn_dills = env_init_fn_dills - self.fps = fps - self.crf = crf - self.n_obs_steps = n_obs_steps - self.n_action_steps = n_action_steps - self.n_latency_steps = n_latency_steps - self.env_n_obs_steps = env_n_obs_steps - self.env_n_action_steps = env_n_action_steps - self.past_action = past_action - self.max_steps = max_steps - self.rotation_transformer = rotation_transformer - self.abs_action = abs_action - self.tqdm_interval_sec = tqdm_interval_sec - - def run(self, policy: BaseLowdimPolicy): - device = policy.device - dtype = policy.dtype - env = self.env - - # plan for rollout - n_envs = len(self.env_fns) - n_inits = len(self.env_init_fn_dills) - n_chunks = math.ceil(n_inits / n_envs) - - # allocate data - all_video_paths = [None] * n_inits - all_rewards = [None] * n_inits - - for chunk_idx in range(n_chunks): - start = chunk_idx * n_envs - end = min(n_inits, start + n_envs) - this_global_slice = slice(start, end) - this_n_active_envs = end - start - this_local_slice = slice(0,this_n_active_envs) - - this_init_fns = self.env_init_fn_dills[this_global_slice] - n_diff = n_envs - len(this_init_fns) - if n_diff > 0: - this_init_fns.extend([self.env_init_fn_dills[0]]*n_diff) - assert len(this_init_fns) == n_envs - - # init envs - env.call_each('run_dill_function', - args_list=[(x,) for x in this_init_fns]) - - # start rollout - obs = env.reset() - past_action = None - policy.reset() - - env_name = self.env_meta['env_name'] - pbar = tqdm.tqdm(total=self.max_steps, desc=f"Eval {env_name}Lowdim {chunk_idx+1}/{n_chunks}", - leave=False, mininterval=self.tqdm_interval_sec) - - done = False - while not done: - # create obs dict - np_obs_dict = { - # handle n_latency_steps by discarding the last n_latency_steps - 'obs': obs[:,:self.n_obs_steps].astype(np.float32) - } - if self.past_action and (past_action is not None): - # TODO: not tested - np_obs_dict['past_action'] = past_action[ - :,-(self.n_obs_steps-1):].astype(np.float32) - - # device transfer - obs_dict = dict_apply(np_obs_dict, - lambda x: torch.from_numpy(x).to( - device=device)) - - # run policy - with torch.no_grad(): - action_dict = policy.predict_action(obs_dict) - - # device_transfer - np_action_dict = dict_apply(action_dict, - lambda x: x.detach().to('cpu').numpy()) - - # handle latency_steps, we discard the first n_latency_steps actions - # to simulate latency - action = np_action_dict['action'][:,self.n_latency_steps:] - if not np.all(np.isfinite(action)): - print(action) - raise RuntimeError("Nan or Inf action") - - # step env - env_action = action - if self.abs_action: - env_action = self.undo_transform_action(action) - - obs, reward, done, info = env.step(env_action) - done = np.all(done) - past_action = action - - # update pbar - pbar.update(action.shape[1]) - pbar.close() - - # collect data for this round - all_video_paths[this_global_slice] = env.render()[this_local_slice] - all_rewards[this_global_slice] = env.call('get_attr', 'reward')[this_local_slice] - - # log - max_rewards = collections.defaultdict(list) - log_data = dict() - # results reported in the paper are generated using the commented out line below - # which will only report and average metrics from first n_envs initial condition and seeds - # fortunately this won't invalidate our conclusion since - # 1. This bug only affects the variance of metrics, not their mean - # 2. All baseline methods are evaluated using the same code - # to completely reproduce reported numbers, uncomment this line: - # for i in range(len(self.env_fns)): - # and comment out this line - for i in range(n_inits): - seed = self.env_seeds[i] - prefix = self.env_prefixs[i] - max_reward = np.max(all_rewards[i]) - max_rewards[prefix].append(max_reward) - log_data[prefix+f'sim_max_reward_{seed}'] = max_reward - - # visualize sim - video_path = all_video_paths[i] - if video_path is not None: - sim_video = wandb.Video(video_path) - log_data[prefix+f'sim_video_{seed}'] = sim_video - - # log aggregate metrics - for prefix, value in max_rewards.items(): - name = prefix+'mean_score' - value = np.mean(value) - log_data[name] = value - - return log_data - - def undo_transform_action(self, action): - raw_shape = action.shape - if raw_shape[-1] == 20: - # dual arm - action = action.reshape(-1,2,10) - - d_rot = action.shape[-1] - 4 - pos = action[...,:3] - rot = action[...,3:3+d_rot] - gripper = action[...,[-1]] - rot = self.rotation_transformer.inverse(rot) - uaction = np.concatenate([ - pos, rot, gripper - ], axis=-1) - - if raw_shape[-1] == 20: - # dual arm - uaction = uaction.reshape(*raw_shape[:-1], 14) - - return uaction diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/async_vector_env.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/async_vector_env.py deleted file mode 100644 index dfb0f620a..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/async_vector_env.py +++ /dev/null @@ -1,671 +0,0 @@ -""" -Back ported methods: call, set_attr from v0.26 -Disabled auto-reset after done -Added render method. -""" - - -import numpy as np -import multiprocessing as mp -import time -import sys -from enum import Enum -from copy import deepcopy - -from gym import logger -from gym.vector.vector_env import VectorEnv -from gym.error import ( - AlreadyPendingCallError, - NoAsyncCallError, - ClosedEnvironmentError, - CustomSpaceError, -) -from gym.vector.utils import ( - create_shared_memory, - create_empty_array, - write_to_shared_memory, - read_from_shared_memory, - concatenate, - CloudpickleWrapper, - clear_mpi_env_vars, -) - -__all__ = ["AsyncVectorEnv"] - - -class AsyncState(Enum): - DEFAULT = "default" - WAITING_RESET = "reset" - WAITING_STEP = "step" - WAITING_CALL = "call" - - -class AsyncVectorEnv(VectorEnv): - """Vectorized environment that runs multiple environments in parallel. It - uses `multiprocessing` processes, and pipes for communication. - Parameters - ---------- - env_fns : iterable of callable - Functions that create the environments. - observation_space : `gym.spaces.Space` instance, optional - Observation space of a single environment. If `None`, then the - observation space of the first environment is taken. - action_space : `gym.spaces.Space` instance, optional - Action space of a single environment. If `None`, then the action space - of the first environment is taken. - shared_memory : bool (default: `True`) - If `True`, then the observations from the worker processes are - communicated back through shared variables. This can improve the - efficiency if the observations are large (e.g. images). - copy : bool (default: `True`) - If `True`, then the `reset` and `step` methods return a copy of the - observations. - context : str, optional - Context for multiprocessing. If `None`, then the default context is used. - Only available in Python 3. - daemon : bool (default: `True`) - If `True`, then subprocesses have `daemon` flag turned on; that is, they - will quit if the head process quits. However, `daemon=True` prevents - subprocesses to spawn children, so for some environments you may want - to have it set to `False` - worker : function, optional - WARNING - advanced mode option! If set, then use that worker in a subprocess - instead of a default one. Can be useful to override some inner vector env - logic, for instance, how resets on done are handled. Provides high - degree of flexibility and a high chance to shoot yourself in the foot; thus, - if you are writing your own worker, it is recommended to start from the code - for `_worker` (or `_worker_shared_memory`) method below, and add changes - """ - - def __init__( - self, - env_fns, - dummy_env_fn=None, - observation_space=None, - action_space=None, - shared_memory=True, - copy=True, - context=None, - daemon=True, - worker=None, - ): - ctx = mp.get_context(context) - self.env_fns = env_fns - self.shared_memory = shared_memory - self.copy = copy - - # Added dummy_env_fn to fix OpenGL error in Mujoco - # disable any OpenGL rendering in dummy_env_fn, since it - # will conflict with OpenGL context in the forked child process - if dummy_env_fn is None: - dummy_env_fn = env_fns[0] - dummy_env = dummy_env_fn() - self.metadata = dummy_env.metadata - - if (observation_space is None) or (action_space is None): - observation_space = observation_space or dummy_env.observation_space - action_space = action_space or dummy_env.action_space - dummy_env.close() - del dummy_env - super(AsyncVectorEnv, self).__init__( - num_envs=len(env_fns), - observation_space=observation_space, - action_space=action_space, - ) - - if self.shared_memory: - try: - _obs_buffer = create_shared_memory( - self.single_observation_space, n=self.num_envs, ctx=ctx - ) - self.observations = read_from_shared_memory( - _obs_buffer, self.single_observation_space, n=self.num_envs - ) - except CustomSpaceError: - raise ValueError( - "Using `shared_memory=True` in `AsyncVectorEnv` " - "is incompatible with non-standard Gym observation spaces " - "(i.e. custom spaces inheriting from `gym.Space`), and is " - "only compatible with default Gym spaces (e.g. `Box`, " - "`Tuple`, `Dict`) for batching. Set `shared_memory=False` " - "if you use custom observation spaces." - ) - else: - _obs_buffer = None - self.observations = create_empty_array( - self.single_observation_space, n=self.num_envs, fn=np.zeros - ) - - self.parent_pipes, self.processes = [], [] - self.error_queue = ctx.Queue() - target = _worker_shared_memory if self.shared_memory else _worker - target = worker or target - with clear_mpi_env_vars(): - for idx, env_fn in enumerate(self.env_fns): - parent_pipe, child_pipe = ctx.Pipe() - process = ctx.Process( - target=target, - name="Worker<{0}>-{1}".format(type(self).__name__, idx), - args=( - idx, - CloudpickleWrapper(env_fn), - child_pipe, - parent_pipe, - _obs_buffer, - self.error_queue, - ), - ) - - self.parent_pipes.append(parent_pipe) - self.processes.append(process) - - process.daemon = daemon - process.start() - child_pipe.close() - - self._state = AsyncState.DEFAULT - self._check_observation_spaces() - - def seed(self, seeds=None): - self._assert_is_running() - if seeds is None: - seeds = [None for _ in range(self.num_envs)] - if isinstance(seeds, int): - seeds = [seeds + i for i in range(self.num_envs)] - assert len(seeds) == self.num_envs - - if self._state != AsyncState.DEFAULT: - raise AlreadyPendingCallError( - "Calling `seed` while waiting " - "for a pending call to `{0}` to complete.".format(self._state.value), - self._state.value, - ) - - for pipe, seed in zip(self.parent_pipes, seeds): - pipe.send(("seed", seed)) - _, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) - self._raise_if_errors(successes) - - def reset_async(self): - self._assert_is_running() - if self._state != AsyncState.DEFAULT: - raise AlreadyPendingCallError( - "Calling `reset_async` while waiting " - "for a pending call to `{0}` to complete".format(self._state.value), - self._state.value, - ) - - for pipe in self.parent_pipes: - pipe.send(("reset", None)) - self._state = AsyncState.WAITING_RESET - - def reset_wait(self, timeout=None): - """ - Parameters - ---------- - timeout : int or float, optional - Number of seconds before the call to `reset_wait` times out. If - `None`, the call to `reset_wait` never times out. - Returns - ------- - observations : sample from `observation_space` - A batch of observations from the vectorized environment. - """ - self._assert_is_running() - if self._state != AsyncState.WAITING_RESET: - raise NoAsyncCallError( - "Calling `reset_wait` without any prior " "call to `reset_async`.", - AsyncState.WAITING_RESET.value, - ) - - if not self._poll(timeout): - self._state = AsyncState.DEFAULT - raise mp.TimeoutError( - "The call to `reset_wait` has timed out after " - "{0} second{1}.".format(timeout, "s" if timeout > 1 else "") - ) - - results, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) - self._raise_if_errors(successes) - self._state = AsyncState.DEFAULT - - if not self.shared_memory: - self.observations = concatenate( - results, self.observations, self.single_observation_space - ) - - return deepcopy(self.observations) if self.copy else self.observations - - def step_async(self, actions): - """ - Parameters - ---------- - actions : iterable of samples from `action_space` - List of actions. - """ - self._assert_is_running() - if self._state != AsyncState.DEFAULT: - raise AlreadyPendingCallError( - "Calling `step_async` while waiting " - "for a pending call to `{0}` to complete.".format(self._state.value), - self._state.value, - ) - - for pipe, action in zip(self.parent_pipes, actions): - pipe.send(("step", action)) - self._state = AsyncState.WAITING_STEP - - def step_wait(self, timeout=None): - """ - Parameters - ---------- - timeout : int or float, optional - Number of seconds before the call to `step_wait` times out. If - `None`, the call to `step_wait` never times out. - Returns - ------- - observations : sample from `observation_space` - A batch of observations from the vectorized environment. - rewards : `np.ndarray` instance (dtype `np.float_`) - A vector of rewards from the vectorized environment. - dones : `np.ndarray` instance (dtype `np.bool_`) - A vector whose entries indicate whether the episode has ended. - infos : list of dict - A list of auxiliary diagnostic information. - """ - self._assert_is_running() - if self._state != AsyncState.WAITING_STEP: - raise NoAsyncCallError( - "Calling `step_wait` without any prior call " "to `step_async`.", - AsyncState.WAITING_STEP.value, - ) - - if not self._poll(timeout): - self._state = AsyncState.DEFAULT - raise mp.TimeoutError( - "The call to `step_wait` has timed out after " - "{0} second{1}.".format(timeout, "s" if timeout > 1 else "") - ) - - results, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) - self._raise_if_errors(successes) - self._state = AsyncState.DEFAULT - observations_list, rewards, dones, infos = zip(*results) - - if not self.shared_memory: - self.observations = concatenate( - observations_list, self.observations, self.single_observation_space - ) - - return ( - deepcopy(self.observations) if self.copy else self.observations, - np.array(rewards), - np.array(dones, dtype=np.bool_), - infos, - ) - - def close_extras(self, timeout=None, terminate=False): - """ - Parameters - ---------- - timeout : int or float, optional - Number of seconds before the call to `close` times out. If `None`, - the call to `close` never times out. If the call to `close` times - out, then all processes are terminated. - terminate : bool (default: `False`) - If `True`, then the `close` operation is forced and all processes - are terminated. - """ - timeout = 0 if terminate else timeout - try: - if self._state != AsyncState.DEFAULT: - logger.warn( - "Calling `close` while waiting for a pending " - "call to `{0}` to complete.".format(self._state.value) - ) - function = getattr(self, "{0}_wait".format(self._state.value)) - function(timeout) - except mp.TimeoutError: - terminate = True - - if terminate: - for process in self.processes: - if process.is_alive(): - process.terminate() - else: - for pipe in self.parent_pipes: - if (pipe is not None) and (not pipe.closed): - pipe.send(("close", None)) - for pipe in self.parent_pipes: - if (pipe is not None) and (not pipe.closed): - pipe.recv() - - for pipe in self.parent_pipes: - if pipe is not None: - pipe.close() - for process in self.processes: - process.join() - - def _poll(self, timeout=None): - self._assert_is_running() - if timeout is None: - return True - end_time = time.perf_counter() + timeout - delta = None - for pipe in self.parent_pipes: - delta = max(end_time - time.perf_counter(), 0) - if pipe is None: - return False - if pipe.closed or (not pipe.poll(delta)): - return False - return True - - def _check_observation_spaces(self): - self._assert_is_running() - for pipe in self.parent_pipes: - pipe.send(("_check_observation_space", self.single_observation_space)) - same_spaces, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) - self._raise_if_errors(successes) - if not all(same_spaces): - raise RuntimeError( - "Some environments have an observation space " - "different from `{0}`. In order to batch observations, the " - "observation spaces from all environments must be " - "equal.".format(self.single_observation_space) - ) - - def _assert_is_running(self): - if self.closed: - raise ClosedEnvironmentError( - "Trying to operate on `{0}`, after a " - "call to `close()`.".format(type(self).__name__) - ) - - def _raise_if_errors(self, successes): - if all(successes): - return - - num_errors = self.num_envs - sum(successes) - assert num_errors > 0 - for _ in range(num_errors): - index, exctype, value = self.error_queue.get() - logger.error( - "Received the following error from Worker-{0}: " - "{1}: {2}".format(index, exctype.__name__, value) - ) - logger.error("Shutting down Worker-{0}.".format(index)) - self.parent_pipes[index].close() - self.parent_pipes[index] = None - - logger.error("Raising the last exception back to the main process.") - raise exctype(value) - - def call_async(self, name: str, *args, **kwargs): - """Calls the method with name asynchronously and apply args and kwargs to the method. - - Args: - name: Name of the method or property to call. - *args: Arguments to apply to the method call. - **kwargs: Keyword arguments to apply to the method call. - - Raises: - ClosedEnvironmentError: If the environment was closed (if :meth:`close` was previously called). - AlreadyPendingCallError: Calling `call_async` while waiting for a pending call to complete - """ - self._assert_is_running() - if self._state != AsyncState.DEFAULT: - raise AlreadyPendingCallError( - "Calling `call_async` while waiting " - f"for a pending call to `{self._state.value}` to complete.", - self._state.value, - ) - - for pipe in self.parent_pipes: - pipe.send(("_call", (name, args, kwargs))) - self._state = AsyncState.WAITING_CALL - - def call_wait(self, timeout = None) -> list: - """Calls all parent pipes and waits for the results. - - Args: - timeout: Number of seconds before the call to `step_wait` times out. - If `None` (default), the call to `step_wait` never times out. - - Returns: - List of the results of the individual calls to the method or property for each environment. - - Raises: - NoAsyncCallError: Calling `call_wait` without any prior call to `call_async`. - TimeoutError: The call to `call_wait` has timed out after timeout second(s). - """ - self._assert_is_running() - if self._state != AsyncState.WAITING_CALL: - raise NoAsyncCallError( - "Calling `call_wait` without any prior call to `call_async`.", - AsyncState.WAITING_CALL.value, - ) - - if not self._poll(timeout): - self._state = AsyncState.DEFAULT - raise mp.TimeoutError( - f"The call to `call_wait` has timed out after {timeout} second(s)." - ) - - results, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) - self._raise_if_errors(successes) - self._state = AsyncState.DEFAULT - - return results - - def call(self, name: str, *args, **kwargs): - """Call a method, or get a property, from each parallel environment. - - Args: - name (str): Name of the method or property to call. - *args: Arguments to apply to the method call. - **kwargs: Keyword arguments to apply to the method call. - - Returns: - List of the results of the individual calls to the method or property for each environment. - """ - self.call_async(name, *args, **kwargs) - return self.call_wait() - - - def call_each(self, name: str, - args_list: list=None, - kwargs_list: list=None, - timeout = None): - n_envs = len(self.parent_pipes) - if args_list is None: - args_list = [[]] * n_envs - assert len(args_list) == n_envs - - if kwargs_list is None: - kwargs_list = [dict()] * n_envs - assert len(kwargs_list) == n_envs - - # send - self._assert_is_running() - if self._state != AsyncState.DEFAULT: - raise AlreadyPendingCallError( - "Calling `call_async` while waiting " - f"for a pending call to `{self._state.value}` to complete.", - self._state.value, - ) - - for i, pipe in enumerate(self.parent_pipes): - pipe.send(("_call", (name, args_list[i], kwargs_list[i]))) - self._state = AsyncState.WAITING_CALL - - # receive - self._assert_is_running() - if self._state != AsyncState.WAITING_CALL: - raise NoAsyncCallError( - "Calling `call_wait` without any prior call to `call_async`.", - AsyncState.WAITING_CALL.value, - ) - - if not self._poll(timeout): - self._state = AsyncState.DEFAULT - raise mp.TimeoutError( - f"The call to `call_wait` has timed out after {timeout} second(s)." - ) - - results, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) - self._raise_if_errors(successes) - self._state = AsyncState.DEFAULT - - return results - - - def set_attr(self, name: str, values): - """Sets an attribute of the sub-environments. - - Args: - name: Name of the property to be set in each individual environment. - values: Values of the property to be set to. If ``values`` is a list or - tuple, then it corresponds to the values for each individual - environment, otherwise a single value is set for all environments. - - Raises: - ValueError: Values must be a list or tuple with length equal to the number of environments. - AlreadyPendingCallError: Calling `set_attr` while waiting for a pending call to complete. - """ - self._assert_is_running() - if not isinstance(values, (list, tuple)): - values = [values for _ in range(self.num_envs)] - if len(values) != self.num_envs: - raise ValueError( - "Values must be a list or tuple with length equal to the " - f"number of environments. Got `{len(values)}` values for " - f"{self.num_envs} environments." - ) - - if self._state != AsyncState.DEFAULT: - raise AlreadyPendingCallError( - "Calling `set_attr` while waiting " - f"for a pending call to `{self._state.value}` to complete.", - self._state.value, - ) - - for pipe, value in zip(self.parent_pipes, values): - pipe.send(("_setattr", (name, value))) - _, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) - self._raise_if_errors(successes) - - def render(self, *args, **kwargs): - return self.call('render', *args, **kwargs) - - - -def _worker(index, env_fn, pipe, parent_pipe, shared_memory, error_queue): - assert shared_memory is None - env = env_fn() - parent_pipe.close() - try: - while True: - command, data = pipe.recv() - if command == "reset": - observation = env.reset() - pipe.send((observation, True)) - elif command == "step": - observation, reward, done, info = env.step(data) - # if done: - # observation = env.reset() - pipe.send(((observation, reward, done, info), True)) - elif command == "seed": - env.seed(data) - pipe.send((None, True)) - elif command == "close": - pipe.send((None, True)) - break - elif command == "_call": - name, args, kwargs = data - if name in ["reset", "step", "seed", "close"]: - raise ValueError( - f"Trying to call function `{name}` with " - f"`_call`. Use `{name}` directly instead." - ) - function = getattr(env, name) - if callable(function): - pipe.send((function(*args, **kwargs), True)) - else: - pipe.send((function, True)) - elif command == "_setattr": - name, value = data - setattr(env, name, value) - pipe.send((None, True)) - - elif command == "_check_observation_space": - pipe.send((data == env.observation_space, True)) - else: - raise RuntimeError( - "Received unknown command `{0}`. Must " - "be one of {`reset`, `step`, `seed`, `close`, " - "`_check_observation_space`}.".format(command) - ) - except (KeyboardInterrupt, Exception): - error_queue.put((index,) + sys.exc_info()[:2]) - pipe.send((None, False)) - finally: - env.close() - - -def _worker_shared_memory(index, env_fn, pipe, parent_pipe, shared_memory, error_queue): - assert shared_memory is not None - env = env_fn() - observation_space = env.observation_space - parent_pipe.close() - try: - while True: - command, data = pipe.recv() - if command == "reset": - observation = env.reset() - write_to_shared_memory( - index, observation, shared_memory, observation_space - ) - pipe.send((None, True)) - elif command == "step": - observation, reward, done, info = env.step(data) - # if done: - # observation = env.reset() - write_to_shared_memory( - index, observation, shared_memory, observation_space - ) - pipe.send(((None, reward, done, info), True)) - elif command == "seed": - env.seed(data) - pipe.send((None, True)) - elif command == "close": - pipe.send((None, True)) - break - elif command == "_call": - name, args, kwargs = data - if name in ["reset", "step", "seed", "close"]: - raise ValueError( - f"Trying to call function `{name}` with " - f"`_call`. Use `{name}` directly instead." - ) - function = getattr(env, name) - if callable(function): - pipe.send((function(*args, **kwargs), True)) - else: - pipe.send((function, True)) - elif command == "_setattr": - name, value = data - setattr(env, name, value) - pipe.send((None, True)) - elif command == "_check_observation_space": - pipe.send((data == observation_space, True)) - else: - raise RuntimeError( - "Received unknown command `{0}`. Must " - "be one of {`reset`, `step`, `seed`, `close`, " - "`_check_observation_space`}.".format(command) - ) - except (KeyboardInterrupt, Exception): - error_queue.put((index,) + sys.exc_info()[:2]) - pipe.send((None, False)) - finally: - env.close() \ No newline at end of file diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/multistep_wrapper.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/multistep_wrapper.py deleted file mode 100644 index cf2cfdac1..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/multistep_wrapper.py +++ /dev/null @@ -1,162 +0,0 @@ -import gym -from gym import spaces -import numpy as np -from collections import defaultdict, deque -import dill - -def stack_repeated(x, n): - return np.repeat(np.expand_dims(x,axis=0),n,axis=0) - -def repeated_box(box_space, n): - return spaces.Box( - low=stack_repeated(box_space.low, n), - high=stack_repeated(box_space.high, n), - shape=(n,) + box_space.shape, - dtype=box_space.dtype - ) - -def repeated_space(space, n): - if isinstance(space, spaces.Box): - return repeated_box(space, n) - elif isinstance(space, spaces.Dict): - result_space = spaces.Dict() - for key, value in space.items(): - result_space[key] = repeated_space(value, n) - return result_space - else: - raise RuntimeError(f'Unsupported space type {type(space)}') - -def take_last_n(x, n): - x = list(x) - n = min(len(x), n) - return np.array(x[-n:]) - -def dict_take_last_n(x, n): - result = dict() - for key, value in x.items(): - result[key] = take_last_n(value, n) - return result - -def aggregate(data, method='max'): - if method == 'max': - # equivalent to any - return np.max(data) - elif method == 'min': - # equivalent to all - return np.min(data) - elif method == 'mean': - return np.mean(data) - elif method == 'sum': - return np.sum(data) - else: - raise NotImplementedError() - -def stack_last_n_obs(all_obs, n_steps): - assert(len(all_obs) > 0) - all_obs = list(all_obs) - result = np.zeros((n_steps,) + all_obs[-1].shape, - dtype=all_obs[-1].dtype) - start_idx = -min(n_steps, len(all_obs)) - result[start_idx:] = np.array(all_obs[start_idx:]) - if n_steps > len(all_obs): - # pad - result[:start_idx] = result[start_idx] - return result - - -class MultiStepWrapper(gym.Wrapper): - def __init__(self, - env, - n_obs_steps, - n_action_steps, - max_episode_steps=None, - reward_agg_method='max' - ): - super().__init__(env) - self._action_space = repeated_space(env.action_space, n_action_steps) - self._observation_space = repeated_space(env.observation_space, n_obs_steps) - self.max_episode_steps = max_episode_steps - self.n_obs_steps = n_obs_steps - self.n_action_steps = n_action_steps - self.reward_agg_method = reward_agg_method - self.n_obs_steps = n_obs_steps - - self.obs = deque(maxlen=n_obs_steps+1) - self.reward = list() - self.done = list() - self.info = defaultdict(lambda : deque(maxlen=n_obs_steps+1)) - - def reset(self): - """Resets the environment using kwargs.""" - obs = super().reset() - - self.obs = deque([obs], maxlen=self.n_obs_steps+1) - self.reward = list() - self.done = list() - self.info = defaultdict(lambda : deque(maxlen=self.n_obs_steps+1)) - - obs = self._get_obs(self.n_obs_steps) - return obs - - def step(self, action): - """ - actions: (n_action_steps,) + action_shape - """ - for act in action: - if len(self.done) > 0 and self.done[-1]: - # termination - break - observation, reward, done, info = super().step(act) - - self.obs.append(observation) - self.reward.append(reward) - if (self.max_episode_steps is not None) \ - and (len(self.reward) >= self.max_episode_steps): - # truncation - done = True - self.done.append(done) - self._add_info(info) - - observation = self._get_obs(self.n_obs_steps) - reward = aggregate(self.reward, self.reward_agg_method) - done = aggregate(self.done, 'max') - info = dict_take_last_n(self.info, self.n_obs_steps) - return observation, reward, done, info - - def _get_obs(self, n_steps=1): - """ - Output (n_steps,) + obs_shape - """ - assert(len(self.obs) > 0) - if isinstance(self.observation_space, spaces.Box): - return stack_last_n_obs(self.obs, n_steps) - elif isinstance(self.observation_space, spaces.Dict): - result = dict() - for key in self.observation_space.keys(): - result[key] = stack_last_n_obs( - [obs[key] for obs in self.obs], - n_steps - ) - return result - else: - raise RuntimeError('Unsupported space type') - - def _add_info(self, info): - for key, value in info.items(): - self.info[key].append(value) - - def get_rewards(self): - return self.reward - - def get_attr(self, name): - return getattr(self, name) - - def run_dill_function(self, dill_fn): - fn = dill.loads(dill_fn) - return fn(self) - - def get_infos(self): - result = dict() - for k, v in self.info.items(): - result[k] = list(v) - return result diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/sync_vector_env.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/sync_vector_env.py deleted file mode 100644 index c85a68223..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/sync_vector_env.py +++ /dev/null @@ -1,182 +0,0 @@ -import numpy as np -from copy import deepcopy - -from gym import logger -from gym.vector.vector_env import VectorEnv -from gym.vector.utils import concatenate, create_empty_array - -__all__ = ["SyncVectorEnv"] - - -class SyncVectorEnv(VectorEnv): - """Vectorized environment that serially runs multiple environments. - Parameters - ---------- - env_fns : iterable of callable - Functions that create the environments. - observation_space : `gym.spaces.Space` instance, optional - Observation space of a single environment. If `None`, then the - observation space of the first environment is taken. - action_space : `gym.spaces.Space` instance, optional - Action space of a single environment. If `None`, then the action space - of the first environment is taken. - copy : bool (default: `True`) - If `True`, then the `reset` and `step` methods return a copy of the - observations. - """ - - def __init__(self, env_fns, observation_space=None, action_space=None, copy=True): - self.env_fns = env_fns - self.envs = [env_fn() for env_fn in env_fns] - self.copy = copy - self.metadata = self.envs[0].metadata - - if (observation_space is None) or (action_space is None): - observation_space = observation_space or self.envs[0].observation_space - action_space = action_space or self.envs[0].action_space - super(SyncVectorEnv, self).__init__( - num_envs=len(env_fns), - observation_space=observation_space, - action_space=action_space, - ) - - self._check_observation_spaces() - self.observations = create_empty_array( - self.single_observation_space, n=self.num_envs, fn=np.zeros - ) - self._rewards = np.zeros((self.num_envs,), dtype=np.float64) - self._dones = np.zeros((self.num_envs,), dtype=np.bool_) - # self._rewards = [0] * self.num_envs - # self._dones = [False] * self.num_envs - self._actions = None - - def seed(self, seeds=None): - if seeds is None: - seeds = [None for _ in range(self.num_envs)] - if isinstance(seeds, int): - seeds = [seeds + i for i in range(self.num_envs)] - assert len(seeds) == self.num_envs - - for env, seed in zip(self.envs, seeds): - env.seed(seed) - - def reset_wait(self): - self._dones[:] = False - observations = [] - for env in self.envs: - observation = env.reset() - observations.append(observation) - self.observations = concatenate( - observations, self.observations, self.single_observation_space - ) - - return deepcopy(self.observations) if self.copy else self.observations - - def step_async(self, actions): - self._actions = actions - - def step_wait(self): - observations, infos = [], [] - for i, (env, action) in enumerate(zip(self.envs, self._actions)): - observation, self._rewards[i], self._dones[i], info = env.step(action) - # if self._dones[i]: - # observation = env.reset() - observations.append(observation) - infos.append(info) - self.observations = concatenate( - observations, self.observations, self.single_observation_space - ) - - return ( - deepcopy(self.observations) if self.copy else self.observations, - np.copy(self._rewards), - np.copy(self._dones), - infos, - ) - - def close_extras(self, **kwargs): - [env.close() for env in self.envs] - - def _check_observation_spaces(self): - for env in self.envs: - if not (env.observation_space == self.single_observation_space): - break - else: - return True - raise RuntimeError( - "Some environments have an observation space " - "different from `{0}`. In order to batch observations, the " - "observation spaces from all environments must be " - "equal.".format(self.single_observation_space) - ) - - def call(self, name, *args, **kwargs) -> tuple: - """Calls the method with name and applies args and kwargs. - - Args: - name: The method name - *args: The method args - **kwargs: The method kwargs - - Returns: - Tuple of results - """ - results = [] - for env in self.envs: - function = getattr(env, name) - if callable(function): - results.append(function(*args, **kwargs)) - else: - results.append(function) - - return tuple(results) - - def call_each(self, name: str, - args_list: list=None, - kwargs_list: list=None): - n_envs = len(self.envs) - if args_list is None: - args_list = [[]] * n_envs - assert len(args_list) == n_envs - - if kwargs_list is None: - kwargs_list = [dict()] * n_envs - assert len(kwargs_list) == n_envs - - results = [] - for i, env in enumerate(self.envs): - function = getattr(env, name) - if callable(function): - results.append(function(*args_list[i], **kwargs_list[i])) - else: - results.append(function) - - return tuple(results) - - - def render(self, *args, **kwargs): - return self.call('render', *args, **kwargs) - - def set_attr(self, name: str, values): - """Sets an attribute of the sub-environments. - - Args: - name: The property name to change - values: Values of the property to be set to. If ``values`` is a list or - tuple, then it corresponds to the values for each individual - environment, otherwise, a single value is set for all environments. - - Raises: - ValueError: Values must be a list or tuple with length equal to the number of environments. - """ - if not isinstance(values, (list, tuple)): - values = [values for _ in range(self.num_envs)] - if len(values) != self.num_envs: - raise ValueError( - "Values must be a list or tuple with length equal to the " - f"number of environments. Got `{len(values)}` values for " - f"{self.num_envs} environments." - ) - - for env, value in zip(self.envs, values): - setattr(env, name, value) \ No newline at end of file diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/video_recording_wrapper.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/video_recording_wrapper.py deleted file mode 100644 index c690f50dc..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/video_recording_wrapper.py +++ /dev/null @@ -1,51 +0,0 @@ -import gym -import numpy as np -from diffusion_policy.real_world.video_recorder import VideoRecorder - -class VideoRecordingWrapper(gym.Wrapper): - def __init__(self, - env, - video_recoder: VideoRecorder, - mode='rgb_array', - file_path=None, - steps_per_render=1, - **kwargs - ): - """ - When file_path is None, don't record. - """ - super().__init__(env) - - self.mode = mode - self.render_kwargs = kwargs - self.steps_per_render = steps_per_render - self.file_path = file_path - self.video_recoder = video_recoder - - self.step_count = 0 - - def reset(self, **kwargs): - obs = super().reset(**kwargs) - self.frames = list() - self.step_count = 1 - self.video_recoder.stop() - return obs - - def step(self, action): - result = super().step(action) - self.step_count += 1 - if self.file_path is not None \ - and ((self.step_count % self.steps_per_render) == 0): - if not self.video_recoder.is_ready(): - self.video_recoder.start(self.file_path) - - frame = self.env.render( - mode=self.mode, **self.render_kwargs) - assert frame.dtype == np.uint8 - self.video_recoder.write_frame(frame) - return result - - def render(self, mode='rgb_array', **kwargs): - if self.video_recoder.is_ready(): - self.video_recoder.stop() - return self.file_path diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/video_wrapper.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/video_wrapper.py deleted file mode 100644 index abfebbefe..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/gym_util/video_wrapper.py +++ /dev/null @@ -1,44 +0,0 @@ -import gym -import numpy as np - -class VideoWrapper(gym.Wrapper): - def __init__(self, - env, - mode='rgb_array', - enabled=True, - steps_per_render=1, - **kwargs - ): - super().__init__(env) - - self.mode = mode - self.enabled = enabled - self.render_kwargs = kwargs - self.steps_per_render = steps_per_render - - self.frames = list() - self.step_count = 0 - - def reset(self, **kwargs): - obs = super().reset(**kwargs) - self.frames = list() - self.step_count = 1 - if self.enabled: - frame = self.env.render( - mode=self.mode, **self.render_kwargs) - assert frame.dtype == np.uint8 - self.frames.append(frame) - return obs - - def step(self, action): - result = super().step(action) - self.step_count += 1 - if self.enabled and ((self.step_count % self.steps_per_render) == 0): - frame = self.env.render( - mode=self.mode, **self.render_kwargs) - assert frame.dtype == np.uint8 - self.frames.append(frame) - return result - - def render(self, mode='rgb_array', **kwargs): - return self.frames diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/action_ae/__init__.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/action_ae/__init__.py deleted file mode 100644 index 9a7b88d8e..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/action_ae/__init__.py +++ /dev/null @@ -1,63 +0,0 @@ -import torch -import torch.nn as nn -from torch.utils.data import DataLoader -import abc - -from typing import Optional, Union - -import diffusion_policy.model.bet.utils as utils - - -class AbstractActionAE(utils.SaveModule, abc.ABC): - @abc.abstractmethod - def fit_model( - self, - input_dataloader: DataLoader, - eval_dataloader: DataLoader, - obs_encoding_net: Optional[nn.Module] = None, - ) -> None: - pass - - @abc.abstractmethod - def encode_into_latent( - self, - input_action: torch.Tensor, - input_rep: Optional[torch.Tensor], - ) -> torch.Tensor: - """ - Given the input action, discretize it. - - Inputs: - input_action (shape: ... x action_dim): The input action to discretize. This can be in a batch, - and is generally assumed that the last dimnesion is the action dimension. - - Outputs: - discretized_action (shape: ... x num_tokens): The discretized action. - """ - raise NotImplementedError - - @abc.abstractmethod - def decode_actions( - self, - latent_action_batch: Optional[torch.Tensor], - input_rep_batch: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """ - Given a discretized action, convert it to a continuous action. - - Inputs: - latent_action_batch (shape: ... x num_tokens): The discretized action - generated by the discretizer. - - Outputs: - continuous_action (shape: ... x action_dim): The continuous action. - """ - raise NotImplementedError - - @property - @abc.abstractmethod - def num_latents(self) -> Union[int, float]: - """ - Number of possible latents for this generator, useful for state priors that use softmax. - """ - return float("inf") diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/action_ae/discretizers/k_means.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/action_ae/discretizers/k_means.py deleted file mode 100644 index 9051df083..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/action_ae/discretizers/k_means.py +++ /dev/null @@ -1,148 +0,0 @@ -import torch -import numpy as np - -import tqdm - -from typing import Optional, Tuple, Union -from diffusion_policy.model.common.dict_of_tensor_mixin import DictOfTensorMixin - - -class KMeansDiscretizer(DictOfTensorMixin): - """ - Simplified and modified version of KMeans algorithm from sklearn. - """ - - def __init__( - self, - action_dim: int, - num_bins: int = 100, - predict_offsets: bool = False, - ): - super().__init__() - self.n_bins = num_bins - self.action_dim = action_dim - self.predict_offsets = predict_offsets - - def fit_discretizer(self, input_actions: torch.Tensor) -> None: - assert ( - self.action_dim == input_actions.shape[-1] - ), f"Input action dimension {self.action_dim} does not match fitted model {input_actions.shape[-1]}" - - flattened_actions = input_actions.view(-1, self.action_dim) - cluster_centers = KMeansDiscretizer._kmeans( - flattened_actions, ncluster=self.n_bins - ) - self.params_dict['bin_centers'] = cluster_centers - - @property - def suggested_actions(self) -> torch.Tensor: - return self.params_dict['bin_centers'] - - @classmethod - def _kmeans(cls, x: torch.Tensor, ncluster: int = 512, niter: int = 50): - """ - Simple k-means clustering algorithm adapted from Karpathy's minGPT library - https://github.com/karpathy/minGPT/blob/master/play_image.ipynb - """ - N, D = x.size() - c = x[torch.randperm(N)[:ncluster]] # init clusters at random - - pbar = tqdm.trange(niter) - pbar.set_description("K-means clustering") - for i in pbar: - # assign all pixels to the closest codebook element - a = ((x[:, None, :] - c[None, :, :]) ** 2).sum(-1).argmin(1) - # move each codebook element to be the mean of the pixels that assigned to it - c = torch.stack([x[a == k].mean(0) for k in range(ncluster)]) - # re-assign any poorly positioned codebook elements - nanix = torch.any(torch.isnan(c), dim=1) - ndead = nanix.sum().item() - if ndead: - tqdm.tqdm.write( - "done step %d/%d, re-initialized %d dead clusters" - % (i + 1, niter, ndead) - ) - c[nanix] = x[torch.randperm(N)[:ndead]] # re-init dead clusters - return c - - def encode_into_latent( - self, input_action: torch.Tensor, input_rep: Optional[torch.Tensor] = None - ) -> torch.Tensor: - """ - Given the input action, discretize it using the k-Means clustering algorithm. - - Inputs: - input_action (shape: ... x action_dim): The input action to discretize. This can be in a batch, - and is generally assumed that the last dimnesion is the action dimension. - - Outputs: - discretized_action (shape: ... x num_tokens): The discretized action. - If self.predict_offsets is True, then the offsets are also returned. - """ - assert ( - input_action.shape[-1] == self.action_dim - ), "Input action dimension does not match fitted model" - - # flatten the input action - flattened_actions = input_action.view(-1, self.action_dim) - - # get the closest cluster center - closest_cluster_center = torch.argmin( - torch.sum( - (flattened_actions[:, None, :] - self.params_dict['bin_centers'][None, :, :]) ** 2, - dim=2, - ), - dim=1, - ) - # Reshape to the original shape - discretized_action = closest_cluster_center.view(input_action.shape[:-1] + (1,)) - - if self.predict_offsets: - # decode from latent and get the difference - reconstructed_action = self.decode_actions(discretized_action) - offsets = input_action - reconstructed_action - return (discretized_action, offsets) - else: - # return the one-hot vector - return discretized_action - - def decode_actions( - self, - latent_action_batch: torch.Tensor, - input_rep_batch: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """ - Given the latent action, reconstruct the original action. - - Inputs: - latent_action (shape: ... x 1): The latent action to reconstruct. This can be in a batch, - and is generally assumed that the last dimension is the action dimension. If the latent_action_batch - is a tuple, then it is assumed to be (discretized_action, offsets). - - Outputs: - reconstructed_action (shape: ... x action_dim): The reconstructed action. - """ - offsets = None - if type(latent_action_batch) == tuple: - latent_action_batch, offsets = latent_action_batch - # get the closest cluster center - closest_cluster_center = self.params_dict['bin_centers'][latent_action_batch] - # Reshape to the original shape - reconstructed_action = closest_cluster_center.view( - latent_action_batch.shape[:-1] + (self.action_dim,) - ) - if offsets is not None: - reconstructed_action += offsets - return reconstructed_action - - @property - def discretized_space(self) -> int: - return self.n_bins - - @property - def latent_dim(self) -> int: - return 1 - - @property - def num_latents(self) -> int: - return self.n_bins diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/latent_generators/latent_generator.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/latent_generators/latent_generator.py deleted file mode 100644 index 072bdc86a..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/latent_generators/latent_generator.py +++ /dev/null @@ -1,71 +0,0 @@ -import abc -import torch -from typing import Tuple, Optional - -import diffusion_policy.model.bet.utils as utils - - -class AbstractLatentGenerator(abc.ABC, utils.SaveModule): - """ - Abstract class for a generative model that can generate latents given observation representations. - - In the probabilisitc sense, this model fits and samples from P(latent|observation) given some observation. - """ - - @abc.abstractmethod - def get_latent_and_loss( - self, - obs_rep: torch.Tensor, - target_latents: torch.Tensor, - seq_masks: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Given a set of observation representation and generated latents, get the encoded latent and the loss. - - Inputs: - input_action: Batch of the actions taken in the multimodal demonstrations. - target_latents: Batch of the latents that the generator should learn to generate the actions from. - seq_masks: Batch of masks that indicate which timesteps are valid. - - Outputs: - latent: The sampled latent from the observation. - loss: The loss of the latent generator. - """ - pass - - @abc.abstractmethod - def generate_latents( - self, seq_obses: torch.Tensor, seq_masks: torch.Tensor - ) -> torch.Tensor: - """ - Given a batch of sequences of observations, generate a batch of sequences of latents. - - Inputs: - seq_obses: Batch of sequences of observations, of shape seq x batch x dim, following the transformer convention. - seq_masks: Batch of sequences of masks, of shape seq x batch, following the transformer convention. - - Outputs: - seq_latents: Batch of sequences of latents of shape seq x batch x latent_dim. - """ - pass - - def get_optimizer( - self, weight_decay: float, learning_rate: float, betas: Tuple[float, float] - ) -> torch.optim.Optimizer: - """ - Default optimizer class. Override this if you want to use a different optimizer. - """ - return torch.optim.Adam( - self.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=betas - ) - - -class LatentGeneratorDataParallel(torch.nn.DataParallel): - def get_latent_and_loss(self, *args, **kwargs): - return self.module.get_latent_and_loss(*args, **kwargs) # type: ignore - - def generate_latents(self, *args, **kwargs): - return self.module.generate_latents(*args, **kwargs) # type: ignore - - def get_optimizer(self, *args, **kwargs): - return self.module.get_optimizer(*args, **kwargs) # type: ignore diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/latent_generators/mingpt.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/latent_generators/mingpt.py deleted file mode 100644 index 4836b2791..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/latent_generators/mingpt.py +++ /dev/null @@ -1,188 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -import einops -import diffusion_policy.model.bet.latent_generators.latent_generator as latent_generator - -import diffusion_policy.model.bet.libraries.mingpt.model as mingpt_model -import diffusion_policy.model.bet.libraries.mingpt.trainer as mingpt_trainer -from diffusion_policy.model.bet.libraries.loss_fn import FocalLoss, soft_cross_entropy - -from typing import Optional, Tuple - - -class MinGPT(latent_generator.AbstractLatentGenerator): - def __init__( - self, - input_dim: int, - n_layer: int = 12, - n_head: int = 12, - n_embd: int = 768, - embd_pdrop: float = 0.1, - resid_pdrop: float = 0.1, - attn_pdrop: float = 0.1, - block_size: int = 128, - vocab_size: int = 50257, - latent_dim: int = 768, # Ignore, used for compatibility with other models. - action_dim: int = 0, - discrete_input: bool = False, - predict_offsets: bool = False, - offset_loss_scale: float = 1.0, - focal_loss_gamma: float = 0.0, - **kwargs - ): - super().__init__() - self.input_size = input_dim - self.n_layer = n_layer - self.n_head = n_head - self.n_embd = n_embd - self.embd_pdrop = embd_pdrop - self.resid_pdrop = resid_pdrop - self.attn_pdrop = attn_pdrop - self.block_size = block_size - self.vocab_size = vocab_size - self.action_dim = action_dim - self.predict_offsets = predict_offsets - self.offset_loss_scale = offset_loss_scale - self.focal_loss_gamma = focal_loss_gamma - for k, v in kwargs.items(): - setattr(self, k, v) - - gpt_config = mingpt_model.GPTConfig( - input_size=self.input_size, - vocab_size=self.vocab_size * (1 + self.action_dim) - if self.predict_offsets - else self.vocab_size, - block_size=self.block_size, - n_layer=n_layer, - n_head=n_head, - n_embd=n_embd, - discrete_input=discrete_input, - embd_pdrop=embd_pdrop, - resid_pdrop=resid_pdrop, - attn_pdrop=attn_pdrop, - ) - - self.model = mingpt_model.GPT(gpt_config) - - def get_latent_and_loss( - self, - obs_rep: torch.Tensor, - target_latents: torch.Tensor, - seq_masks: Optional[torch.Tensor] = None, - return_loss_components: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor]: - # Unlike torch.transformers, GPT takes in batch x seq_len x embd_dim - # obs_rep = einops.rearrange(obs_rep, "seq batch embed -> batch seq embed") - # target_latents = einops.rearrange( - # target_latents, "seq batch embed -> batch seq embed" - # ) - # While this has been trained autoregressively, - # there is no reason why it needs to be so. - # We can just use the observation as the input and the next latent as the target. - if self.predict_offsets: - target_latents, target_offsets = target_latents - is_soft_target = (target_latents.shape[-1] == self.vocab_size) and ( - self.vocab_size != 1 - ) - if is_soft_target: - target_latents = target_latents.view(-1, target_latents.size(-1)) - criterion = soft_cross_entropy - else: - target_latents = target_latents.view(-1) - if self.vocab_size == 1: - # unify k-means (target_class == 0) and GMM (target_prob == 1) - target_latents = torch.zeros_like(target_latents) - criterion = FocalLoss(gamma=self.focal_loss_gamma) - if self.predict_offsets: - output, _ = self.model(obs_rep) - logits = output[:, :, : self.vocab_size] - offsets = output[:, :, self.vocab_size :] - batch = logits.shape[0] - seq = logits.shape[1] - offsets = einops.rearrange( - offsets, - "N T (V A) -> (N T) V A", # N = batch, T = seq - V=self.vocab_size, - A=self.action_dim, - ) - # calculate (optionally soft) cross entropy and offset losses - class_loss = criterion(logits.view(-1, logits.size(-1)), target_latents) - # offset loss is only calculated on the target class - # if soft targets, argmax is considered the target class - selected_offsets = offsets[ - torch.arange(offsets.size(0)), - target_latents.argmax(dim=-1).view(-1) - if is_soft_target - else target_latents.view(-1), - ] - offset_loss = self.offset_loss_scale * F.mse_loss( - selected_offsets, target_offsets.view(-1, self.action_dim) - ) - loss = offset_loss + class_loss - logits = einops.rearrange(logits, "batch seq classes -> seq batch classes") - offsets = einops.rearrange( - offsets, - "(N T) V A -> T N V A", # ? N, T order? Anyway does not affect loss and training (might affect visualization) - N=batch, - T=seq, - ) - if return_loss_components: - return ( - (logits, offsets), - loss, - {"offset": offset_loss, "class": class_loss, "total": loss}, - ) - else: - return (logits, offsets), loss - else: - logits, _ = self.model(obs_rep) - loss = criterion(logits.view(-1, logits.size(-1)), target_latents) - logits = einops.rearrange( - logits, "batch seq classes -> seq batch classes" - ) # ? N, T order? Anyway does not affect loss and training (might affect visualization) - if return_loss_components: - return logits, loss, {"class": loss, "total": loss} - else: - return logits, loss - - def generate_latents( - self, obs_rep: torch.Tensor - ) -> torch.Tensor: - batch, seq, embed = obs_rep.shape - - output, _ = self.model(obs_rep, None) - if self.predict_offsets: - logits = output[:, :, : self.vocab_size] - offsets = output[:, :, self.vocab_size :] - offsets = einops.rearrange( - offsets, - "N T (V A) -> (N T) V A", # N = batch, T = seq - V=self.vocab_size, - A=self.action_dim, - ) - else: - logits = output - probs = F.softmax(logits, dim=-1) - batch, seq, choices = probs.shape - # Sample from the multinomial distribution, one per row. - sampled_data = torch.multinomial(probs.view(-1, choices), num_samples=1) - sampled_data = einops.rearrange( - sampled_data, "(batch seq) 1 -> batch seq 1", batch=batch, seq=seq - ) - if self.predict_offsets: - sampled_offsets = offsets[ - torch.arange(offsets.shape[0]), sampled_data.flatten() - ].view(batch, seq, self.action_dim) - - return (sampled_data, sampled_offsets) - else: - return sampled_data - - def get_optimizer( - self, weight_decay: float, learning_rate: float, betas: Tuple[float, float] - ) -> torch.optim.Optimizer: - trainer_cfg = mingpt_trainer.TrainerConfig( - weight_decay=weight_decay, learning_rate=learning_rate, betas=betas - ) - return self.model.configure_optimizers(trainer_cfg) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/latent_generators/transformer.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/latent_generators/transformer.py deleted file mode 100644 index 76af7cced..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/latent_generators/transformer.py +++ /dev/null @@ -1,108 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -import einops -import diffusion_policy.model.bet.latent_generators.latent_generator as latent_generator - -from diffusion_policy.model.diffusion.transformer_for_diffusion import TransformerForDiffusion -from diffusion_policy.model.bet.libraries.loss_fn import FocalLoss, soft_cross_entropy - -from typing import Optional, Tuple - -class Transformer(latent_generator.AbstractLatentGenerator): - def __init__( - self, - input_dim: int, - num_bins: int, - action_dim: int, - horizon: int, - focal_loss_gamma: float, - offset_loss_scale: float, - **kwargs - ): - super().__init__() - self.model = TransformerForDiffusion( - input_dim=input_dim, - output_dim=num_bins * (1 + action_dim), - horizon=horizon, - **kwargs - ) - self.vocab_size = num_bins - self.focal_loss_gamma = focal_loss_gamma - self.offset_loss_scale = offset_loss_scale - self.action_dim = action_dim - - def get_optimizer(self, **kwargs) -> torch.optim.Optimizer: - return self.model.configure_optimizers(**kwargs) - - def get_latent_and_loss(self, - obs_rep: torch.Tensor, - target_latents: torch.Tensor, - return_loss_components=True, - ) -> Tuple[torch.Tensor, torch.Tensor]: - target_latents, target_offsets = target_latents - target_latents = target_latents.view(-1) - criterion = FocalLoss(gamma=self.focal_loss_gamma) - - t = torch.tensor(0, device=self.model.device) - output = self.model(obs_rep, t) - logits = output[:, :, : self.vocab_size] - offsets = output[:, :, self.vocab_size :] - batch = logits.shape[0] - seq = logits.shape[1] - offsets = einops.rearrange( - offsets, - "N T (V A) -> (N T) V A", # N = batch, T = seq - V=self.vocab_size, - A=self.action_dim, - ) - # calculate (optionally soft) cross entropy and offset losses - class_loss = criterion(logits.view(-1, logits.size(-1)), target_latents) - # offset loss is only calculated on the target class - # if soft targets, argmax is considered the target class - selected_offsets = offsets[ - torch.arange(offsets.size(0)), - target_latents.view(-1), - ] - offset_loss = self.offset_loss_scale * F.mse_loss( - selected_offsets, target_offsets.view(-1, self.action_dim) - ) - loss = offset_loss + class_loss - logits = einops.rearrange(logits, "batch seq classes -> seq batch classes") - offsets = einops.rearrange( - offsets, - "(N T) V A -> T N V A", # ? N, T order? Anyway does not affect loss and training (might affect visualization) - N=batch, - T=seq, - ) - return ( - (logits, offsets), - loss, - {"offset": offset_loss, "class": class_loss, "total": loss}, - ) - - def generate_latents( - self, obs_rep: torch.Tensor - ) -> torch.Tensor: - t = torch.tensor(0, device=self.model.device) - output = self.model(obs_rep, t) - logits = output[:, :, : self.vocab_size] - offsets = output[:, :, self.vocab_size :] - offsets = einops.rearrange( - offsets, - "N T (V A) -> (N T) V A", # N = batch, T = seq - V=self.vocab_size, - A=self.action_dim, - ) - - probs = F.softmax(logits, dim=-1) - batch, seq, choices = probs.shape - # Sample from the multinomial distribution, one per row. - sampled_data = torch.multinomial(probs.view(-1, choices), num_samples=1) - sampled_data = einops.rearrange( - sampled_data, "(batch seq) 1 -> batch seq 1", batch=batch, seq=seq - ) - sampled_offsets = offsets[ - torch.arange(offsets.shape[0]), sampled_data.flatten() - ].view(batch, seq, self.action_dim) - return (sampled_data, sampled_offsets) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/loss_fn.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/loss_fn.py deleted file mode 100644 index 7873ae1b4..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/loss_fn.py +++ /dev/null @@ -1,168 +0,0 @@ -from typing import Optional, Sequence - -import torch -from torch import Tensor -from torch import nn -from torch.nn import functional as F - -# Reference: https://github.com/pytorch/pytorch/issues/11959 -def soft_cross_entropy( - input: torch.Tensor, - target: torch.Tensor, -) -> torch.Tensor: - """ - Args: - input: (batch_size, num_classes): tensor of raw logits - target: (batch_size, num_classes): tensor of class probability; sum(target) == 1 - - Returns: - loss: (batch_size,) - """ - log_probs = torch.log_softmax(input, dim=-1) - # target is a distribution - loss = F.kl_div(log_probs, target, reduction="batchmean") - return loss - - -# Focal loss implementation -# Source: https://github.com/AdeelH/pytorch-multi-class-focal-loss/blob/master/focal_loss.py -# MIT License -# -# Copyright (c) 2020 Adeel Hassan -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -class FocalLoss(nn.Module): - """Focal Loss, as described in https://arxiv.org/abs/1708.02002. - It is essentially an enhancement to cross entropy loss and is - useful for classification tasks when there is a large class imbalance. - x is expected to contain raw, unnormalized scores for each class. - y is expected to contain class labels. - Shape: - - x: (batch_size, C) or (batch_size, C, d1, d2, ..., dK), K > 0. - - y: (batch_size,) or (batch_size, d1, d2, ..., dK), K > 0. - """ - - def __init__( - self, - alpha: Optional[Tensor] = None, - gamma: float = 0.0, - reduction: str = "mean", - ignore_index: int = -100, - ): - """Constructor. - Args: - alpha (Tensor, optional): Weights for each class. Defaults to None. - gamma (float, optional): A constant, as described in the paper. - Defaults to 0. - reduction (str, optional): 'mean', 'sum' or 'none'. - Defaults to 'mean'. - ignore_index (int, optional): class label to ignore. - Defaults to -100. - """ - if reduction not in ("mean", "sum", "none"): - raise ValueError('Reduction must be one of: "mean", "sum", "none".') - - super().__init__() - self.alpha = alpha - self.gamma = gamma - self.ignore_index = ignore_index - self.reduction = reduction - - self.nll_loss = nn.NLLLoss( - weight=alpha, reduction="none", ignore_index=ignore_index - ) - - def __repr__(self): - arg_keys = ["alpha", "gamma", "ignore_index", "reduction"] - arg_vals = [self.__dict__[k] for k in arg_keys] - arg_strs = [f"{k}={v}" for k, v in zip(arg_keys, arg_vals)] - arg_str = ", ".join(arg_strs) - return f"{type(self).__name__}({arg_str})" - - def forward(self, x: Tensor, y: Tensor) -> Tensor: - if x.ndim > 2: - # (N, C, d1, d2, ..., dK) --> (N * d1 * ... * dK, C) - c = x.shape[1] - x = x.permute(0, *range(2, x.ndim), 1).reshape(-1, c) - # (N, d1, d2, ..., dK) --> (N * d1 * ... * dK,) - y = y.view(-1) - - unignored_mask = y != self.ignore_index - y = y[unignored_mask] - if len(y) == 0: - return 0.0 - x = x[unignored_mask] - - # compute weighted cross entropy term: -alpha * log(pt) - # (alpha is already part of self.nll_loss) - log_p = F.log_softmax(x, dim=-1) - ce = self.nll_loss(log_p, y) - - # get true class column from each row - all_rows = torch.arange(len(x)) - log_pt = log_p[all_rows, y] - - # compute focal term: (1 - pt)^gamma - pt = log_pt.exp() - focal_term = (1 - pt) ** self.gamma - - # the full loss: -alpha * ((1 - pt)^gamma) * log(pt) - loss = focal_term * ce - - if self.reduction == "mean": - loss = loss.mean() - elif self.reduction == "sum": - loss = loss.sum() - - return loss - - -def focal_loss( - alpha: Optional[Sequence] = None, - gamma: float = 0.0, - reduction: str = "mean", - ignore_index: int = -100, - device="cpu", - dtype=torch.float32, -) -> FocalLoss: - """Factory function for FocalLoss. - Args: - alpha (Sequence, optional): Weights for each class. Will be converted - to a Tensor if not None. Defaults to None. - gamma (float, optional): A constant, as described in the paper. - Defaults to 0. - reduction (str, optional): 'mean', 'sum' or 'none'. - Defaults to 'mean'. - ignore_index (int, optional): class label to ignore. - Defaults to -100. - device (str, optional): Device to move alpha to. Defaults to 'cpu'. - dtype (torch.dtype, optional): dtype to cast alpha to. - Defaults to torch.float32. - Returns: - A FocalLoss object - """ - if alpha is not None: - if not isinstance(alpha, Tensor): - alpha = torch.tensor(alpha) - alpha = alpha.to(device=device, dtype=dtype) - - fl = FocalLoss( - alpha=alpha, gamma=gamma, reduction=reduction, ignore_index=ignore_index - ) - return fl diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/LICENSE b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/LICENSE deleted file mode 100644 index 87bfb153c..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/LICENSE +++ /dev/null @@ -1,8 +0,0 @@ -The MIT License (MIT) Copyright (c) 2020 Andrej Karpathy - -Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. - diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/__init__.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/model.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/model.py deleted file mode 100644 index 06308e391..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/model.py +++ /dev/null @@ -1,250 +0,0 @@ -""" -GPT model: -- the initial stem consists of a combination of token encoding and a positional encoding -- the meat of it is a uniform sequence of Transformer blocks - - each Transformer is a sequential combination of a 1-hidden-layer MLP block and a self-attention block - - all blocks feed into a central residual pathway similar to resnets -- the final decoder is a linear projection into a vanilla Softmax classifier -""" - -import math -import logging - -import torch -import torch.nn as nn -from torch.nn import functional as F - -logger = logging.getLogger(__name__) - - -class GPTConfig: - """base GPT config, params common to all GPT versions""" - - embd_pdrop = 0.1 - resid_pdrop = 0.1 - attn_pdrop = 0.1 - discrete_input = False - input_size = 10 - n_embd = 768 - n_layer = 12 - - def __init__(self, vocab_size, block_size, **kwargs): - self.vocab_size = vocab_size - self.block_size = block_size - for k, v in kwargs.items(): - setattr(self, k, v) - - -class GPT1Config(GPTConfig): - """GPT-1 like network roughly 125M params""" - - n_layer = 12 - n_head = 12 - n_embd = 768 - - -class CausalSelfAttention(nn.Module): - """ - A vanilla multi-head masked self-attention layer with a projection at the end. - It is possible to use torch.nn.MultiheadAttention here but I am including an - explicit implementation here to show that there is nothing too scary here. - """ - - def __init__(self, config): - super().__init__() - assert config.n_embd % config.n_head == 0 - # key, query, value projections for all heads - self.key = nn.Linear(config.n_embd, config.n_embd) - self.query = nn.Linear(config.n_embd, config.n_embd) - self.value = nn.Linear(config.n_embd, config.n_embd) - # regularization - self.attn_drop = nn.Dropout(config.attn_pdrop) - self.resid_drop = nn.Dropout(config.resid_pdrop) - # output projection - self.proj = nn.Linear(config.n_embd, config.n_embd) - # causal mask to ensure that attention is only applied to the left in the input sequence - self.register_buffer( - "mask", - torch.tril(torch.ones(config.block_size, config.block_size)).view( - 1, 1, config.block_size, config.block_size - ), - ) - self.n_head = config.n_head - - def forward(self, x): - ( - B, - T, - C, - ) = x.size() # batch size, sequence length, embedding dimensionality (n_embd) - - # calculate query, key, values for all heads in batch and move head forward to be the batch dim - k = ( - self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) - ) # (B, nh, T, hs) - q = ( - self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) - ) # (B, nh, T, hs) - v = ( - self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) - ) # (B, nh, T, hs) - - # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T) - att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) - att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf")) - att = F.softmax(att, dim=-1) - att = self.attn_drop(att) - y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) - y = ( - y.transpose(1, 2).contiguous().view(B, T, C) - ) # re-assemble all head outputs side by side - - # output projection - y = self.resid_drop(self.proj(y)) - return y - - -class Block(nn.Module): - """an unassuming Transformer block""" - - def __init__(self, config): - super().__init__() - self.ln1 = nn.LayerNorm(config.n_embd) - self.ln2 = nn.LayerNorm(config.n_embd) - self.attn = CausalSelfAttention(config) - self.mlp = nn.Sequential( - nn.Linear(config.n_embd, 4 * config.n_embd), - nn.GELU(), - nn.Linear(4 * config.n_embd, config.n_embd), - nn.Dropout(config.resid_pdrop), - ) - - def forward(self, x): - x = x + self.attn(self.ln1(x)) - x = x + self.mlp(self.ln2(x)) - return x - - -class GPT(nn.Module): - """the full GPT language model, with a context size of block_size""" - - def __init__(self, config: GPTConfig): - super().__init__() - - # input embedding stem - if config.discrete_input: - self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd) - else: - self.tok_emb = nn.Linear(config.input_size, config.n_embd) - self.discrete_input = config.discrete_input - self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd)) - self.drop = nn.Dropout(config.embd_pdrop) - # transformer - self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)]) - # decoder head - self.ln_f = nn.LayerNorm(config.n_embd) - self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False) - - self.block_size = config.block_size - self.apply(self._init_weights) - - logger.info( - "number of parameters: %e", sum(p.numel() for p in self.parameters()) - ) - - def get_block_size(self): - return self.block_size - - def _init_weights(self, module): - if isinstance(module, (nn.Linear, nn.Embedding)): - torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) - if isinstance(module, nn.Linear) and module.bias is not None: - torch.nn.init.zeros_(module.bias) - elif isinstance(module, nn.LayerNorm): - torch.nn.init.zeros_(module.bias) - torch.nn.init.ones_(module.weight) - elif isinstance(module, GPT): - torch.nn.init.normal_(module.pos_emb, mean=0.0, std=0.02) - - def configure_optimizers(self, train_config): - """ - This long function is unfortunately doing something very simple and is being very defensive: - We are separating out all parameters of the model into two buckets: those that will experience - weight decay for regularization and those that won't (biases, and layernorm/embedding weights). - We are then returning the PyTorch optimizer object. - """ - - # separate out all parameters to those that will and won't experience regularizing weight decay - decay = set() - no_decay = set() - whitelist_weight_modules = (torch.nn.Linear,) - blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding) - for mn, m in self.named_modules(): - for pn, p in m.named_parameters(): - fpn = "%s.%s" % (mn, pn) if mn else pn # full param name - - if pn.endswith("bias"): - # all biases will not be decayed - no_decay.add(fpn) - elif pn.endswith("weight") and isinstance(m, whitelist_weight_modules): - # weights of whitelist modules will be weight decayed - decay.add(fpn) - elif pn.endswith("weight") and isinstance(m, blacklist_weight_modules): - # weights of blacklist modules will NOT be weight decayed - no_decay.add(fpn) - - # special case the position embedding parameter in the root GPT module as not decayed - no_decay.add("pos_emb") - - # validate that we considered every parameter - param_dict = {pn: p for pn, p in self.named_parameters()} - inter_params = decay & no_decay - union_params = decay | no_decay - assert ( - len(inter_params) == 0 - ), "parameters %s made it into both decay/no_decay sets!" % (str(inter_params),) - assert ( - len(param_dict.keys() - union_params) == 0 - ), "parameters %s were not separated into either decay/no_decay set!" % ( - str(param_dict.keys() - union_params), - ) - - # create the pytorch optimizer object - optim_groups = [ - { - "params": [param_dict[pn] for pn in sorted(list(decay))], - "weight_decay": train_config.weight_decay, - }, - { - "params": [param_dict[pn] for pn in sorted(list(no_decay))], - "weight_decay": 0.0, - }, - ] - optimizer = torch.optim.AdamW( - optim_groups, lr=train_config.learning_rate, betas=train_config.betas - ) - return optimizer - - def forward(self, idx, targets=None): - if self.discrete_input: - b, t = idx.size() - else: - b, t, dim = idx.size() - assert t <= self.block_size, "Cannot forward, model block size is exhausted." - - # forward the GPT model - token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector - position_embeddings = self.pos_emb[ - :, :t, : - ] # each position maps to a (learnable) vector - x = self.drop(token_embeddings + position_embeddings) - x = self.blocks(x) - x = self.ln_f(x) - logits = self.head(x) - - # if we are given some desired targets also calculate the loss - loss = None - if targets is not None: - loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) - - return logits, loss diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/trainer.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/trainer.py deleted file mode 100644 index 64fec8ee9..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/trainer.py +++ /dev/null @@ -1,162 +0,0 @@ -""" -Simple training loop; Boilerplate that could apply to any arbitrary neural network, -so nothing in this file really has anything to do with GPT specifically. -""" - -import math -import logging - -from tqdm import tqdm -import numpy as np - -import torch -import torch.optim as optim -from torch.optim.lr_scheduler import LambdaLR -from torch.utils.data.dataloader import DataLoader - -logger = logging.getLogger(__name__) - - -class TrainerConfig: - # optimization parameters - max_epochs = 10 - batch_size = 64 - learning_rate = 3e-4 - betas = (0.9, 0.95) - grad_norm_clip = 1.0 - weight_decay = 0.1 # only applied on matmul weights - # learning rate decay params: linear warmup followed by cosine decay to 10% of original - lr_decay = False - warmup_tokens = 375e6 # these two numbers come from the GPT-3 paper, but may not be good defaults elsewhere - final_tokens = 260e9 # (at what point we reach 10% of original LR) - # checkpoint settings - ckpt_path = None - num_workers = 0 # for DataLoader - - def __init__(self, **kwargs): - for k, v in kwargs.items(): - setattr(self, k, v) - - -class Trainer: - def __init__(self, model, train_dataset, test_dataset, config): - self.model = model - self.train_dataset = train_dataset - self.test_dataset = test_dataset - self.config = config - - # take over whatever gpus are on the system - self.device = "cpu" - if torch.cuda.is_available(): - self.device = torch.cuda.current_device() - self.model = torch.nn.DataParallel(self.model).to(self.device) - - def save_checkpoint(self): - # DataParallel wrappers keep raw model object in .module attribute - raw_model = self.model.module if hasattr(self.model, "module") else self.model - logger.info("saving %s", self.config.ckpt_path) - torch.save(raw_model.state_dict(), self.config.ckpt_path) - - def train(self): - model, config = self.model, self.config - raw_model = model.module if hasattr(self.model, "module") else model - optimizer = raw_model.configure_optimizers(config) - - def run_epoch(loader, is_train): - model.train(is_train) - - losses = [] - pbar = ( - tqdm(enumerate(loader), total=len(loader)) - if is_train - else enumerate(loader) - ) - for it, (x, y) in pbar: - - # place data on the correct device - x = x.to(self.device) - y = y.to(self.device) - - # forward the model - with torch.set_grad_enabled(is_train): - logits, loss = model(x, y) - loss = ( - loss.mean() - ) # collapse all losses if they are scattered on multiple gpus - losses.append(loss.item()) - - if is_train: - - # backprop and update the parameters - model.zero_grad() - loss.backward() - torch.nn.utils.clip_grad_norm_( - model.parameters(), config.grad_norm_clip - ) - optimizer.step() - - # decay the learning rate based on our progress - if config.lr_decay: - self.tokens += ( - y >= 0 - ).sum() # number of tokens processed this step (i.e. label is not -100) - if self.tokens < config.warmup_tokens: - # linear warmup - lr_mult = float(self.tokens) / float( - max(1, config.warmup_tokens) - ) - else: - # cosine learning rate decay - progress = float( - self.tokens - config.warmup_tokens - ) / float( - max(1, config.final_tokens - config.warmup_tokens) - ) - lr_mult = max( - 0.1, 0.5 * (1.0 + math.cos(math.pi * progress)) - ) - lr = config.learning_rate * lr_mult - for param_group in optimizer.param_groups: - param_group["lr"] = lr - else: - lr = config.learning_rate - - # report progress - pbar.set_description( # type: ignore - f"epoch {epoch+1} iter {it}: train loss {loss.item():.5f}. lr {lr:e}" - ) - - if not is_train: - test_loss = float(np.mean(losses)) - logger.info("test loss: %f", test_loss) - return test_loss - - best_loss = float("inf") - self.tokens = 0 # counter used for learning rate decay - - train_loader = DataLoader( - self.train_dataset, - shuffle=True, - pin_memory=True, - batch_size=config.batch_size, - num_workers=config.num_workers, - ) - if self.test_dataset is not None: - test_loader = DataLoader( - self.test_dataset, - shuffle=True, - pin_memory=True, - batch_size=config.batch_size, - num_workers=config.num_workers, - ) - - for epoch in range(config.max_epochs): - run_epoch(train_loader, is_train=True) - if self.test_dataset is not None: - test_loss = run_epoch(test_loader, is_train=False) - - # supports early stopping based on the test loss, or just save always if no test set is provided - good_model = self.test_dataset is None or test_loss < best_loss - if self.config.ckpt_path is not None and good_model: - best_loss = test_loss - self.save_checkpoint() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/utils.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/utils.py deleted file mode 100644 index 620b9d8c3..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/libraries/mingpt/utils.py +++ /dev/null @@ -1,52 +0,0 @@ -import random -import numpy as np -import torch -import torch.nn as nn -from torch.nn import functional as F - - -def set_seed(seed): - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -def top_k_logits(logits, k): - v, ix = torch.topk(logits, k) - out = logits.clone() - out[out < v[:, [-1]]] = -float("Inf") - return out - - -@torch.no_grad() -def sample(model, x, steps, temperature=1.0, sample=False, top_k=None): - """ - take a conditioning sequence of indices in x (of shape (b,t)) and predict the next token in - the sequence, feeding the predictions back into the model each time. Clearly the sampling - has quadratic complexity unlike an RNN that is only linear, and has a finite context window - of block_size, unlike an RNN that has an infinite context window. - """ - block_size = model.get_block_size() - model.eval() - for k in range(steps): - x_cond = ( - x if x.size(1) <= block_size else x[:, -block_size:] - ) # crop context if needed - logits, _ = model(x_cond) - # pluck the logits at the final step and scale by temperature - logits = logits[:, -1, :] / temperature - # optionally crop probabilities to only the top k options - if top_k is not None: - logits = top_k_logits(logits, top_k) - # apply softmax to convert to probabilities - probs = F.softmax(logits, dim=-1) - # sample from the distribution or take the most likely - if sample: - ix = torch.multinomial(probs, num_samples=1) - else: - _, ix = torch.topk(probs, k=1, dim=-1) - # append to the sequence and continue - x = torch.cat((x, ix), dim=1) - - return x diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/utils.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/utils.py deleted file mode 100644 index e021a0b22..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/bet/utils.py +++ /dev/null @@ -1,131 +0,0 @@ -import os -import random -from collections import OrderedDict -from typing import List, Optional - -import einops -import numpy as np -import torch -import torch.nn as nn - -from torch.utils.data import random_split -import wandb - - -def mlp(input_dim, hidden_dim, output_dim, hidden_depth, output_mod=None): - if hidden_depth == 0: - mods = [nn.Linear(input_dim, output_dim)] - else: - mods = [nn.Linear(input_dim, hidden_dim), nn.ReLU(inplace=True)] - for i in range(hidden_depth - 1): - mods += [nn.Linear(hidden_dim, hidden_dim), nn.ReLU(inplace=True)] - mods.append(nn.Linear(hidden_dim, output_dim)) - if output_mod is not None: - mods.append(output_mod) - trunk = nn.Sequential(*mods) - return trunk - - -class eval_mode: - def __init__(self, *models, no_grad=False): - self.models = models - self.no_grad = no_grad - self.no_grad_context = torch.no_grad() - - def __enter__(self): - self.prev_states = [] - for model in self.models: - self.prev_states.append(model.training) - model.train(False) - if self.no_grad: - self.no_grad_context.__enter__() - - def __exit__(self, *args): - if self.no_grad: - self.no_grad_context.__exit__(*args) - for model, state in zip(self.models, self.prev_states): - model.train(state) - return False - - -def freeze_module(module: nn.Module) -> nn.Module: - for param in module.parameters(): - param.requires_grad = False - module.eval() - return module - - -def set_seed_everywhere(seed): - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - np.random.seed(seed) - random.seed(seed) - - -def shuffle_along_axis(a, axis): - idx = np.random.rand(*a.shape).argsort(axis=axis) - return np.take_along_axis(a, idx, axis=axis) - - -def transpose_batch_timestep(*args): - return (einops.rearrange(arg, "b t ... -> t b ...") for arg in args) - - -class TrainWithLogger: - def reset_log(self): - self.log_components = OrderedDict() - - def log_append(self, log_key, length, loss_components): - for key, value in loss_components.items(): - key_name = f"{log_key}/{key}" - count, sum = self.log_components.get(key_name, (0, 0.0)) - self.log_components[key_name] = ( - count + length, - sum + (length * value.detach().cpu().item()), - ) - - def flush_log(self, epoch, iterator=None): - log_components = OrderedDict() - iterator_log_component = OrderedDict() - for key, value in self.log_components.items(): - count, sum = value - to_log = sum / count - log_components[key] = to_log - # Set the iterator status - log_key, name_key = key.split("/") - iterator_log_name = f"{log_key[0]}{name_key[0]}".upper() - iterator_log_component[iterator_log_name] = to_log - postfix = ",".join( - "{}:{:.2e}".format(key, iterator_log_component[key]) - for key in iterator_log_component.keys() - ) - if iterator is not None: - iterator.set_postfix_str(postfix) - wandb.log(log_components, step=epoch) - self.log_components = OrderedDict() - - -class SaveModule(nn.Module): - def set_snapshot_path(self, path): - self.snapshot_path = path - print(f"Setting snapshot path to {self.snapshot_path}") - - def save_snapshot(self): - os.makedirs(self.snapshot_path, exist_ok=True) - torch.save(self.state_dict(), self.snapshot_path / "snapshot.pth") - - def load_snapshot(self): - self.load_state_dict(torch.load(self.snapshot_path / "snapshot.pth")) - - -def split_datasets(dataset, train_fraction=0.95, random_seed=42): - dataset_length = len(dataset) - lengths = [ - int(train_fraction * dataset_length), - dataset_length - int(train_fraction * dataset_length), - ] - train_set, val_set = random_split( - dataset, lengths, generator=torch.Generator().manual_seed(random_seed) - ) - return train_set, val_set diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/dict_of_tensor_mixin.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/dict_of_tensor_mixin.py deleted file mode 100644 index 9d08bc135..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/dict_of_tensor_mixin.py +++ /dev/null @@ -1,38 +0,0 @@ -import torch -import torch.nn as nn - -class DictOfTensorMixin(nn.Module): - def __init__(self, params_dict=None): - super().__init__() - if params_dict is None: - params_dict = nn.ParameterDict() - self.params_dict = params_dict - - @property - def device(self): - return next(iter(self.parameters())).device - - def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): - def dfs_add(dest, keys, value: torch.Tensor): - if len(keys) == 1: - dest[keys[0]] = value - return - - if keys[0] not in dest: - dest[keys[0]] = nn.ParameterDict() - dfs_add(dest[keys[0]], keys[1:], value) - - def load_dict(state_dict, prefix): - out_dict = nn.ParameterDict() - for key, value in state_dict.items(): - value: torch.Tensor - if key.startswith(prefix): - param_keys = key[len(prefix):].split('.')[1:] - # if len(param_keys) == 0: - # import pdb; pdb.set_trace() - dfs_add(out_dict, param_keys, value.clone()) - return out_dict - - self.params_dict = load_dict(state_dict, prefix + 'params_dict') - self.params_dict.requires_grad_(False) - return diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/lr_scheduler.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/lr_scheduler.py deleted file mode 100644 index f4c053363..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/lr_scheduler.py +++ /dev/null @@ -1,46 +0,0 @@ -from diffusers.optimization import ( - Union, SchedulerType, Optional, - Optimizer, TYPE_TO_SCHEDULER_FUNCTION -) - -def get_scheduler( - name: Union[str, SchedulerType], - optimizer: Optimizer, - num_warmup_steps: Optional[int] = None, - num_training_steps: Optional[int] = None, - **kwargs -): - """ - Added kwargs vs diffuser's original implementation - - Unified API to get any scheduler from its name. - - Args: - name (`str` or `SchedulerType`): - The name of the scheduler to use. - optimizer (`torch.optim.Optimizer`): - The optimizer that will be used during training. - num_warmup_steps (`int`, *optional*): - The number of warmup steps to do. This is not required by all schedulers (hence the argument being - optional), the function will raise an error if it's unset and the scheduler type requires it. - num_training_steps (`int``, *optional*): - The number of training steps to do. This is not required by all schedulers (hence the argument being - optional), the function will raise an error if it's unset and the scheduler type requires it. - """ - name = SchedulerType(name) - schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name] - if name == SchedulerType.CONSTANT: - return schedule_func(optimizer, **kwargs) - - # All other schedulers require `num_warmup_steps` - if num_warmup_steps is None: - raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.") - - if name == SchedulerType.CONSTANT_WITH_WARMUP: - return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, **kwargs) - - # All other schedulers require `num_training_steps` - if num_training_steps is None: - raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.") - - return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, **kwargs) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/module_attr_mixin.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/module_attr_mixin.py deleted file mode 100644 index 8cbdf7099..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/module_attr_mixin.py +++ /dev/null @@ -1,14 +0,0 @@ -import torch.nn as nn - -class ModuleAttrMixin(nn.Module): - def __init__(self): - super().__init__() - self._dummy_variable = nn.Parameter() - - @property - def device(self): - return next(iter(self.parameters())).device - - @property - def dtype(self): - return next(iter(self.parameters())).dtype diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/normalizer.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/normalizer.py deleted file mode 100644 index a26469290..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/normalizer.py +++ /dev/null @@ -1,353 +0,0 @@ -from typing import Union, Dict - -import unittest -import zarr -import numpy as np -import torch -import torch.nn as nn -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.model.common.dict_of_tensor_mixin import DictOfTensorMixin - - -class LinearNormalizer(DictOfTensorMixin): - avaliable_modes = ['limits', 'gaussian'] - - @torch.no_grad() - def fit(self, - data: Union[Dict, torch.Tensor, np.ndarray, zarr.Array], - last_n_dims=1, - dtype=torch.float32, - mode='limits', - output_max=1., - output_min=-1., - range_eps=1e-4, - fit_offset=True): - if isinstance(data, dict): - for key, value in data.items(): - self.params_dict[key] = _fit(value, - last_n_dims=last_n_dims, - dtype=dtype, - mode=mode, - output_max=output_max, - output_min=output_min, - range_eps=range_eps, - fit_offset=fit_offset) - else: - self.params_dict['_default'] = _fit(data, - last_n_dims=last_n_dims, - dtype=dtype, - mode=mode, - output_max=output_max, - output_min=output_min, - range_eps=range_eps, - fit_offset=fit_offset) - - def __call__(self, x: Union[Dict, torch.Tensor, np.ndarray]) -> torch.Tensor: - return self.normalize(x) - - def __getitem__(self, key: str): - return SingleFieldLinearNormalizer(self.params_dict[key]) - - def __setitem__(self, key: str , value: 'SingleFieldLinearNormalizer'): - self.params_dict[key] = value.params_dict - - def _normalize_impl(self, x, forward=True): - if isinstance(x, dict): - result = dict() - for key, value in x.items(): - params = self.params_dict[key] - result[key] = _normalize(value, params, forward=forward) - return result - else: - if '_default' not in self.params_dict: - raise RuntimeError("Not initialized") - params = self.params_dict['_default'] - return _normalize(x, params, forward=forward) - - def normalize(self, x: Union[Dict, torch.Tensor, np.ndarray]) -> torch.Tensor: - return self._normalize_impl(x, forward=True) - - def unnormalize(self, x: Union[Dict, torch.Tensor, np.ndarray]) -> torch.Tensor: - return self._normalize_impl(x, forward=False) - - def get_input_stats(self) -> Dict: - if len(self.params_dict) == 0: - raise RuntimeError("Not initialized") - if len(self.params_dict) == 1 and '_default' in self.params_dict: - return self.params_dict['_default']['input_stats'] - - result = dict() - for key, value in self.params_dict.items(): - if key != '_default': - result[key] = value['input_stats'] - return result - - - def get_output_stats(self, key='_default'): - input_stats = self.get_input_stats() - if 'min' in input_stats: - # no dict - return dict_apply(input_stats, self.normalize) - - result = dict() - for key, group in input_stats.items(): - this_dict = dict() - for name, value in group.items(): - this_dict[name] = self.normalize({key:value})[key] - result[key] = this_dict - return result - - -class SingleFieldLinearNormalizer(DictOfTensorMixin): - avaliable_modes = ['limits', 'gaussian'] - - @torch.no_grad() - def fit(self, - data: Union[torch.Tensor, np.ndarray, zarr.Array], - last_n_dims=1, - dtype=torch.float32, - mode='limits', - output_max=1., - output_min=-1., - range_eps=1e-4, - fit_offset=True): - self.params_dict = _fit(data, - last_n_dims=last_n_dims, - dtype=dtype, - mode=mode, - output_max=output_max, - output_min=output_min, - range_eps=range_eps, - fit_offset=fit_offset) - - @classmethod - def create_fit(cls, data: Union[torch.Tensor, np.ndarray, zarr.Array], **kwargs): - obj = cls() - obj.fit(data, **kwargs) - return obj - - @classmethod - def create_manual(cls, - scale: Union[torch.Tensor, np.ndarray], - offset: Union[torch.Tensor, np.ndarray], - input_stats_dict: Dict[str, Union[torch.Tensor, np.ndarray]]): - def to_tensor(x): - if not isinstance(x, torch.Tensor): - x = torch.from_numpy(x) - x = x.flatten() - return x - - # check - for x in [offset] + list(input_stats_dict.values()): - assert x.shape == scale.shape - assert x.dtype == scale.dtype - - params_dict = nn.ParameterDict({ - 'scale': to_tensor(scale), - 'offset': to_tensor(offset), - 'input_stats': nn.ParameterDict( - dict_apply(input_stats_dict, to_tensor)) - }) - return cls(params_dict) - - @classmethod - def create_identity(cls, dtype=torch.float32): - scale = torch.tensor([1], dtype=dtype) - offset = torch.tensor([0], dtype=dtype) - input_stats_dict = { - 'min': torch.tensor([-1], dtype=dtype), - 'max': torch.tensor([1], dtype=dtype), - 'mean': torch.tensor([0], dtype=dtype), - 'std': torch.tensor([1], dtype=dtype) - } - return cls.create_manual(scale, offset, input_stats_dict) - - def normalize(self, x: Union[torch.Tensor, np.ndarray]) -> torch.Tensor: - return _normalize(x, self.params_dict, forward=True) - - def unnormalize(self, x: Union[torch.Tensor, np.ndarray]) -> torch.Tensor: - return _normalize(x, self.params_dict, forward=False) - - def get_input_stats(self): - return self.params_dict['input_stats'] - - def get_output_stats(self): - return dict_apply(self.params_dict['input_stats'], self.normalize) - - def __call__(self, x: Union[torch.Tensor, np.ndarray]) -> torch.Tensor: - return self.normalize(x) - - - -def _fit(data: Union[torch.Tensor, np.ndarray, zarr.Array], - last_n_dims=1, - dtype=torch.float32, - mode='limits', - output_max=1., - output_min=-1., - range_eps=1e-4, - fit_offset=True): - assert mode in ['limits', 'gaussian'] - assert last_n_dims >= 0 - assert output_max > output_min - - # convert data to torch and type - if isinstance(data, zarr.Array): - data = data[:] - if isinstance(data, np.ndarray): - data = torch.from_numpy(data) - if dtype is not None: - data = data.type(dtype) - - # convert shape - dim = 1 - if last_n_dims > 0: - dim = np.prod(data.shape[-last_n_dims:]) - data = data.reshape(-1,dim) - - # compute input stats min max mean std - input_min, _ = data.min(axis=0) - input_max, _ = data.max(axis=0) - input_mean = data.mean(axis=0) - input_std = data.std(axis=0) - - # compute scale and offset - if mode == 'limits': - if fit_offset: - # unit scale - input_range = input_max - input_min - ignore_dim = input_range < range_eps - input_range[ignore_dim] = output_max - output_min - scale = (output_max - output_min) / input_range - offset = output_min - scale * input_min - offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] - # ignore dims scaled to mean of output max and min - else: - # use this when data is pre-zero-centered. - assert output_max > 0 - assert output_min < 0 - # unit abs - output_abs = min(abs(output_min), abs(output_max)) - input_abs = torch.maximum(torch.abs(input_min), torch.abs(input_max)) - ignore_dim = input_abs < range_eps - input_abs[ignore_dim] = output_abs - # don't scale constant channels - scale = output_abs / input_abs - offset = torch.zeros_like(input_mean) - elif mode == 'gaussian': - ignore_dim = input_std < range_eps - scale = input_std.clone() - scale[ignore_dim] = 1 - scale = 1 / scale - - if fit_offset: - offset = - input_mean * scale - else: - offset = torch.zeros_like(input_mean) - - # save - this_params = nn.ParameterDict({ - 'scale': scale, - 'offset': offset, - 'input_stats': nn.ParameterDict({ - 'min': input_min, - 'max': input_max, - 'mean': input_mean, - 'std': input_std - }) - }) - for p in this_params.parameters(): - p.requires_grad_(False) - return this_params - - -def _normalize(x, params, forward=True): - assert 'scale' in params - if isinstance(x, np.ndarray): - x = torch.from_numpy(x) - scale = params['scale'] - offset = params['offset'] - x = x.to(device=scale.device, dtype=scale.dtype) - src_shape = x.shape - x = x.reshape(-1, scale.shape[0]) - if forward: - x = x * scale + offset - else: - x = (x - offset) / scale - x = x.reshape(src_shape) - return x - - -def test(): - data = torch.zeros((100,10,9,2)).uniform_() - data[...,0,0] = 0 - - normalizer = SingleFieldLinearNormalizer() - normalizer.fit(data, mode='limits', last_n_dims=2) - datan = normalizer.normalize(data) - assert datan.shape == data.shape - assert np.allclose(datan.max(), 1.) - assert np.allclose(datan.min(), -1.) - dataun = normalizer.unnormalize(datan) - assert torch.allclose(data, dataun, atol=1e-7) - - input_stats = normalizer.get_input_stats() - output_stats = normalizer.get_output_stats() - - normalizer = SingleFieldLinearNormalizer() - normalizer.fit(data, mode='limits', last_n_dims=1, fit_offset=False) - datan = normalizer.normalize(data) - assert datan.shape == data.shape - assert np.allclose(datan.max(), 1., atol=1e-3) - assert np.allclose(datan.min(), 0., atol=1e-3) - dataun = normalizer.unnormalize(datan) - assert torch.allclose(data, dataun, atol=1e-7) - - data = torch.zeros((100,10,9,2)).uniform_() - normalizer = SingleFieldLinearNormalizer() - normalizer.fit(data, mode='gaussian', last_n_dims=0) - datan = normalizer.normalize(data) - assert datan.shape == data.shape - assert np.allclose(datan.mean(), 0., atol=1e-3) - assert np.allclose(datan.std(), 1., atol=1e-3) - dataun = normalizer.unnormalize(datan) - assert torch.allclose(data, dataun, atol=1e-7) - - - # dict - data = torch.zeros((100,10,9,2)).uniform_() - data[...,0,0] = 0 - - normalizer = LinearNormalizer() - normalizer.fit(data, mode='limits', last_n_dims=2) - datan = normalizer.normalize(data) - assert datan.shape == data.shape - assert np.allclose(datan.max(), 1.) - assert np.allclose(datan.min(), -1.) - dataun = normalizer.unnormalize(datan) - assert torch.allclose(data, dataun, atol=1e-7) - - input_stats = normalizer.get_input_stats() - output_stats = normalizer.get_output_stats() - - data = { - 'obs': torch.zeros((1000,128,9,2)).uniform_() * 512, - 'action': torch.zeros((1000,128,2)).uniform_() * 512 - } - normalizer = LinearNormalizer() - normalizer.fit(data) - datan = normalizer.normalize(data) - dataun = normalizer.unnormalize(datan) - for key in data: - assert torch.allclose(data[key], dataun[key], atol=1e-4) - - input_stats = normalizer.get_input_stats() - output_stats = normalizer.get_output_stats() - - state_dict = normalizer.state_dict() - n = LinearNormalizer() - n.load_state_dict(state_dict) - datan = n.normalize(data) - dataun = n.unnormalize(datan) - for key in data: - assert torch.allclose(data[key], dataun[key], atol=1e-4) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/rotation_transformer.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/rotation_transformer.py deleted file mode 100644 index a801d285f..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/rotation_transformer.py +++ /dev/null @@ -1,103 +0,0 @@ -from typing import Union -import pytorch3d.transforms as pt -import torch -import numpy as np -import functools - -class RotationTransformer: - valid_reps = [ - 'axis_angle', - 'euler_angles', - 'quaternion', - 'rotation_6d', - 'matrix' - ] - - def __init__(self, - from_rep='axis_angle', - to_rep='rotation_6d', - from_convention=None, - to_convention=None): - """ - Valid representations - - Always use matrix as intermediate representation. - """ - assert from_rep != to_rep - assert from_rep in self.valid_reps - assert to_rep in self.valid_reps - if from_rep == 'euler_angles': - assert from_convention is not None - if to_rep == 'euler_angles': - assert to_convention is not None - - forward_funcs = list() - inverse_funcs = list() - - if from_rep != 'matrix': - funcs = [ - getattr(pt, f'{from_rep}_to_matrix'), - getattr(pt, f'matrix_to_{from_rep}') - ] - if from_convention is not None: - funcs = [functools.partial(func, convention=from_convention) - for func in funcs] - forward_funcs.append(funcs[0]) - inverse_funcs.append(funcs[1]) - - if to_rep != 'matrix': - funcs = [ - getattr(pt, f'matrix_to_{to_rep}'), - getattr(pt, f'{to_rep}_to_matrix') - ] - if to_convention is not None: - funcs = [functools.partial(func, convention=to_convention) - for func in funcs] - forward_funcs.append(funcs[0]) - inverse_funcs.append(funcs[1]) - - inverse_funcs = inverse_funcs[::-1] - - self.forward_funcs = forward_funcs - self.inverse_funcs = inverse_funcs - - @staticmethod - def _apply_funcs(x: Union[np.ndarray, torch.Tensor], funcs: list) -> Union[np.ndarray, torch.Tensor]: - x_ = x - if isinstance(x, np.ndarray): - x_ = torch.from_numpy(x) - x_: torch.Tensor - for func in funcs: - x_ = func(x_) - y = x_ - if isinstance(x, np.ndarray): - y = x_.numpy() - return y - - def forward(self, x: Union[np.ndarray, torch.Tensor] - ) -> Union[np.ndarray, torch.Tensor]: - return self._apply_funcs(x, self.forward_funcs) - - def inverse(self, x: Union[np.ndarray, torch.Tensor] - ) -> Union[np.ndarray, torch.Tensor]: - return self._apply_funcs(x, self.inverse_funcs) - - -def test(): - tf = RotationTransformer() - - rotvec = np.random.uniform(-2*np.pi,2*np.pi,size=(1000,3)) - rot6d = tf.forward(rotvec) - new_rotvec = tf.inverse(rot6d) - - from scipy.spatial.transform import Rotation - diff = Rotation.from_rotvec(rotvec) * Rotation.from_rotvec(new_rotvec).inv() - dist = diff.magnitude() - assert dist.max() < 1e-7 - - tf = RotationTransformer('rotation_6d', 'matrix') - rot6d_wrong = rot6d + np.random.normal(scale=0.1, size=rot6d.shape) - mat = tf.forward(rot6d_wrong) - mat_det = np.linalg.det(mat) - assert np.allclose(mat_det, 1) - # rotaiton_6d will be normalized to rotation matrix diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/shape_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/shape_util.py deleted file mode 100644 index e1786c174..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/shape_util.py +++ /dev/null @@ -1,23 +0,0 @@ -from typing import Dict, List, Tuple, Callable -import torch -import torch.nn as nn - -def get_module_device(m: nn.Module): - device = torch.device('cpu') - try: - param = next(iter(m.parameters())) - device = param.device - except StopIteration: - pass - return device - -@torch.no_grad() -def get_output_shape( - input_shape: Tuple[int], - net: Callable[[torch.Tensor], torch.Tensor] - ): - device = get_module_device(net) - test_input = torch.zeros((1,)+tuple(input_shape), device=device) - test_output = net(test_input) - output_shape = tuple(test_output.shape[1:]) - return output_shape diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/tensor_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/tensor_util.py deleted file mode 100644 index 7d6cbffc5..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/common/tensor_util.py +++ /dev/null @@ -1,960 +0,0 @@ -""" -A collection of utilities for working with nested tensor structures consisting -of numpy arrays and torch tensors. -""" -import collections -import numpy as np -import torch - - -def recursive_dict_list_tuple_apply(x, type_func_dict): - """ - Recursively apply functions to a nested dictionary or list or tuple, given a dictionary of - {data_type: function_to_apply}. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - type_func_dict (dict): a mapping from data types to the functions to be - applied for each data type. - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - assert(list not in type_func_dict) - assert(tuple not in type_func_dict) - assert(dict not in type_func_dict) - - if isinstance(x, (dict, collections.OrderedDict)): - new_x = collections.OrderedDict() if isinstance(x, collections.OrderedDict) else dict() - for k, v in x.items(): - new_x[k] = recursive_dict_list_tuple_apply(v, type_func_dict) - return new_x - elif isinstance(x, (list, tuple)): - ret = [recursive_dict_list_tuple_apply(v, type_func_dict) for v in x] - if isinstance(x, tuple): - ret = tuple(ret) - return ret - else: - for t, f in type_func_dict.items(): - if isinstance(x, t): - return f(x) - else: - raise NotImplementedError( - 'Cannot handle data type %s' % str(type(x))) - - -def map_tensor(x, func): - """ - Apply function @func to torch.Tensor objects in a nested dictionary or - list or tuple. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - func (function): function to apply to each tensor - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: func, - type(None): lambda x: x, - } - ) - - -def map_ndarray(x, func): - """ - Apply function @func to np.ndarray objects in a nested dictionary or - list or tuple. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - func (function): function to apply to each array - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - np.ndarray: func, - type(None): lambda x: x, - } - ) - - -def map_tensor_ndarray(x, tensor_func, ndarray_func): - """ - Apply function @tensor_func to torch.Tensor objects and @ndarray_func to - np.ndarray objects in a nested dictionary or list or tuple. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - tensor_func (function): function to apply to each tensor - ndarray_Func (function): function to apply to each array - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: tensor_func, - np.ndarray: ndarray_func, - type(None): lambda x: x, - } - ) - - -def clone(x): - """ - Clones all torch tensors and numpy arrays in nested dictionary or list - or tuple and returns a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x: x.clone(), - np.ndarray: lambda x: x.copy(), - type(None): lambda x: x, - } - ) - - -def detach(x): - """ - Detaches all torch tensors in nested dictionary or list - or tuple and returns a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x: x.detach(), - } - ) - - -def to_batch(x): - """ - Introduces a leading batch dimension of 1 for all torch tensors and numpy - arrays in nested dictionary or list or tuple and returns a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x: x[None, ...], - np.ndarray: lambda x: x[None, ...], - type(None): lambda x: x, - } - ) - - -def to_sequence(x): - """ - Introduces a time dimension of 1 at dimension 1 for all torch tensors and numpy - arrays in nested dictionary or list or tuple and returns a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x: x[:, None, ...], - np.ndarray: lambda x: x[:, None, ...], - type(None): lambda x: x, - } - ) - - -def index_at_time(x, ind): - """ - Indexes all torch tensors and numpy arrays in dimension 1 with index @ind in - nested dictionary or list or tuple and returns a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - ind (int): index - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x: x[:, ind, ...], - np.ndarray: lambda x: x[:, ind, ...], - type(None): lambda x: x, - } - ) - - -def unsqueeze(x, dim): - """ - Adds dimension of size 1 at dimension @dim in all torch tensors and numpy arrays - in nested dictionary or list or tuple and returns a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - dim (int): dimension - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x: x.unsqueeze(dim=dim), - np.ndarray: lambda x: np.expand_dims(x, axis=dim), - type(None): lambda x: x, - } - ) - - -def contiguous(x): - """ - Makes all torch tensors and numpy arrays contiguous in nested dictionary or - list or tuple and returns a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x: x.contiguous(), - np.ndarray: lambda x: np.ascontiguousarray(x), - type(None): lambda x: x, - } - ) - - -def to_device(x, device): - """ - Sends all torch tensors in nested dictionary or list or tuple to device - @device, and returns a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - device (torch.Device): device to send tensors to - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x, d=device: x.to(d), - type(None): lambda x: x, - } - ) - - -def to_tensor(x): - """ - Converts all numpy arrays in nested dictionary or list or tuple to - torch tensors (and leaves existing torch Tensors as-is), and returns - a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x: x, - np.ndarray: lambda x: torch.from_numpy(x), - type(None): lambda x: x, - } - ) - - -def to_numpy(x): - """ - Converts all torch tensors in nested dictionary or list or tuple to - numpy (and leaves existing numpy arrays as-is), and returns - a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - def f(tensor): - if tensor.is_cuda: - return tensor.detach().cpu().numpy() - else: - return tensor.detach().numpy() - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: f, - np.ndarray: lambda x: x, - type(None): lambda x: x, - } - ) - - -def to_list(x): - """ - Converts all torch tensors and numpy arrays in nested dictionary or list - or tuple to a list, and returns a new nested structure. Useful for - json encoding. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - def f(tensor): - if tensor.is_cuda: - return tensor.detach().cpu().numpy().tolist() - else: - return tensor.detach().numpy().tolist() - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: f, - np.ndarray: lambda x: x.tolist(), - type(None): lambda x: x, - } - ) - - -def to_float(x): - """ - Converts all torch tensors and numpy arrays in nested dictionary or list - or tuple to float type entries, and returns a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x: x.float(), - np.ndarray: lambda x: x.astype(np.float32), - type(None): lambda x: x, - } - ) - - -def to_uint8(x): - """ - Converts all torch tensors and numpy arrays in nested dictionary or list - or tuple to uint8 type entries, and returns a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x: x.byte(), - np.ndarray: lambda x: x.astype(np.uint8), - type(None): lambda x: x, - } - ) - - -def to_torch(x, device): - """ - Converts all numpy arrays and torch tensors in nested dictionary or list or tuple to - torch tensors on device @device and returns a new nested structure. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - device (torch.Device): device to send tensors to - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return to_device(to_float(to_tensor(x)), device) - - -def to_one_hot_single(tensor, num_class): - """ - Convert tensor to one-hot representation, assuming a certain number of total class labels. - - Args: - tensor (torch.Tensor): tensor containing integer labels - num_class (int): number of classes - - Returns: - x (torch.Tensor): tensor containing one-hot representation of labels - """ - x = torch.zeros(tensor.size() + (num_class,)).to(tensor.device) - x.scatter_(-1, tensor.unsqueeze(-1), 1) - return x - - -def to_one_hot(tensor, num_class): - """ - Convert all tensors in nested dictionary or list or tuple to one-hot representation, - assuming a certain number of total class labels. - - Args: - tensor (dict or list or tuple): a possibly nested dictionary or list or tuple - num_class (int): number of classes - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return map_tensor(tensor, func=lambda x, nc=num_class: to_one_hot_single(x, nc)) - - -def flatten_single(x, begin_axis=1): - """ - Flatten a tensor in all dimensions from @begin_axis onwards. - - Args: - x (torch.Tensor): tensor to flatten - begin_axis (int): which axis to flatten from - - Returns: - y (torch.Tensor): flattened tensor - """ - fixed_size = x.size()[:begin_axis] - _s = list(fixed_size) + [-1] - return x.reshape(*_s) - - -def flatten(x, begin_axis=1): - """ - Flatten all tensors in nested dictionary or list or tuple, from @begin_axis onwards. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - begin_axis (int): which axis to flatten from - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x, b=begin_axis: flatten_single(x, begin_axis=b), - } - ) - - -def reshape_dimensions_single(x, begin_axis, end_axis, target_dims): - """ - Reshape selected dimensions in a tensor to a target dimension. - - Args: - x (torch.Tensor): tensor to reshape - begin_axis (int): begin dimension - end_axis (int): end dimension - target_dims (tuple or list): target shape for the range of dimensions - (@begin_axis, @end_axis) - - Returns: - y (torch.Tensor): reshaped tensor - """ - assert(begin_axis <= end_axis) - assert(begin_axis >= 0) - assert(end_axis < len(x.shape)) - assert(isinstance(target_dims, (tuple, list))) - s = x.shape - final_s = [] - for i in range(len(s)): - if i == begin_axis: - final_s.extend(target_dims) - elif i < begin_axis or i > end_axis: - final_s.append(s[i]) - return x.reshape(*final_s) - - -def reshape_dimensions(x, begin_axis, end_axis, target_dims): - """ - Reshape selected dimensions for all tensors in nested dictionary or list or tuple - to a target dimension. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - begin_axis (int): begin dimension - end_axis (int): end dimension - target_dims (tuple or list): target shape for the range of dimensions - (@begin_axis, @end_axis) - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x, b=begin_axis, e=end_axis, t=target_dims: reshape_dimensions_single( - x, begin_axis=b, end_axis=e, target_dims=t), - np.ndarray: lambda x, b=begin_axis, e=end_axis, t=target_dims: reshape_dimensions_single( - x, begin_axis=b, end_axis=e, target_dims=t), - type(None): lambda x: x, - } - ) - - -def join_dimensions(x, begin_axis, end_axis): - """ - Joins all dimensions between dimensions (@begin_axis, @end_axis) into a flat dimension, for - all tensors in nested dictionary or list or tuple. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - begin_axis (int): begin dimension - end_axis (int): end dimension - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x, b=begin_axis, e=end_axis: reshape_dimensions_single( - x, begin_axis=b, end_axis=e, target_dims=[-1]), - np.ndarray: lambda x, b=begin_axis, e=end_axis: reshape_dimensions_single( - x, begin_axis=b, end_axis=e, target_dims=[-1]), - type(None): lambda x: x, - } - ) - - -def expand_at_single(x, size, dim): - """ - Expand a tensor at a single dimension @dim by @size - - Args: - x (torch.Tensor): input tensor - size (int): size to expand - dim (int): dimension to expand - - Returns: - y (torch.Tensor): expanded tensor - """ - assert dim < x.ndimension() - assert x.shape[dim] == 1 - expand_dims = [-1] * x.ndimension() - expand_dims[dim] = size - return x.expand(*expand_dims) - - -def expand_at(x, size, dim): - """ - Expand all tensors in nested dictionary or list or tuple at a single - dimension @dim by @size. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - size (int): size to expand - dim (int): dimension to expand - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return map_tensor(x, lambda t, s=size, d=dim: expand_at_single(t, s, d)) - - -def unsqueeze_expand_at(x, size, dim): - """ - Unsqueeze and expand a tensor at a dimension @dim by @size. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - size (int): size to expand - dim (int): dimension to unsqueeze and expand - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - x = unsqueeze(x, dim) - return expand_at(x, size, dim) - - -def repeat_by_expand_at(x, repeats, dim): - """ - Repeat a dimension by combining expand and reshape operations. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - repeats (int): number of times to repeat the target dimension - dim (int): dimension to repeat on - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - x = unsqueeze_expand_at(x, repeats, dim + 1) - return join_dimensions(x, dim, dim + 1) - - -def named_reduce_single(x, reduction, dim): - """ - Reduce tensor at a dimension by named reduction functions. - - Args: - x (torch.Tensor): tensor to be reduced - reduction (str): one of ["sum", "max", "mean", "flatten"] - dim (int): dimension to be reduced (or begin axis for flatten) - - Returns: - y (torch.Tensor): reduced tensor - """ - assert x.ndimension() > dim - assert reduction in ["sum", "max", "mean", "flatten"] - if reduction == "flatten": - x = flatten(x, begin_axis=dim) - elif reduction == "max": - x = torch.max(x, dim=dim)[0] # [B, D] - elif reduction == "sum": - x = torch.sum(x, dim=dim) - else: - x = torch.mean(x, dim=dim) - return x - - -def named_reduce(x, reduction, dim): - """ - Reduces all tensors in nested dictionary or list or tuple at a dimension - using a named reduction function. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - reduction (str): one of ["sum", "max", "mean", "flatten"] - dim (int): dimension to be reduced (or begin axis for flatten) - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return map_tensor(x, func=lambda t, r=reduction, d=dim: named_reduce_single(t, r, d)) - - -def gather_along_dim_with_dim_single(x, target_dim, source_dim, indices): - """ - This function indexes out a target dimension of a tensor in a structured way, - by allowing a different value to be selected for each member of a flat index - tensor (@indices) corresponding to a source dimension. This can be interpreted - as moving along the source dimension, using the corresponding index value - in @indices to select values for all other dimensions outside of the - source and target dimensions. A common use case is to gather values - in target dimension 1 for each batch member (target dimension 0). - - Args: - x (torch.Tensor): tensor to gather values for - target_dim (int): dimension to gather values along - source_dim (int): dimension to hold constant and use for gathering values - from the other dimensions - indices (torch.Tensor): flat index tensor with same shape as tensor @x along - @source_dim - - Returns: - y (torch.Tensor): gathered tensor, with dimension @target_dim indexed out - """ - assert len(indices.shape) == 1 - assert x.shape[source_dim] == indices.shape[0] - - # unsqueeze in all dimensions except the source dimension - new_shape = [1] * x.ndimension() - new_shape[source_dim] = -1 - indices = indices.reshape(*new_shape) - - # repeat in all dimensions - but preserve shape of source dimension, - # and make sure target_dimension has singleton dimension - expand_shape = list(x.shape) - expand_shape[source_dim] = -1 - expand_shape[target_dim] = 1 - indices = indices.expand(*expand_shape) - - out = x.gather(dim=target_dim, index=indices) - return out.squeeze(target_dim) - - -def gather_along_dim_with_dim(x, target_dim, source_dim, indices): - """ - Apply @gather_along_dim_with_dim_single to all tensors in a nested - dictionary or list or tuple. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - target_dim (int): dimension to gather values along - source_dim (int): dimension to hold constant and use for gathering values - from the other dimensions - indices (torch.Tensor): flat index tensor with same shape as tensor @x along - @source_dim - - Returns: - y (dict or list or tuple): new nested dict-list-tuple - """ - return map_tensor(x, - lambda y, t=target_dim, s=source_dim, i=indices: gather_along_dim_with_dim_single(y, t, s, i)) - - -def gather_sequence_single(seq, indices): - """ - Given a tensor with leading dimensions [B, T, ...], gather an element from each sequence in - the batch given an index for each sequence. - - Args: - seq (torch.Tensor): tensor with leading dimensions [B, T, ...] - indices (torch.Tensor): tensor indices of shape [B] - - Return: - y (torch.Tensor): indexed tensor of shape [B, ....] - """ - return gather_along_dim_with_dim_single(seq, target_dim=1, source_dim=0, indices=indices) - - -def gather_sequence(seq, indices): - """ - Given a nested dictionary or list or tuple, gathers an element from each sequence of the batch - for tensors with leading dimensions [B, T, ...]. - - Args: - seq (dict or list or tuple): a possibly nested dictionary or list or tuple with tensors - of leading dimensions [B, T, ...] - indices (torch.Tensor): tensor indices of shape [B] - - Returns: - y (dict or list or tuple): new nested dict-list-tuple with tensors of shape [B, ...] - """ - return gather_along_dim_with_dim(seq, target_dim=1, source_dim=0, indices=indices) - - -def pad_sequence_single(seq, padding, batched=False, pad_same=True, pad_values=None): - """ - Pad input tensor or array @seq in the time dimension (dimension 1). - - Args: - seq (np.ndarray or torch.Tensor): sequence to be padded - padding (tuple): begin and end padding, e.g. [1, 1] pads both begin and end of the sequence by 1 - batched (bool): if sequence has the batch dimension - pad_same (bool): if pad by duplicating - pad_values (scalar or (ndarray, Tensor)): values to be padded if not pad_same - - Returns: - padded sequence (np.ndarray or torch.Tensor) - """ - assert isinstance(seq, (np.ndarray, torch.Tensor)) - assert pad_same or pad_values is not None - if pad_values is not None: - assert isinstance(pad_values, float) - repeat_func = np.repeat if isinstance(seq, np.ndarray) else torch.repeat_interleave - concat_func = np.concatenate if isinstance(seq, np.ndarray) else torch.cat - ones_like_func = np.ones_like if isinstance(seq, np.ndarray) else torch.ones_like - seq_dim = 1 if batched else 0 - - begin_pad = [] - end_pad = [] - - if padding[0] > 0: - pad = seq[[0]] if pad_same else ones_like_func(seq[[0]]) * pad_values - begin_pad.append(repeat_func(pad, padding[0], seq_dim)) - if padding[1] > 0: - pad = seq[[-1]] if pad_same else ones_like_func(seq[[-1]]) * pad_values - end_pad.append(repeat_func(pad, padding[1], seq_dim)) - - return concat_func(begin_pad + [seq] + end_pad, seq_dim) - - -def pad_sequence(seq, padding, batched=False, pad_same=True, pad_values=None): - """ - Pad a nested dictionary or list or tuple of sequence tensors in the time dimension (dimension 1). - - Args: - seq (dict or list or tuple): a possibly nested dictionary or list or tuple with tensors - of leading dimensions [B, T, ...] - padding (tuple): begin and end padding, e.g. [1, 1] pads both begin and end of the sequence by 1 - batched (bool): if sequence has the batch dimension - pad_same (bool): if pad by duplicating - pad_values (scalar or (ndarray, Tensor)): values to be padded if not pad_same - - Returns: - padded sequence (dict or list or tuple) - """ - return recursive_dict_list_tuple_apply( - seq, - { - torch.Tensor: lambda x, p=padding, b=batched, ps=pad_same, pv=pad_values: - pad_sequence_single(x, p, b, ps, pv), - np.ndarray: lambda x, p=padding, b=batched, ps=pad_same, pv=pad_values: - pad_sequence_single(x, p, b, ps, pv), - type(None): lambda x: x, - } - ) - - -def assert_size_at_dim_single(x, size, dim, msg): - """ - Ensure that array or tensor @x has size @size in dim @dim. - - Args: - x (np.ndarray or torch.Tensor): input array or tensor - size (int): size that tensors should have at @dim - dim (int): dimension to check - msg (str): text to display if assertion fails - """ - assert x.shape[dim] == size, msg - - -def assert_size_at_dim(x, size, dim, msg): - """ - Ensure that arrays and tensors in nested dictionary or list or tuple have - size @size in dim @dim. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - size (int): size that tensors should have at @dim - dim (int): dimension to check - """ - map_tensor(x, lambda t, s=size, d=dim, m=msg: assert_size_at_dim_single(t, s, d, m)) - - -def get_shape(x): - """ - Get all shapes of arrays and tensors in nested dictionary or list or tuple. - - Args: - x (dict or list or tuple): a possibly nested dictionary or list or tuple - - Returns: - y (dict or list or tuple): new nested dict-list-tuple that contains each array or - tensor's shape - """ - return recursive_dict_list_tuple_apply( - x, - { - torch.Tensor: lambda x: x.shape, - np.ndarray: lambda x: x.shape, - type(None): lambda x: x, - } - ) - - -def list_of_flat_dict_to_dict_of_list(list_of_dict): - """ - Helper function to go from a list of flat dictionaries to a dictionary of lists. - By "flat" we mean that none of the values are dictionaries, but are numpy arrays, - floats, etc. - - Args: - list_of_dict (list): list of flat dictionaries - - Returns: - dict_of_list (dict): dictionary of lists - """ - assert isinstance(list_of_dict, list) - dic = collections.OrderedDict() - for i in range(len(list_of_dict)): - for k in list_of_dict[i]: - if k not in dic: - dic[k] = [] - dic[k].append(list_of_dict[i][k]) - return dic - - -def flatten_nested_dict_list(d, parent_key='', sep='_', item_key=''): - """ - Flatten a nested dict or list to a list. - - For example, given a dict - { - a: 1 - b: { - c: 2 - } - c: 3 - } - - the function would return [(a, 1), (b_c, 2), (c, 3)] - - Args: - d (dict, list): a nested dict or list to be flattened - parent_key (str): recursion helper - sep (str): separator for nesting keys - item_key (str): recursion helper - Returns: - list: a list of (key, value) tuples - """ - items = [] - if isinstance(d, (tuple, list)): - new_key = parent_key + sep + item_key if len(parent_key) > 0 else item_key - for i, v in enumerate(d): - items.extend(flatten_nested_dict_list(v, new_key, sep=sep, item_key=str(i))) - return items - elif isinstance(d, dict): - new_key = parent_key + sep + item_key if len(parent_key) > 0 else item_key - for k, v in d.items(): - assert isinstance(k, str) - items.extend(flatten_nested_dict_list(v, new_key, sep=sep, item_key=k)) - return items - else: - new_key = parent_key + sep + item_key if len(parent_key) > 0 else item_key - return [(new_key, d)] - - -def time_distributed(inputs, op, activation=None, inputs_as_kwargs=False, inputs_as_args=False, **kwargs): - """ - Apply function @op to all tensors in nested dictionary or list or tuple @inputs in both the - batch (B) and time (T) dimension, where the tensors are expected to have shape [B, T, ...]. - Will do this by reshaping tensors to [B * T, ...], passing through the op, and then reshaping - outputs to [B, T, ...]. - - Args: - inputs (list or tuple or dict): a possibly nested dictionary or list or tuple with tensors - of leading dimensions [B, T, ...] - op: a layer op that accepts inputs - activation: activation to apply at the output - inputs_as_kwargs (bool): whether to feed input as a kwargs dict to the op - inputs_as_args (bool) whether to feed input as a args list to the op - kwargs (dict): other kwargs to supply to the op - - Returns: - outputs (dict or list or tuple): new nested dict-list-tuple with tensors of leading dimension [B, T]. - """ - batch_size, seq_len = flatten_nested_dict_list(inputs)[0][1].shape[:2] - inputs = join_dimensions(inputs, 0, 1) - if inputs_as_kwargs: - outputs = op(**inputs, **kwargs) - elif inputs_as_args: - outputs = op(*inputs, **kwargs) - else: - outputs = op(inputs, **kwargs) - - if activation is not None: - outputs = map_tensor(outputs, activation) - outputs = reshape_dimensions(outputs, begin_axis=0, end_axis=0, target_dims=(batch_size, seq_len)) - return outputs diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/conditional_unet1d.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/conditional_unet1d.py deleted file mode 100644 index fea22f5ea..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/conditional_unet1d.py +++ /dev/null @@ -1,242 +0,0 @@ -from typing import Union -import logging -import torch -import torch.nn as nn -import einops -from einops.layers.torch import Rearrange - -from diffusion_policy.model.diffusion.conv1d_components import ( - Downsample1d, Upsample1d, Conv1dBlock) -from diffusion_policy.model.diffusion.positional_embedding import SinusoidalPosEmb - -logger = logging.getLogger(__name__) - -class ConditionalResidualBlock1D(nn.Module): - def __init__(self, - in_channels, - out_channels, - cond_dim, - kernel_size=3, - n_groups=8, - cond_predict_scale=False): - super().__init__() - - self.blocks = nn.ModuleList([ - Conv1dBlock(in_channels, out_channels, kernel_size, n_groups=n_groups), - Conv1dBlock(out_channels, out_channels, kernel_size, n_groups=n_groups), - ]) - - # FiLM modulation https://arxiv.org/abs/1709.07871 - # predicts per-channel scale and bias - cond_channels = out_channels - if cond_predict_scale: - cond_channels = out_channels * 2 - self.cond_predict_scale = cond_predict_scale - self.out_channels = out_channels - self.cond_encoder = nn.Sequential( - nn.Mish(), - nn.Linear(cond_dim, cond_channels), - Rearrange('batch t -> batch t 1'), - ) - - # make sure dimensions compatible - self.residual_conv = nn.Conv1d(in_channels, out_channels, 1) \ - if in_channels != out_channels else nn.Identity() - - def forward(self, x, cond): - ''' - x : [ batch_size x in_channels x horizon ] - cond : [ batch_size x cond_dim] - - returns: - out : [ batch_size x out_channels x horizon ] - ''' - out = self.blocks[0](x) - embed = self.cond_encoder(cond) - if self.cond_predict_scale: - embed = embed.reshape( - embed.shape[0], 2, self.out_channels, 1) - scale = embed[:,0,...] - bias = embed[:,1,...] - out = scale * out + bias - else: - out = out + embed - out = self.blocks[1](out) - out = out + self.residual_conv(x) - return out - - -class ConditionalUnet1D(nn.Module): - def __init__(self, - input_dim, - local_cond_dim=None, - global_cond_dim=None, - diffusion_step_embed_dim=256, - down_dims=[256,512,1024], - kernel_size=3, - n_groups=8, - cond_predict_scale=False - ): - super().__init__() - all_dims = [input_dim] + list(down_dims) - start_dim = down_dims[0] - - dsed = diffusion_step_embed_dim - diffusion_step_encoder = nn.Sequential( - SinusoidalPosEmb(dsed), - nn.Linear(dsed, dsed * 4), - nn.Mish(), - nn.Linear(dsed * 4, dsed), - ) - cond_dim = dsed - if global_cond_dim is not None: - cond_dim += global_cond_dim - - in_out = list(zip(all_dims[:-1], all_dims[1:])) - - local_cond_encoder = None - if local_cond_dim is not None: - _, dim_out = in_out[0] - dim_in = local_cond_dim - local_cond_encoder = nn.ModuleList([ - # down encoder - ConditionalResidualBlock1D( - dim_in, dim_out, cond_dim=cond_dim, - kernel_size=kernel_size, n_groups=n_groups, - cond_predict_scale=cond_predict_scale), - # up encoder - ConditionalResidualBlock1D( - dim_in, dim_out, cond_dim=cond_dim, - kernel_size=kernel_size, n_groups=n_groups, - cond_predict_scale=cond_predict_scale) - ]) - - mid_dim = all_dims[-1] - self.mid_modules = nn.ModuleList([ - ConditionalResidualBlock1D( - mid_dim, mid_dim, cond_dim=cond_dim, - kernel_size=kernel_size, n_groups=n_groups, - cond_predict_scale=cond_predict_scale - ), - ConditionalResidualBlock1D( - mid_dim, mid_dim, cond_dim=cond_dim, - kernel_size=kernel_size, n_groups=n_groups, - cond_predict_scale=cond_predict_scale - ), - ]) - - down_modules = nn.ModuleList([]) - for ind, (dim_in, dim_out) in enumerate(in_out): - is_last = ind >= (len(in_out) - 1) - down_modules.append(nn.ModuleList([ - ConditionalResidualBlock1D( - dim_in, dim_out, cond_dim=cond_dim, - kernel_size=kernel_size, n_groups=n_groups, - cond_predict_scale=cond_predict_scale), - ConditionalResidualBlock1D( - dim_out, dim_out, cond_dim=cond_dim, - kernel_size=kernel_size, n_groups=n_groups, - cond_predict_scale=cond_predict_scale), - Downsample1d(dim_out) if not is_last else nn.Identity() - ])) - - up_modules = nn.ModuleList([]) - for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): - is_last = ind >= (len(in_out) - 1) - up_modules.append(nn.ModuleList([ - ConditionalResidualBlock1D( - dim_out*2, dim_in, cond_dim=cond_dim, - kernel_size=kernel_size, n_groups=n_groups, - cond_predict_scale=cond_predict_scale), - ConditionalResidualBlock1D( - dim_in, dim_in, cond_dim=cond_dim, - kernel_size=kernel_size, n_groups=n_groups, - cond_predict_scale=cond_predict_scale), - Upsample1d(dim_in) if not is_last else nn.Identity() - ])) - - final_conv = nn.Sequential( - Conv1dBlock(start_dim, start_dim, kernel_size=kernel_size), - nn.Conv1d(start_dim, input_dim, 1), - ) - - self.diffusion_step_encoder = diffusion_step_encoder - self.local_cond_encoder = local_cond_encoder - self.up_modules = up_modules - self.down_modules = down_modules - self.final_conv = final_conv - - logger.info( - "number of parameters: %e", sum(p.numel() for p in self.parameters()) - ) - - def forward(self, - sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], - local_cond=None, global_cond=None, **kwargs): - """ - x: (B,T,input_dim) - timestep: (B,) or int, diffusion step - local_cond: (B,T,local_cond_dim) - global_cond: (B,global_cond_dim) - output: (B,T,input_dim) - """ - sample = einops.rearrange(sample, 'b h t -> b t h') - - # 1. time - timesteps = timestep - if not torch.is_tensor(timesteps): - # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can - timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device) - elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - - global_feature = self.diffusion_step_encoder(timesteps) - - if global_cond is not None: - global_feature = torch.cat([ - global_feature, global_cond - ], axis=-1) - - # encode local features - h_local = list() - if local_cond is not None: - local_cond = einops.rearrange(local_cond, 'b h t -> b t h') - resnet, resnet2 = self.local_cond_encoder - x = resnet(local_cond, global_feature) - h_local.append(x) - x = resnet2(local_cond, global_feature) - h_local.append(x) - - x = sample - h = [] - for idx, (resnet, resnet2, downsample) in enumerate(self.down_modules): - x = resnet(x, global_feature) - if idx == 0 and len(h_local) > 0: - x = x + h_local[0] - x = resnet2(x, global_feature) - h.append(x) - x = downsample(x) - - for mid_module in self.mid_modules: - x = mid_module(x, global_feature) - - for idx, (resnet, resnet2, upsample) in enumerate(self.up_modules): - x = torch.cat((x, h.pop()), dim=1) - x = resnet(x, global_feature) - # The correct condition should be: - # if idx == (len(self.up_modules)-1) and len(h_local) > 0: - # However this change will break compatibility with published checkpoints. - # Therefore it is left as a comment. - if idx == len(self.up_modules) and len(h_local) > 0: - x = x + h_local[1] - x = resnet2(x, global_feature) - x = upsample(x) - - x = self.final_conv(x) - - x = einops.rearrange(x, 'b t h -> b h t') - return x - diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/conv1d_components.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/conv1d_components.py deleted file mode 100644 index 1c4cfc925..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/conv1d_components.py +++ /dev/null @@ -1,46 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -# from einops.layers.torch import Rearrange - - -class Downsample1d(nn.Module): - def __init__(self, dim): - super().__init__() - self.conv = nn.Conv1d(dim, dim, 3, 2, 1) - - def forward(self, x): - return self.conv(x) - -class Upsample1d(nn.Module): - def __init__(self, dim): - super().__init__() - self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1) - - def forward(self, x): - return self.conv(x) - -class Conv1dBlock(nn.Module): - ''' - Conv1d --> GroupNorm --> Mish - ''' - - def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8): - super().__init__() - - self.block = nn.Sequential( - nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2), - # Rearrange('batch channels horizon -> batch channels 1 horizon'), - nn.GroupNorm(n_groups, out_channels), - # Rearrange('batch channels 1 horizon -> batch channels horizon'), - nn.Mish(), - ) - - def forward(self, x): - return self.block(x) - - -def test(): - cb = Conv1dBlock(256, 128, kernel_size=3) - x = torch.zeros((1,256,16)) - o = cb(x) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/ema_model.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/ema_model.py deleted file mode 100644 index d916c2b12..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/ema_model.py +++ /dev/null @@ -1,88 +0,0 @@ -import copy -import torch -from torch.nn.modules.batchnorm import _BatchNorm - -class EMAModel: - """ - Exponential Moving Average of models weights - """ - - def __init__( - self, - model, - update_after_step=0, - inv_gamma=1.0, - power=2 / 3, - min_value=0.0, - max_value=0.9999 - ): - """ - @crowsonkb's notes on EMA Warmup: - If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan - to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps), - gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 - at 215.4k steps). - Args: - inv_gamma (float): Inverse multiplicative factor of EMA warmup. Default: 1. - power (float): Exponential factor of EMA warmup. Default: 2/3. - min_value (float): The minimum EMA decay rate. Default: 0. - """ - - self.averaged_model = model - self.averaged_model.eval() - self.averaged_model.requires_grad_(False) - - self.update_after_step = update_after_step - self.inv_gamma = inv_gamma - self.power = power - self.min_value = min_value - self.max_value = max_value - - self.decay = 0.0 - self.optimization_step = 0 - - def get_decay(self, optimization_step): - """ - Compute the decay factor for the exponential moving average. - """ - step = max(0, optimization_step - self.update_after_step - 1) - value = 1 - (1 + step / self.inv_gamma) ** -self.power - - if step <= 0: - return 0.0 - - return max(self.min_value, min(value, self.max_value)) - - @torch.no_grad() - def step(self, new_model): - self.decay = self.get_decay(self.optimization_step) - - # old_all_dataptrs = set() - # for param in new_model.parameters(): - # data_ptr = param.data_ptr() - # if data_ptr != 0: - # old_all_dataptrs.add(data_ptr) - - all_dataptrs = set() - for module, ema_module in zip(new_model.modules(), self.averaged_model.modules()): - for param, ema_param in zip(module.parameters(recurse=False), ema_module.parameters(recurse=False)): - # iterative over immediate parameters only. - if isinstance(param, dict): - raise RuntimeError('Dict parameter not supported') - - # data_ptr = param.data_ptr() - # if data_ptr != 0: - # all_dataptrs.add(data_ptr) - - if isinstance(module, _BatchNorm): - # skip batchnorms - ema_param.copy_(param.to(dtype=ema_param.dtype).data) - elif not param.requires_grad: - ema_param.copy_(param.to(dtype=ema_param.dtype).data) - else: - ema_param.mul_(self.decay) - ema_param.add_(param.data.to(dtype=ema_param.dtype), alpha=1 - self.decay) - - # verify that iterating over module and then parameters is identical to parameters recursively. - # assert old_all_dataptrs == all_dataptrs - self.optimization_step += 1 diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/mask_generator.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/mask_generator.py deleted file mode 100644 index fc5bb7d46..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/mask_generator.py +++ /dev/null @@ -1,221 +0,0 @@ -from typing import Sequence, Optional -import torch -from torch import nn -from diffusion_policy.model.common.module_attr_mixin import ModuleAttrMixin - - -def get_intersection_slice_mask( - shape: tuple, - dim_slices: Sequence[slice], - device: Optional[torch.device]=None - ): - assert(len(shape) == len(dim_slices)) - mask = torch.zeros(size=shape, dtype=torch.bool, device=device) - mask[dim_slices] = True - return mask - - -def get_union_slice_mask( - shape: tuple, - dim_slices: Sequence[slice], - device: Optional[torch.device]=None - ): - assert(len(shape) == len(dim_slices)) - mask = torch.zeros(size=shape, dtype=torch.bool, device=device) - for i in range(len(dim_slices)): - this_slices = [slice(None)] * len(shape) - this_slices[i] = dim_slices[i] - mask[this_slices] = True - return mask - - -class DummyMaskGenerator(ModuleAttrMixin): - def __init__(self): - super().__init__() - - @torch.no_grad() - def forward(self, shape): - device = self.device - mask = torch.ones(size=shape, dtype=torch.bool, device=device) - return mask - - -class LowdimMaskGenerator(ModuleAttrMixin): - def __init__(self, - action_dim, obs_dim, - # obs mask setup - max_n_obs_steps=2, - fix_obs_steps=True, - # action mask - action_visible=False - ): - super().__init__() - self.action_dim = action_dim - self.obs_dim = obs_dim - self.max_n_obs_steps = max_n_obs_steps - self.fix_obs_steps = fix_obs_steps - self.action_visible = action_visible - - @torch.no_grad() - def forward(self, shape, seed=None): - device = self.device - B, T, D = shape - assert D == (self.action_dim + self.obs_dim) - - # create all tensors on this device - rng = torch.Generator(device=device) - if seed is not None: - rng = rng.manual_seed(seed) - - # generate dim mask - dim_mask = torch.zeros(size=shape, - dtype=torch.bool, device=device) - is_action_dim = dim_mask.clone() - is_action_dim[...,:self.action_dim] = True - is_obs_dim = ~is_action_dim - - # generate obs mask - if self.fix_obs_steps: - obs_steps = torch.full((B,), - fill_value=self.max_n_obs_steps, device=device) - else: - obs_steps = torch.randint( - low=1, high=self.max_n_obs_steps+1, - size=(B,), generator=rng, device=device) - - steps = torch.arange(0, T, device=device).reshape(1,T).expand(B,T) - obs_mask = (steps.T < obs_steps).T.reshape(B,T,1).expand(B,T,D) - obs_mask = obs_mask & is_obs_dim - - # generate action mask - if self.action_visible: - action_steps = torch.maximum( - obs_steps - 1, - torch.tensor(0, - dtype=obs_steps.dtype, - device=obs_steps.device)) - action_mask = (steps.T < action_steps).T.reshape(B,T,1).expand(B,T,D) - action_mask = action_mask & is_action_dim - - mask = obs_mask - if self.action_visible: - mask = mask | action_mask - - return mask - - -class KeypointMaskGenerator(ModuleAttrMixin): - def __init__(self, - # dimensions - action_dim, keypoint_dim, - # obs mask setup - max_n_obs_steps=2, fix_obs_steps=True, - # keypoint mask setup - keypoint_visible_rate=0.7, time_independent=False, - # action mask - action_visible=False, - context_dim=0, # dim for context - n_context_steps=1 - ): - super().__init__() - self.action_dim = action_dim - self.keypoint_dim = keypoint_dim - self.context_dim = context_dim - self.max_n_obs_steps = max_n_obs_steps - self.fix_obs_steps = fix_obs_steps - self.keypoint_visible_rate = keypoint_visible_rate - self.time_independent = time_independent - self.action_visible = action_visible - self.n_context_steps = n_context_steps - - @torch.no_grad() - def forward(self, shape, seed=None): - device = self.device - B, T, D = shape - all_keypoint_dims = D - self.action_dim - self.context_dim - n_keypoints = all_keypoint_dims // self.keypoint_dim - - # create all tensors on this device - rng = torch.Generator(device=device) - if seed is not None: - rng = rng.manual_seed(seed) - - # generate dim mask - dim_mask = torch.zeros(size=shape, - dtype=torch.bool, device=device) - is_action_dim = dim_mask.clone() - is_action_dim[...,:self.action_dim] = True - is_context_dim = dim_mask.clone() - if self.context_dim > 0: - is_context_dim[...,-self.context_dim:] = True - is_obs_dim = ~(is_action_dim | is_context_dim) - # assumption trajectory=cat([action, keypoints, context], dim=-1) - - # generate obs mask - if self.fix_obs_steps: - obs_steps = torch.full((B,), - fill_value=self.max_n_obs_steps, device=device) - else: - obs_steps = torch.randint( - low=1, high=self.max_n_obs_steps+1, - size=(B,), generator=rng, device=device) - - steps = torch.arange(0, T, device=device).reshape(1,T).expand(B,T) - obs_mask = (steps.T < obs_steps).T.reshape(B,T,1).expand(B,T,D) - obs_mask = obs_mask & is_obs_dim - - # generate action mask - if self.action_visible: - action_steps = torch.maximum( - obs_steps - 1, - torch.tensor(0, - dtype=obs_steps.dtype, - device=obs_steps.device)) - action_mask = (steps.T < action_steps).T.reshape(B,T,1).expand(B,T,D) - action_mask = action_mask & is_action_dim - - # generate keypoint mask - if self.time_independent: - visible_kps = torch.rand(size=(B, T, n_keypoints), - generator=rng, device=device) < self.keypoint_visible_rate - visible_dims = torch.repeat_interleave(visible_kps, repeats=self.keypoint_dim, dim=-1) - visible_dims_mask = torch.cat([ - torch.ones((B, T, self.action_dim), - dtype=torch.bool, device=device), - visible_dims, - torch.ones((B, T, self.context_dim), - dtype=torch.bool, device=device), - ], axis=-1) - keypoint_mask = visible_dims_mask - else: - visible_kps = torch.rand(size=(B,n_keypoints), - generator=rng, device=device) < self.keypoint_visible_rate - visible_dims = torch.repeat_interleave(visible_kps, repeats=self.keypoint_dim, dim=-1) - visible_dims_mask = torch.cat([ - torch.ones((B, self.action_dim), - dtype=torch.bool, device=device), - visible_dims, - torch.ones((B, self.context_dim), - dtype=torch.bool, device=device), - ], axis=-1) - keypoint_mask = visible_dims_mask.reshape(B,1,D).expand(B,T,D) - keypoint_mask = keypoint_mask & is_obs_dim - - # generate context mask - context_mask = is_context_dim.clone() - context_mask[:,self.n_context_steps:,:] = False - - mask = obs_mask & keypoint_mask - if self.action_visible: - mask = mask | action_mask - if self.context_dim > 0: - mask = mask | context_mask - - return mask - - -def test(): - # kmg = KeypointMaskGenerator(2,2, random_obs_steps=True) - # self = KeypointMaskGenerator(2,2,context_dim=2, action_visible=True) - # self = KeypointMaskGenerator(2,2,context_dim=0, action_visible=True) - self = LowdimMaskGenerator(2,20, max_n_obs_steps=3, action_visible=True) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/positional_embedding.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/positional_embedding.py deleted file mode 100644 index fe7fdcd3f..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/positional_embedding.py +++ /dev/null @@ -1,17 +0,0 @@ -import math -import torch -import torch.nn as nn - -class SinusoidalPosEmb(nn.Module): - def __init__(self, dim): - super().__init__() - self.dim = dim - - def forward(self, x): - device = x.device - half_dim = self.dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, device=device) * -emb) - emb = x[:, None] * emb[None, :] - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/transformer_for_diffusion.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/transformer_for_diffusion.py deleted file mode 100644 index 2c533c3d3..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/diffusion/transformer_for_diffusion.py +++ /dev/null @@ -1,418 +0,0 @@ -from typing import Union, Optional, Tuple -import logging -import torch -import torch.nn as nn -from diffusion_policy.model.diffusion.positional_embedding import SinusoidalPosEmb -from diffusion_policy.model.common.module_attr_mixin import ModuleAttrMixin - -logger = logging.getLogger(__name__) - -class TransformerForDiffusion(ModuleAttrMixin): - def __init__(self, - input_dim: int, - output_dim: int, - horizon: int, - n_obs_steps: int = None, - cond_dim: int = 0, - n_layer: int = 12, - n_head: int = 12, - n_emb: int = 768, - p_drop_emb: float = 0.1, - p_drop_attn: float = 0.1, - causal_attn: bool=False, - time_as_cond: bool=True, - obs_as_cond: bool=False, - n_cond_layers: int = 0 - ) -> None: - super().__init__() - - # compute number of tokens for main trunk and condition encoder - if n_obs_steps is None: - n_obs_steps = horizon - - T = horizon - T_cond = 1 - if not time_as_cond: - T += 1 - T_cond -= 1 - obs_as_cond = cond_dim > 0 - if obs_as_cond: - assert time_as_cond - T_cond += n_obs_steps - - # input embedding stem - self.input_emb = nn.Linear(input_dim, n_emb) - self.pos_emb = nn.Parameter(torch.zeros(1, T, n_emb)) - self.drop = nn.Dropout(p_drop_emb) - - # cond encoder - self.time_emb = SinusoidalPosEmb(n_emb) - self.cond_obs_emb = None - - if obs_as_cond: - self.cond_obs_emb = nn.Linear(cond_dim, n_emb) - - self.cond_pos_emb = None - self.encoder = None - self.decoder = None - encoder_only = False - if T_cond > 0: - self.cond_pos_emb = nn.Parameter(torch.zeros(1, T_cond, n_emb)) - if n_cond_layers > 0: - encoder_layer = nn.TransformerEncoderLayer( - d_model=n_emb, - nhead=n_head, - dim_feedforward=4*n_emb, - dropout=p_drop_attn, - activation='gelu', - batch_first=True, - norm_first=True - ) - self.encoder = nn.TransformerEncoder( - encoder_layer=encoder_layer, - num_layers=n_cond_layers - ) - else: - self.encoder = nn.Sequential( - nn.Linear(n_emb, 4 * n_emb), - nn.Mish(), - nn.Linear(4 * n_emb, n_emb) - ) - # decoder - decoder_layer = nn.TransformerDecoderLayer( - d_model=n_emb, - nhead=n_head, - dim_feedforward=4*n_emb, - dropout=p_drop_attn, - activation='gelu', - batch_first=True, - norm_first=True # important for stability - ) - self.decoder = nn.TransformerDecoder( - decoder_layer=decoder_layer, - num_layers=n_layer - ) - else: - # encoder only BERT - encoder_only = True - - encoder_layer = nn.TransformerEncoderLayer( - d_model=n_emb, - nhead=n_head, - dim_feedforward=4*n_emb, - dropout=p_drop_attn, - activation='gelu', - batch_first=True, - norm_first=True - ) - self.encoder = nn.TransformerEncoder( - encoder_layer=encoder_layer, - num_layers=n_layer - ) - - # attention mask - if causal_attn: - # causal mask to ensure that attention is only applied to the left in the input sequence - # torch.nn.Transformer uses additive mask as opposed to multiplicative mask in minGPT - # therefore, the upper triangle should be -inf and others (including diag) should be 0. - sz = T - mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1) - mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) - self.register_buffer("mask", mask) - - if time_as_cond and obs_as_cond: - S = T_cond - t, s = torch.meshgrid( - torch.arange(T), - torch.arange(S), - indexing='ij' - ) - mask = t >= (s-1) # add one dimension since time is the first token in cond - mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) - self.register_buffer('memory_mask', mask) - else: - self.memory_mask = None - else: - self.mask = None - self.memory_mask = None - - # decoder head - self.ln_f = nn.LayerNorm(n_emb) - self.head = nn.Linear(n_emb, output_dim) - - # constants - self.T = T - self.T_cond = T_cond - self.horizon = horizon - self.time_as_cond = time_as_cond - self.obs_as_cond = obs_as_cond - self.encoder_only = encoder_only - - # init - self.apply(self._init_weights) - logger.info( - "number of parameters: %e", sum(p.numel() for p in self.parameters()) - ) - - def _init_weights(self, module): - ignore_types = (nn.Dropout, - SinusoidalPosEmb, - nn.TransformerEncoderLayer, - nn.TransformerDecoderLayer, - nn.TransformerEncoder, - nn.TransformerDecoder, - nn.ModuleList, - nn.Mish, - nn.Sequential) - if isinstance(module, (nn.Linear, nn.Embedding)): - torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) - if isinstance(module, nn.Linear) and module.bias is not None: - torch.nn.init.zeros_(module.bias) - elif isinstance(module, nn.MultiheadAttention): - weight_names = [ - 'in_proj_weight', 'q_proj_weight', 'k_proj_weight', 'v_proj_weight'] - for name in weight_names: - weight = getattr(module, name) - if weight is not None: - torch.nn.init.normal_(weight, mean=0.0, std=0.02) - - bias_names = ['in_proj_bias', 'bias_k', 'bias_v'] - for name in bias_names: - bias = getattr(module, name) - if bias is not None: - torch.nn.init.zeros_(bias) - elif isinstance(module, nn.LayerNorm): - torch.nn.init.zeros_(module.bias) - torch.nn.init.ones_(module.weight) - elif isinstance(module, TransformerForDiffusion): - torch.nn.init.normal_(module.pos_emb, mean=0.0, std=0.02) - if module.cond_obs_emb is not None: - torch.nn.init.normal_(module.cond_pos_emb, mean=0.0, std=0.02) - elif isinstance(module, ignore_types): - # no param - pass - else: - raise RuntimeError("Unaccounted module {}".format(module)) - - def get_optim_groups(self, weight_decay: float=1e-3): - """ - This long function is unfortunately doing something very simple and is being very defensive: - We are separating out all parameters of the model into two buckets: those that will experience - weight decay for regularization and those that won't (biases, and layernorm/embedding weights). - We are then returning the PyTorch optimizer object. - """ - - # separate out all parameters to those that will and won't experience regularizing weight decay - decay = set() - no_decay = set() - whitelist_weight_modules = (torch.nn.Linear, torch.nn.MultiheadAttention) - blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding) - for mn, m in self.named_modules(): - for pn, p in m.named_parameters(): - fpn = "%s.%s" % (mn, pn) if mn else pn # full param name - - if pn.endswith("bias"): - # all biases will not be decayed - no_decay.add(fpn) - elif pn.startswith("bias"): - # MultiheadAttention bias starts with "bias" - no_decay.add(fpn) - elif pn.endswith("weight") and isinstance(m, whitelist_weight_modules): - # weights of whitelist modules will be weight decayed - decay.add(fpn) - elif pn.endswith("weight") and isinstance(m, blacklist_weight_modules): - # weights of blacklist modules will NOT be weight decayed - no_decay.add(fpn) - - # special case the position embedding parameter in the root GPT module as not decayed - no_decay.add("pos_emb") - no_decay.add("_dummy_variable") - if self.cond_pos_emb is not None: - no_decay.add("cond_pos_emb") - - # validate that we considered every parameter - param_dict = {pn: p for pn, p in self.named_parameters()} - inter_params = decay & no_decay - union_params = decay | no_decay - assert ( - len(inter_params) == 0 - ), "parameters %s made it into both decay/no_decay sets!" % (str(inter_params),) - assert ( - len(param_dict.keys() - union_params) == 0 - ), "parameters %s were not separated into either decay/no_decay set!" % ( - str(param_dict.keys() - union_params), - ) - - # create the pytorch optimizer object - optim_groups = [ - { - "params": [param_dict[pn] for pn in sorted(list(decay))], - "weight_decay": weight_decay, - }, - { - "params": [param_dict[pn] for pn in sorted(list(no_decay))], - "weight_decay": 0.0, - }, - ] - return optim_groups - - - def configure_optimizers(self, - learning_rate: float=1e-4, - weight_decay: float=1e-3, - betas: Tuple[float, float]=(0.9,0.95)): - optim_groups = self.get_optim_groups(weight_decay=weight_decay) - optimizer = torch.optim.AdamW( - optim_groups, lr=learning_rate, betas=betas - ) - return optimizer - - def forward(self, - sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], - cond: Optional[torch.Tensor]=None, **kwargs): - """ - x: (B,T,input_dim) - timestep: (B,) or int, diffusion step - cond: (B,T',cond_dim) - output: (B,T,input_dim) - """ - # 1. time - timesteps = timestep - if not torch.is_tensor(timesteps): - # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can - timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device) - elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - time_emb = self.time_emb(timesteps).unsqueeze(1) - # (B,1,n_emb) - - # process input - input_emb = self.input_emb(sample) - - if self.encoder_only: - # BERT - token_embeddings = torch.cat([time_emb, input_emb], dim=1) - t = token_embeddings.shape[1] - position_embeddings = self.pos_emb[ - :, :t, : - ] # each position maps to a (learnable) vector - x = self.drop(token_embeddings + position_embeddings) - # (B,T+1,n_emb) - x = self.encoder(src=x, mask=self.mask) - # (B,T+1,n_emb) - x = x[:,1:,:] - # (B,T,n_emb) - else: - # encoder - cond_embeddings = time_emb - if self.obs_as_cond: - cond_obs_emb = self.cond_obs_emb(cond) - # (B,To,n_emb) - cond_embeddings = torch.cat([cond_embeddings, cond_obs_emb], dim=1) - tc = cond_embeddings.shape[1] - position_embeddings = self.cond_pos_emb[ - :, :tc, : - ] # each position maps to a (learnable) vector - x = self.drop(cond_embeddings + position_embeddings) - x = self.encoder(x) - memory = x - # (B,T_cond,n_emb) - - # decoder - token_embeddings = input_emb - t = token_embeddings.shape[1] - position_embeddings = self.pos_emb[ - :, :t, : - ] # each position maps to a (learnable) vector - x = self.drop(token_embeddings + position_embeddings) - # (B,T,n_emb) - x = self.decoder( - tgt=x, - memory=memory, - tgt_mask=self.mask, - memory_mask=self.memory_mask - ) - # (B,T,n_emb) - - # head - x = self.ln_f(x) - x = self.head(x) - # (B,T,n_out) - return x - - -def test(): - # GPT with time embedding - transformer = TransformerForDiffusion( - input_dim=16, - output_dim=16, - horizon=8, - n_obs_steps=4, - # cond_dim=10, - causal_attn=True, - # time_as_cond=False, - # n_cond_layers=4 - ) - opt = transformer.configure_optimizers() - - timestep = torch.tensor(0) - sample = torch.zeros((4,8,16)) - out = transformer(sample, timestep) - - - # GPT with time embedding and obs cond - transformer = TransformerForDiffusion( - input_dim=16, - output_dim=16, - horizon=8, - n_obs_steps=4, - cond_dim=10, - causal_attn=True, - # time_as_cond=False, - # n_cond_layers=4 - ) - opt = transformer.configure_optimizers() - - timestep = torch.tensor(0) - sample = torch.zeros((4,8,16)) - cond = torch.zeros((4,4,10)) - out = transformer(sample, timestep, cond) - - # GPT with time embedding and obs cond and encoder - transformer = TransformerForDiffusion( - input_dim=16, - output_dim=16, - horizon=8, - n_obs_steps=4, - cond_dim=10, - causal_attn=True, - # time_as_cond=False, - n_cond_layers=4 - ) - opt = transformer.configure_optimizers() - - timestep = torch.tensor(0) - sample = torch.zeros((4,8,16)) - cond = torch.zeros((4,4,10)) - out = transformer(sample, timestep, cond) - - # BERT with time embedding token - transformer = TransformerForDiffusion( - input_dim=16, - output_dim=16, - horizon=8, - n_obs_steps=4, - # cond_dim=10, - # causal_attn=True, - time_as_cond=False, - # n_cond_layers=4 - ) - opt = transformer.configure_optimizers() - - timestep = torch.tensor(0) - sample = torch.zeros((4,8,16)) - out = transformer(sample, timestep) - diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/vision/crop_randomizer.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/vision/crop_randomizer.py deleted file mode 100644 index 907957413..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/vision/crop_randomizer.py +++ /dev/null @@ -1,288 +0,0 @@ -import torch -import torch.nn as nn -import torchvision.transforms.functional as ttf -import diffusion_policy.model.common.tensor_util as tu - -class CropRandomizer(nn.Module): - """ - Randomly sample crops at input, and then average across crop features at output. - """ - def __init__( - self, - input_shape, - crop_height, - crop_width, - num_crops=1, - pos_enc=False, - ): - """ - Args: - input_shape (tuple, list): shape of input (not including batch dimension) - crop_height (int): crop height - crop_width (int): crop width - num_crops (int): number of random crops to take - pos_enc (bool): if True, add 2 channels to the output to encode the spatial - location of the cropped pixels in the source image - """ - super().__init__() - - assert len(input_shape) == 3 # (C, H, W) - assert crop_height < input_shape[1] - assert crop_width < input_shape[2] - - self.input_shape = input_shape - self.crop_height = crop_height - self.crop_width = crop_width - self.num_crops = num_crops - self.pos_enc = pos_enc - - def output_shape_in(self, input_shape=None): - """ - Function to compute output shape from inputs to this module. Corresponds to - the @forward_in operation, where raw inputs (usually observation modalities) - are passed in. - - Args: - input_shape (iterable of int): shape of input. Does not include batch dimension. - Some modules may not need this argument, if their output does not depend - on the size of the input, or if they assume fixed size input. - - Returns: - out_shape ([int]): list of integers corresponding to output shape - """ - - # outputs are shape (C, CH, CW), or maybe C + 2 if using position encoding, because - # the number of crops are reshaped into the batch dimension, increasing the batch - # size from B to B * N - out_c = self.input_shape[0] + 2 if self.pos_enc else self.input_shape[0] - return [out_c, self.crop_height, self.crop_width] - - def output_shape_out(self, input_shape=None): - """ - Function to compute output shape from inputs to this module. Corresponds to - the @forward_out operation, where processed inputs (usually encoded observation - modalities) are passed in. - - Args: - input_shape (iterable of int): shape of input. Does not include batch dimension. - Some modules may not need this argument, if their output does not depend - on the size of the input, or if they assume fixed size input. - - Returns: - out_shape ([int]): list of integers corresponding to output shape - """ - - # since the forward_out operation splits [B * N, ...] -> [B, N, ...] - # and then pools to result in [B, ...], only the batch dimension changes, - # and so the other dimensions retain their shape. - return list(input_shape) - - def forward_in(self, inputs): - """ - Samples N random crops for each input in the batch, and then reshapes - inputs to [B * N, ...]. - """ - assert len(inputs.shape) >= 3 # must have at least (C, H, W) dimensions - if self.training: - # generate random crops - out, _ = sample_random_image_crops( - images=inputs, - crop_height=self.crop_height, - crop_width=self.crop_width, - num_crops=self.num_crops, - pos_enc=self.pos_enc, - ) - # [B, N, ...] -> [B * N, ...] - return tu.join_dimensions(out, 0, 1) - else: - # take center crop during eval - out = ttf.center_crop(img=inputs, output_size=( - self.crop_height, self.crop_width)) - if self.num_crops > 1: - B,C,H,W = out.shape - out = out.unsqueeze(1).expand(B,self.num_crops,C,H,W).reshape(-1,C,H,W) - # [B * N, ...] - return out - - def forward_out(self, inputs): - """ - Splits the outputs from shape [B * N, ...] -> [B, N, ...] and then average across N - to result in shape [B, ...] to make sure the network output is consistent with - what would have happened if there were no randomization. - """ - if self.num_crops <= 1: - return inputs - else: - batch_size = (inputs.shape[0] // self.num_crops) - out = tu.reshape_dimensions(inputs, begin_axis=0, end_axis=0, - target_dims=(batch_size, self.num_crops)) - return out.mean(dim=1) - - def forward(self, inputs): - return self.forward_in(inputs) - - def __repr__(self): - """Pretty print network.""" - header = '{}'.format(str(self.__class__.__name__)) - msg = header + "(input_shape={}, crop_size=[{}, {}], num_crops={})".format( - self.input_shape, self.crop_height, self.crop_width, self.num_crops) - return msg - - -def crop_image_from_indices(images, crop_indices, crop_height, crop_width): - """ - Crops images at the locations specified by @crop_indices. Crops will be - taken across all channels. - - Args: - images (torch.Tensor): batch of images of shape [..., C, H, W] - - crop_indices (torch.Tensor): batch of indices of shape [..., N, 2] where - N is the number of crops to take per image and each entry corresponds - to the pixel height and width of where to take the crop. Note that - the indices can also be of shape [..., 2] if only 1 crop should - be taken per image. Leading dimensions must be consistent with - @images argument. Each index specifies the top left of the crop. - Values must be in range [0, H - CH - 1] x [0, W - CW - 1] where - H and W are the height and width of @images and CH and CW are - @crop_height and @crop_width. - - crop_height (int): height of crop to take - - crop_width (int): width of crop to take - - Returns: - crops (torch.Tesnor): cropped images of shape [..., C, @crop_height, @crop_width] - """ - - # make sure length of input shapes is consistent - assert crop_indices.shape[-1] == 2 - ndim_im_shape = len(images.shape) - ndim_indices_shape = len(crop_indices.shape) - assert (ndim_im_shape == ndim_indices_shape + 1) or (ndim_im_shape == ndim_indices_shape + 2) - - # maybe pad so that @crop_indices is shape [..., N, 2] - is_padded = False - if ndim_im_shape == ndim_indices_shape + 2: - crop_indices = crop_indices.unsqueeze(-2) - is_padded = True - - # make sure leading dimensions between images and indices are consistent - assert images.shape[:-3] == crop_indices.shape[:-2] - - device = images.device - image_c, image_h, image_w = images.shape[-3:] - num_crops = crop_indices.shape[-2] - - # make sure @crop_indices are in valid range - assert (crop_indices[..., 0] >= 0).all().item() - assert (crop_indices[..., 0] < (image_h - crop_height)).all().item() - assert (crop_indices[..., 1] >= 0).all().item() - assert (crop_indices[..., 1] < (image_w - crop_width)).all().item() - - # convert each crop index (ch, cw) into a list of pixel indices that correspond to the entire window. - - # 2D index array with columns [0, 1, ..., CH - 1] and shape [CH, CW] - crop_ind_grid_h = torch.arange(crop_height).to(device) - crop_ind_grid_h = tu.unsqueeze_expand_at(crop_ind_grid_h, size=crop_width, dim=-1) - # 2D index array with rows [0, 1, ..., CW - 1] and shape [CH, CW] - crop_ind_grid_w = torch.arange(crop_width).to(device) - crop_ind_grid_w = tu.unsqueeze_expand_at(crop_ind_grid_w, size=crop_height, dim=0) - # combine into shape [CH, CW, 2] - crop_in_grid = torch.cat((crop_ind_grid_h.unsqueeze(-1), crop_ind_grid_w.unsqueeze(-1)), dim=-1) - - # Add above grid with the offset index of each sampled crop to get 2d indices for each crop. - # After broadcasting, this will be shape [..., N, CH, CW, 2] and each crop has a [CH, CW, 2] - # shape array that tells us which pixels from the corresponding source image to grab. - grid_reshape = [1] * len(crop_indices.shape[:-1]) + [crop_height, crop_width, 2] - all_crop_inds = crop_indices.unsqueeze(-2).unsqueeze(-2) + crop_in_grid.reshape(grid_reshape) - - # For using @torch.gather, convert to flat indices from 2D indices, and also - # repeat across the channel dimension. To get flat index of each pixel to grab for - # each sampled crop, we just use the mapping: ind = h_ind * @image_w + w_ind - all_crop_inds = all_crop_inds[..., 0] * image_w + all_crop_inds[..., 1] # shape [..., N, CH, CW] - all_crop_inds = tu.unsqueeze_expand_at(all_crop_inds, size=image_c, dim=-3) # shape [..., N, C, CH, CW] - all_crop_inds = tu.flatten(all_crop_inds, begin_axis=-2) # shape [..., N, C, CH * CW] - - # Repeat and flatten the source images -> [..., N, C, H * W] and then use gather to index with crop pixel inds - images_to_crop = tu.unsqueeze_expand_at(images, size=num_crops, dim=-4) - images_to_crop = tu.flatten(images_to_crop, begin_axis=-2) - crops = torch.gather(images_to_crop, dim=-1, index=all_crop_inds) - # [..., N, C, CH * CW] -> [..., N, C, CH, CW] - reshape_axis = len(crops.shape) - 1 - crops = tu.reshape_dimensions(crops, begin_axis=reshape_axis, end_axis=reshape_axis, - target_dims=(crop_height, crop_width)) - - if is_padded: - # undo padding -> [..., C, CH, CW] - crops = crops.squeeze(-4) - return crops - -def sample_random_image_crops(images, crop_height, crop_width, num_crops, pos_enc=False): - """ - For each image, randomly sample @num_crops crops of size (@crop_height, @crop_width), from - @images. - - Args: - images (torch.Tensor): batch of images of shape [..., C, H, W] - - crop_height (int): height of crop to take - - crop_width (int): width of crop to take - - num_crops (n): number of crops to sample - - pos_enc (bool): if True, also add 2 channels to the outputs that gives a spatial - encoding of the original source pixel locations. This means that the - output crops will contain information about where in the source image - it was sampled from. - - Returns: - crops (torch.Tensor): crops of shape (..., @num_crops, C, @crop_height, @crop_width) - if @pos_enc is False, otherwise (..., @num_crops, C + 2, @crop_height, @crop_width) - - crop_inds (torch.Tensor): sampled crop indices of shape (..., N, 2) - """ - device = images.device - - # maybe add 2 channels of spatial encoding to the source image - source_im = images - if pos_enc: - # spatial encoding [y, x] in [0, 1] - h, w = source_im.shape[-2:] - pos_y, pos_x = torch.meshgrid(torch.arange(h), torch.arange(w)) - pos_y = pos_y.float().to(device) / float(h) - pos_x = pos_x.float().to(device) / float(w) - position_enc = torch.stack((pos_y, pos_x)) # shape [C, H, W] - - # unsqueeze and expand to match leading dimensions -> shape [..., C, H, W] - leading_shape = source_im.shape[:-3] - position_enc = position_enc[(None,) * len(leading_shape)] - position_enc = position_enc.expand(*leading_shape, -1, -1, -1) - - # concat across channel dimension with input - source_im = torch.cat((source_im, position_enc), dim=-3) - - # make sure sample boundaries ensure crops are fully within the images - image_c, image_h, image_w = source_im.shape[-3:] - max_sample_h = image_h - crop_height - max_sample_w = image_w - crop_width - - # Sample crop locations for all tensor dimensions up to the last 3, which are [C, H, W]. - # Each gets @num_crops samples - typically this will just be the batch dimension (B), so - # we will sample [B, N] indices, but this supports having more than one leading dimension, - # or possibly no leading dimension. - # - # Trick: sample in [0, 1) with rand, then re-scale to [0, M) and convert to long to get sampled ints - crop_inds_h = (max_sample_h * torch.rand(*source_im.shape[:-3], num_crops).to(device)).long() - crop_inds_w = (max_sample_w * torch.rand(*source_im.shape[:-3], num_crops).to(device)).long() - crop_inds = torch.cat((crop_inds_h.unsqueeze(-1), crop_inds_w.unsqueeze(-1)), dim=-1) # shape [..., N, 2] - - crops = crop_image_from_indices( - images=source_im, - crop_indices=crop_inds, - crop_height=crop_height, - crop_width=crop_width, - ) - - return crops, crop_inds diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/vision/model_getter.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/vision/model_getter.py deleted file mode 100644 index 6cacd7349..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/vision/model_getter.py +++ /dev/null @@ -1,28 +0,0 @@ -import torch -import torchvision - -def get_resnet(name, weights=None, **kwargs): - """ - name: resnet18, resnet34, resnet50 - weights: "IMAGENET1K_V1", "r3m" - """ - # load r3m weights - if (weights == "r3m") or (weights == "R3M"): - return get_r3m(name=name, **kwargs) - - func = getattr(torchvision.models, name) - resnet = func(weights=weights, **kwargs) - resnet.fc = torch.nn.Identity() - return resnet - -def get_r3m(name, **kwargs): - """ - name: resnet18, resnet34, resnet50 - """ - import r3m - r3m.device = 'cpu' - model = r3m.load_r3m(name) - r3m_model = model.module - resnet_model = r3m_model.convnet - resnet_model = resnet_model.to('cpu') - return resnet_model diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/vision/multi_image_obs_encoder.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/vision/multi_image_obs_encoder.py deleted file mode 100644 index de6aa6589..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/model/vision/multi_image_obs_encoder.py +++ /dev/null @@ -1,195 +0,0 @@ -from typing import Dict, Tuple, Union -import copy -import torch -import torch.nn as nn -import torchvision -from diffusion_policy.model.vision.crop_randomizer import CropRandomizer -from diffusion_policy.model.common.module_attr_mixin import ModuleAttrMixin -from diffusion_policy.common.pytorch_util import dict_apply, replace_submodules - - -class MultiImageObsEncoder(ModuleAttrMixin): - def __init__(self, - shape_meta: dict, - rgb_model: Union[nn.Module, Dict[str,nn.Module]], - resize_shape: Union[Tuple[int,int], Dict[str,tuple], None]=None, - crop_shape: Union[Tuple[int,int], Dict[str,tuple], None]=None, - random_crop: bool=True, - # replace BatchNorm with GroupNorm - use_group_norm: bool=False, - # use single rgb model for all rgb inputs - share_rgb_model: bool=False, - # renormalize rgb input with imagenet normalization - # assuming input in [0,1] - imagenet_norm: bool=False - ): - """ - Assumes rgb input: B,C,H,W - Assumes low_dim input: B,D - """ - super().__init__() - - rgb_keys = list() - low_dim_keys = list() - key_model_map = nn.ModuleDict() - key_transform_map = nn.ModuleDict() - key_shape_map = dict() - - # handle sharing vision backbone - if share_rgb_model: - assert isinstance(rgb_model, nn.Module) - key_model_map['rgb'] = rgb_model - - obs_shape_meta = shape_meta['obs'] - for key, attr in obs_shape_meta.items(): - shape = tuple(attr['shape']) - type = attr.get('type', 'low_dim') - key_shape_map[key] = shape - if type == 'rgb': - rgb_keys.append(key) - # configure model for this key - this_model = None - if not share_rgb_model: - if isinstance(rgb_model, dict): - # have provided model for each key - this_model = rgb_model[key] - else: - assert isinstance(rgb_model, nn.Module) - # have a copy of the rgb model - this_model = copy.deepcopy(rgb_model) - - if this_model is not None: - if use_group_norm: - this_model = replace_submodules( - root_module=this_model, - predicate=lambda x: isinstance(x, nn.BatchNorm2d), - func=lambda x: nn.GroupNorm( - num_groups=x.num_features//16, - num_channels=x.num_features) - ) - key_model_map[key] = this_model - - # configure resize - input_shape = shape - this_resizer = nn.Identity() - if resize_shape is not None: - if isinstance(resize_shape, dict): - h, w = resize_shape[key] - else: - h, w = resize_shape - this_resizer = torchvision.transforms.Resize( - size=(h,w) - ) - input_shape = (shape[0],h,w) - - # configure randomizer - this_randomizer = nn.Identity() - if crop_shape is not None: - if isinstance(crop_shape, dict): - h, w = crop_shape[key] - else: - h, w = crop_shape - if random_crop: - this_randomizer = CropRandomizer( - input_shape=input_shape, - crop_height=h, - crop_width=w, - num_crops=1, - pos_enc=False - ) - else: - this_normalizer = torchvision.transforms.CenterCrop( - size=(h,w) - ) - # configure normalizer - this_normalizer = nn.Identity() - if imagenet_norm: - this_normalizer = torchvision.transforms.Normalize( - mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) - - this_transform = nn.Sequential(this_resizer, this_randomizer, this_normalizer) - key_transform_map[key] = this_transform - elif type == 'low_dim': - low_dim_keys.append(key) - else: - raise RuntimeError(f"Unsupported obs type: {type}") - rgb_keys = sorted(rgb_keys) - low_dim_keys = sorted(low_dim_keys) - - self.shape_meta = shape_meta - self.key_model_map = key_model_map - self.key_transform_map = key_transform_map - self.share_rgb_model = share_rgb_model - self.rgb_keys = rgb_keys - self.low_dim_keys = low_dim_keys - self.key_shape_map = key_shape_map - - def forward(self, obs_dict): - batch_size = None - features = list() - # process rgb input - if self.share_rgb_model: - # pass all rgb obs to rgb model - imgs = list() - for key in self.rgb_keys: - img = obs_dict[key] - if batch_size is None: - batch_size = img.shape[0] - else: - assert batch_size == img.shape[0] - assert img.shape[1:] == self.key_shape_map[key] - img = self.key_transform_map[key](img) - imgs.append(img) - # (N*B,C,H,W) - imgs = torch.cat(imgs, dim=0) - # (N*B,D) - feature = self.key_model_map['rgb'](imgs) - # (N,B,D) - feature = feature.reshape(-1,batch_size,*feature.shape[1:]) - # (B,N,D) - feature = torch.moveaxis(feature,0,1) - # (B,N*D) - feature = feature.reshape(batch_size,-1) - features.append(feature) - else: - # run each rgb obs to independent models - for key in self.rgb_keys: - img = obs_dict[key] - if batch_size is None: - batch_size = img.shape[0] - else: - assert batch_size == img.shape[0] - assert img.shape[1:] == self.key_shape_map[key] - img = self.key_transform_map[key](img) - feature = self.key_model_map[key](img) - features.append(feature) - - # process lowdim input - for key in self.low_dim_keys: - data = obs_dict[key] - if batch_size is None: - batch_size = data.shape[0] - else: - assert batch_size == data.shape[0] - assert data.shape[1:] == self.key_shape_map[key] - features.append(data) - - # concatenate all features - result = torch.cat(features, dim=-1) - return result - - @torch.no_grad() - def output_shape(self): - example_obs_dict = dict() - obs_shape_meta = self.shape_meta['obs'] - batch_size = 1 - for key, attr in obs_shape_meta.items(): - shape = tuple(attr['shape']) - this_obs = torch.zeros( - (batch_size,) + shape, - dtype=self.dtype, - device=self.device) - example_obs_dict[key] = this_obs - example_output = self.forward(example_obs_dict) - output_shape = example_output.shape[1:] - return output_shape diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/base_image_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/base_image_policy.py deleted file mode 100644 index 383e18082..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/base_image_policy.py +++ /dev/null @@ -1,25 +0,0 @@ -from typing import Dict -import torch -import torch.nn as nn -from diffusion_policy.model.common.module_attr_mixin import ModuleAttrMixin -from diffusion_policy.model.common.normalizer import LinearNormalizer - -class BaseImagePolicy(ModuleAttrMixin): - # init accepts keyword argument shape_meta, see config/task/*_image.yaml - - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """ - obs_dict: - str: B,To,* - return: B,Ta,Da - """ - raise NotImplementedError() - - # reset state for stateful policies - def reset(self): - pass - - # ========== training =========== - # no standard training interface except setting normalizer - def set_normalizer(self, normalizer: LinearNormalizer): - raise NotImplementedError() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/base_lowdim_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/base_lowdim_policy.py deleted file mode 100644 index 8bbaafc0a..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/base_lowdim_policy.py +++ /dev/null @@ -1,36 +0,0 @@ -from typing import Dict -import torch -import torch.nn as nn -from diffusion_policy.model.common.module_attr_mixin import ModuleAttrMixin -from diffusion_policy.model.common.normalizer import LinearNormalizer - -class BaseLowdimPolicy(ModuleAttrMixin): - # ========= inference ============ - # also as self.device and self.dtype for inference device transfer - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """ - obs_dict: - obs: B,To,Do - return: - action: B,Ta,Da - To = 3 - Ta = 4 - T = 6 - |o|o|o| - | | |a|a|a|a| - |o|o| - | |a|a|a|a|a| - | | | | |a|a| - """ - raise NotImplementedError() - - # reset state for stateful policies - def reset(self): - pass - - # ========== training =========== - # no standard training interface except setting normalizer - def set_normalizer(self, normalizer: LinearNormalizer): - raise NotImplementedError() - - \ No newline at end of file diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/bet_lowdim_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/bet_lowdim_policy.py deleted file mode 100644 index 8464a2eac..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/bet_lowdim_policy.py +++ /dev/null @@ -1,127 +0,0 @@ -from typing import Dict, Tuple -import torch -import torch.nn as nn -from omegaconf import OmegaConf -import torch.nn.functional as F - -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.policy.base_lowdim_policy import BaseLowdimPolicy -from diffusion_policy.model.bet.action_ae.discretizers.k_means import KMeansDiscretizer -from diffusion_policy.model.bet.latent_generators.mingpt import MinGPT -from diffusion_policy.model.bet.utils import eval_mode - -class BETLowdimPolicy(BaseLowdimPolicy): - def __init__(self, - action_ae: KMeansDiscretizer, - obs_encoding_net: nn.Module, - state_prior: MinGPT, - horizon, - n_action_steps, - n_obs_steps): - super().__init__() - - self.normalizer = LinearNormalizer() - self.action_ae = action_ae - self.obs_encoding_net = obs_encoding_net - self.state_prior = state_prior - self.horizon = horizon - self.n_action_steps = n_action_steps - self.n_obs_steps = n_obs_steps - - # ========= inference ============ - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """ - obs_dict: must include "obs" key - result: must include "action" key - """ - assert 'obs' in obs_dict - assert 'past_action' not in obs_dict # not implemented yet - nobs = self.normalizer['obs'].normalize(obs_dict['obs']) - B, _, Do = nobs.shape - To = self.n_obs_steps - T = self.horizon - - # pad To to T - obs = torch.full((B,T,Do), -2, dtype=nobs.dtype, device=nobs.device) - obs[:,:To,:] = nobs[:,:To,:] - - # (B,T,Do) - enc_obs = self.obs_encoding_net(obs) - - # Sample latents from the prior - latents, offsets = self.state_prior.generate_latents(enc_obs) - - # un-descritize - naction_pred = self.action_ae.decode_actions( - latent_action_batch=(latents, offsets) - ) - # (B,T,Da) - - # un-normalize - action_pred = self.normalizer['action'].unnormalize(naction_pred) - - # get action - start = To - 1 - end = start + self.n_action_steps - action = action_pred[:,start:end] - result = { - 'action': action, - 'action_pred': action_pred - } - return result - - # ========= training ============ - def set_normalizer(self, normalizer: LinearNormalizer): - self.normalizer.load_state_dict(normalizer.state_dict()) - - def fit_action_ae(self, input_actions: torch.Tensor): - self.action_ae.fit_discretizer(input_actions=input_actions) - - def get_latents(self, latent_collection_loader): - training_latents = list() - with eval_mode(self.action_ae, self.obs_encoding_net, no_grad=True): - for observations, action, mask in latent_collection_loader: - obs, act = observations.to(self.device, non_blocking=True), action.to(self.device, non_blocking=True) - enc_obs = self.obs_encoding_net(obs) - latent = self.action_ae.encode_into_latent(act, enc_obs) - reconstructed_action = self.action_ae.decode_actions( - latent, - enc_obs, - ) - total_mse_loss += F.mse_loss(act, reconstructed_action, reduction="sum") - if type(latent) == tuple: - # serialize into tensor; assumes last dim is latent dim - detached_latents = tuple(x.detach() for x in latent) - training_latents.append(torch.cat(detached_latents, dim=-1)) - else: - training_latents.append(latent.detach()) - training_latents_tensor = torch.cat(training_latents, dim=0) - return training_latents_tensor - - def get_optimizer( - self, weight_decay: float, learning_rate: float, betas: Tuple[float, float] - ) -> torch.optim.Optimizer: - return self.state_prior.get_optimizer( - weight_decay=weight_decay, - learning_rate=learning_rate, - betas=tuple(betas)) - - def compute_loss(self, batch): - # normalize input - assert 'valid_mask' not in batch - nbatch = self.normalizer.normalize(batch) - obs = nbatch['obs'] - action = nbatch['action'] - - # mask out observations after n_obs_steps - To = self.n_obs_steps - obs[:,To:,:] = -2 # (normal obs range [-1,1]) - - enc_obs = self.obs_encoding_net(obs) - latent = self.action_ae.encode_into_latent(action, enc_obs) - _, loss, loss_components = self.state_prior.get_latent_and_loss( - obs_rep=enc_obs, - target_latents=latent, - return_loss_components=True, - ) - return loss, loss_components diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_transformer_hybrid_image_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_transformer_hybrid_image_policy.py deleted file mode 100644 index 114d8447b..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_transformer_hybrid_image_policy.py +++ /dev/null @@ -1,385 +0,0 @@ -from typing import Dict, Tuple -import math -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange, reduce -from diffusers.schedulers.scheduling_ddpm import DDPMScheduler - -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.policy.base_image_policy import BaseImagePolicy -from diffusion_policy.model.diffusion.transformer_for_diffusion import TransformerForDiffusion -from diffusion_policy.model.diffusion.mask_generator import LowdimMaskGenerator -from diffusion_policy.common.robomimic_config_util import get_robomimic_config -from robomimic.algo import algo_factory -from robomimic.algo.algo import PolicyAlgo -import robomimic.utils.obs_utils as ObsUtils -import robomimic.models.base_nets as rmbn -import diffusion_policy.model.vision.crop_randomizer as dmvc -from diffusion_policy.common.pytorch_util import dict_apply, replace_submodules - - -class DiffusionTransformerHybridImagePolicy(BaseImagePolicy): - def __init__(self, - shape_meta: dict, - noise_scheduler: DDPMScheduler, - # task params - horizon, - n_action_steps, - n_obs_steps, - num_inference_steps=None, - # image - crop_shape=(76, 76), - obs_encoder_group_norm=False, - eval_fixed_crop=False, - # arch - n_layer=8, - n_cond_layers=0, - n_head=4, - n_emb=256, - p_drop_emb=0.0, - p_drop_attn=0.3, - causal_attn=True, - time_as_cond=True, - obs_as_cond=True, - pred_action_steps_only=False, - # parameters passed to step - **kwargs): - super().__init__() - - # parse shape_meta - action_shape = shape_meta['action']['shape'] - assert len(action_shape) == 1 - action_dim = action_shape[0] - obs_shape_meta = shape_meta['obs'] - obs_config = { - 'low_dim': [], - 'rgb': [], - 'depth': [], - 'scan': [] - } - obs_key_shapes = dict() - for key, attr in obs_shape_meta.items(): - shape = attr['shape'] - obs_key_shapes[key] = list(shape) - - type = attr.get('type', 'low_dim') - if type == 'rgb': - obs_config['rgb'].append(key) - elif type == 'low_dim': - obs_config['low_dim'].append(key) - else: - raise RuntimeError(f"Unsupported obs type: {type}") - - # get raw robomimic config - config = get_robomimic_config( - algo_name='bc_rnn', - hdf5_type='image', - task_name='square', - dataset_type='ph') - - with config.unlocked(): - # set config with shape_meta - config.observation.modalities.obs = obs_config - - if crop_shape is None: - for key, modality in config.observation.encoder.items(): - if modality.obs_randomizer_class == 'CropRandomizer': - modality['obs_randomizer_class'] = None - else: - # set random crop parameter - ch, cw = crop_shape - for key, modality in config.observation.encoder.items(): - if modality.obs_randomizer_class == 'CropRandomizer': - modality.obs_randomizer_kwargs.crop_height = ch - modality.obs_randomizer_kwargs.crop_width = cw - - # init global state - ObsUtils.initialize_obs_utils_with_config(config) - - # load model - policy: PolicyAlgo = algo_factory( - algo_name=config.algo_name, - config=config, - obs_key_shapes=obs_key_shapes, - ac_dim=action_dim, - device='cpu', - ) - - obs_encoder = policy.nets['policy'].nets['encoder'].nets['obs'] - - if obs_encoder_group_norm: - # replace batch norm with group norm - replace_submodules( - root_module=obs_encoder, - predicate=lambda x: isinstance(x, nn.BatchNorm2d), - func=lambda x: nn.GroupNorm( - num_groups=x.num_features//16, - num_channels=x.num_features) - ) - # obs_encoder.obs_nets['agentview_image'].nets[0].nets - - # obs_encoder.obs_randomizers['agentview_image'] - if eval_fixed_crop: - replace_submodules( - root_module=obs_encoder, - predicate=lambda x: isinstance(x, rmbn.CropRandomizer), - func=lambda x: dmvc.CropRandomizer( - input_shape=x.input_shape, - crop_height=x.crop_height, - crop_width=x.crop_width, - num_crops=x.num_crops, - pos_enc=x.pos_enc - ) - ) - - # create diffusion model - obs_feature_dim = obs_encoder.output_shape()[0] - input_dim = action_dim if obs_as_cond else (obs_feature_dim + action_dim) - output_dim = input_dim - cond_dim = obs_feature_dim if obs_as_cond else 0 - - model = TransformerForDiffusion( - input_dim=input_dim, - output_dim=output_dim, - horizon=horizon, - n_obs_steps=n_obs_steps, - cond_dim=cond_dim, - n_layer=n_layer, - n_head=n_head, - n_emb=n_emb, - p_drop_emb=p_drop_emb, - p_drop_attn=p_drop_attn, - causal_attn=causal_attn, - time_as_cond=time_as_cond, - obs_as_cond=obs_as_cond, - n_cond_layers=n_cond_layers - ) - - self.obs_encoder = obs_encoder - self.model = model - self.noise_scheduler = noise_scheduler - self.mask_generator = LowdimMaskGenerator( - action_dim=action_dim, - obs_dim=0 if (obs_as_cond) else obs_feature_dim, - max_n_obs_steps=n_obs_steps, - fix_obs_steps=True, - action_visible=False - ) - self.normalizer = LinearNormalizer() - self.horizon = horizon - self.obs_feature_dim = obs_feature_dim - self.action_dim = action_dim - self.n_action_steps = n_action_steps - self.n_obs_steps = n_obs_steps - self.obs_as_cond = obs_as_cond - self.pred_action_steps_only = pred_action_steps_only - self.kwargs = kwargs - - if num_inference_steps is None: - num_inference_steps = noise_scheduler.config.num_train_timesteps - self.num_inference_steps = num_inference_steps - - # ========= inference ============ - def conditional_sample(self, - condition_data, condition_mask, - cond=None, generator=None, - # keyword arguments to scheduler.step - **kwargs - ): - model = self.model - scheduler = self.noise_scheduler - - trajectory = torch.randn( - size=condition_data.shape, - dtype=condition_data.dtype, - device=condition_data.device, - generator=generator) - - # set step values - scheduler.set_timesteps(self.num_inference_steps) - - for t in scheduler.timesteps: - # 1. apply conditioning - trajectory[condition_mask] = condition_data[condition_mask] - - # 2. predict model output - model_output = model(trajectory, t, cond) - - # 3. compute previous image: x_t -> x_t-1 - trajectory = scheduler.step( - model_output, t, trajectory, - generator=generator, - **kwargs - ).prev_sample - - # finally make sure conditioning is enforced - trajectory[condition_mask] = condition_data[condition_mask] - - return trajectory - - - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """ - obs_dict: must include "obs" key - result: must include "action" key - """ - assert 'past_action' not in obs_dict # not implemented yet - # normalize input - nobs = self.normalizer.normalize(obs_dict) - value = next(iter(nobs.values())) - B, To = value.shape[:2] - T = self.horizon - Da = self.action_dim - Do = self.obs_feature_dim - To = self.n_obs_steps - - # build input - device = self.device - dtype = self.dtype - - # handle different ways of passing observation - cond = None - cond_data = None - cond_mask = None - if self.obs_as_cond: - this_nobs = dict_apply(nobs, lambda x: x[:,:To,...].reshape(-1,*x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, To, Do - cond = nobs_features.reshape(B, To, -1) - shape = (B, T, Da) - if self.pred_action_steps_only: - shape = (B, self.n_action_steps, Da) - cond_data = torch.zeros(size=shape, device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - else: - # condition through impainting - this_nobs = dict_apply(nobs, lambda x: x[:,:To,...].reshape(-1,*x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, To, Do - nobs_features = nobs_features.reshape(B, To, -1) - shape = (B, T, Da+Do) - cond_data = torch.zeros(size=shape, device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - cond_data[:,:To,Da:] = nobs_features - cond_mask[:,:To,Da:] = True - - # run sampling - nsample = self.conditional_sample( - cond_data, - cond_mask, - cond=cond, - **self.kwargs) - - # unnormalize prediction - naction_pred = nsample[...,:Da] - action_pred = self.normalizer['action'].unnormalize(naction_pred) - - # get action - if self.pred_action_steps_only: - action = action_pred - else: - start = To - 1 - end = start + self.n_action_steps - action = action_pred[:,start:end] - - result = { - 'action': action, - 'action_pred': action_pred - } - return result - - # ========= training ============ - def set_normalizer(self, normalizer: LinearNormalizer): - self.normalizer.load_state_dict(normalizer.state_dict()) - - def get_optimizer( - self, - transformer_weight_decay: float, - obs_encoder_weight_decay: float, - learning_rate: float, - betas: Tuple[float, float] - ) -> torch.optim.Optimizer: - optim_groups = self.model.get_optim_groups( - weight_decay=transformer_weight_decay) - optim_groups.append({ - "params": self.obs_encoder.parameters(), - "weight_decay": obs_encoder_weight_decay - }) - optimizer = torch.optim.AdamW( - optim_groups, lr=learning_rate, betas=betas - ) - return optimizer - - def compute_loss(self, batch): - # normalize input - assert 'valid_mask' not in batch - nobs = self.normalizer.normalize(batch['obs']) - nactions = self.normalizer['action'].normalize(batch['action']) - batch_size = nactions.shape[0] - horizon = nactions.shape[1] - To = self.n_obs_steps - - # handle different ways of passing observation - cond = None - trajectory = nactions - if self.obs_as_cond: - # reshape B, T, ... to B*T - this_nobs = dict_apply(nobs, - lambda x: x[:,:To,...].reshape(-1,*x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, T, Do - cond = nobs_features.reshape(batch_size, To, -1) - if self.pred_action_steps_only: - start = To - 1 - end = start + self.n_action_steps - trajectory = nactions[:,start:end] - else: - # reshape B, T, ... to B*T - this_nobs = dict_apply(nobs, lambda x: x.reshape(-1, *x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, T, Do - nobs_features = nobs_features.reshape(batch_size, horizon, -1) - trajectory = torch.cat([nactions, nobs_features], dim=-1).detach() - - # generate impainting mask - if self.pred_action_steps_only: - condition_mask = torch.zeros_like(trajectory, dtype=torch.bool) - else: - condition_mask = self.mask_generator(trajectory.shape) - - # Sample noise that we'll add to the images - noise = torch.randn(trajectory.shape, device=trajectory.device) - bsz = trajectory.shape[0] - # Sample a random timestep for each image - timesteps = torch.randint( - 0, self.noise_scheduler.config.num_train_timesteps, - (bsz,), device=trajectory.device - ).long() - # Add noise to the clean images according to the noise magnitude at each timestep - # (this is the forward diffusion process) - noisy_trajectory = self.noise_scheduler.add_noise( - trajectory, noise, timesteps) - - # compute loss mask - loss_mask = ~condition_mask - - # apply conditioning - noisy_trajectory[condition_mask] = trajectory[condition_mask] - - # Predict the noise residual - pred = self.model(noisy_trajectory, timesteps, cond) - - pred_type = self.noise_scheduler.config.prediction_type - if pred_type == 'epsilon': - target = noise - elif pred_type == 'sample': - target = trajectory - else: - raise ValueError(f"Unsupported prediction type {pred_type}") - - loss = F.mse_loss(pred, target, reduction='none') - loss = loss * loss_mask.type(loss.dtype) - loss = reduce(loss, 'b ... -> b (...)', 'mean') - loss = loss.mean() - return loss diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_transformer_lowdim_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_transformer_lowdim_policy.py deleted file mode 100644 index 7d0eff8a3..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_transformer_lowdim_policy.py +++ /dev/null @@ -1,234 +0,0 @@ -from typing import Dict, Tuple -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange, reduce -from diffusers.schedulers.scheduling_ddpm import DDPMScheduler - -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.policy.base_lowdim_policy import BaseLowdimPolicy -from diffusion_policy.model.diffusion.transformer_for_diffusion import TransformerForDiffusion -from diffusion_policy.model.diffusion.mask_generator import LowdimMaskGenerator - -class DiffusionTransformerLowdimPolicy(BaseLowdimPolicy): - def __init__(self, - model: TransformerForDiffusion, - noise_scheduler: DDPMScheduler, - horizon, - obs_dim, - action_dim, - n_action_steps, - n_obs_steps, - num_inference_steps=None, - obs_as_cond=False, - pred_action_steps_only=False, - # parameters passed to step - **kwargs): - super().__init__() - if pred_action_steps_only: - assert obs_as_cond - - self.model = model - self.noise_scheduler = noise_scheduler - self.mask_generator = LowdimMaskGenerator( - action_dim=action_dim, - obs_dim=0 if (obs_as_cond) else obs_dim, - max_n_obs_steps=n_obs_steps, - fix_obs_steps=True, - action_visible=False - ) - self.normalizer = LinearNormalizer() - self.horizon = horizon - self.obs_dim = obs_dim - self.action_dim = action_dim - self.n_action_steps = n_action_steps - self.n_obs_steps = n_obs_steps - self.obs_as_cond = obs_as_cond - self.pred_action_steps_only = pred_action_steps_only - self.kwargs = kwargs - - if num_inference_steps is None: - num_inference_steps = noise_scheduler.config.num_train_timesteps - self.num_inference_steps = num_inference_steps - - # ========= inference ============ - def conditional_sample(self, - condition_data, condition_mask, - cond=None, generator=None, - # keyword arguments to scheduler.step - **kwargs - ): - model = self.model - scheduler = self.noise_scheduler - - trajectory = torch.randn( - size=condition_data.shape, - dtype=condition_data.dtype, - device=condition_data.device, - generator=generator) - - # set step values - scheduler.set_timesteps(self.num_inference_steps) - - for t in scheduler.timesteps: - # 1. apply conditioning - trajectory[condition_mask] = condition_data[condition_mask] - - # 2. predict model output - model_output = model(trajectory, t, cond) - - # 3. compute previous image: x_t -> x_t-1 - trajectory = scheduler.step( - model_output, t, trajectory, - generator=generator, - **kwargs - ).prev_sample - - # finally make sure conditioning is enforced - trajectory[condition_mask] = condition_data[condition_mask] - - return trajectory - - - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """ - obs_dict: must include "obs" key - result: must include "action" key - """ - - assert 'obs' in obs_dict - assert 'past_action' not in obs_dict # not implemented yet - nobs = self.normalizer['obs'].normalize(obs_dict['obs']) - B, _, Do = nobs.shape - To = self.n_obs_steps - assert Do == self.obs_dim - T = self.horizon - Da = self.action_dim - - # build input - device = self.device - dtype = self.dtype - - # handle different ways of passing observation - cond = None - cond_data = None - cond_mask = None - if self.obs_as_cond: - cond = nobs[:,:To] - shape = (B, T, Da) - if self.pred_action_steps_only: - shape = (B, self.n_action_steps, Da) - cond_data = torch.zeros(size=shape, device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - else: - # condition through impainting - shape = (B, T, Da+Do) - cond_data = torch.zeros(size=shape, device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - cond_data[:,:To,Da:] = nobs[:,:To] - cond_mask[:,:To,Da:] = True - - # run sampling - nsample = self.conditional_sample( - cond_data, - cond_mask, - cond=cond, - **self.kwargs) - - # unnormalize prediction - naction_pred = nsample[...,:Da] - action_pred = self.normalizer['action'].unnormalize(naction_pred) - - # get action - if self.pred_action_steps_only: - action = action_pred - else: - start = To - 1 - end = start + self.n_action_steps - action = action_pred[:,start:end] - - result = { - 'action': action, - 'action_pred': action_pred - } - if not self.obs_as_cond: - nobs_pred = nsample[...,Da:] - obs_pred = self.normalizer['obs'].unnormalize(nobs_pred) - action_obs_pred = obs_pred[:,start:end] - result['action_obs_pred'] = action_obs_pred - result['obs_pred'] = obs_pred - return result - - # ========= training ============ - def set_normalizer(self, normalizer: LinearNormalizer): - self.normalizer.load_state_dict(normalizer.state_dict()) - - def get_optimizer( - self, weight_decay: float, learning_rate: float, betas: Tuple[float, float] - ) -> torch.optim.Optimizer: - return self.model.configure_optimizers( - weight_decay=weight_decay, - learning_rate=learning_rate, - betas=tuple(betas)) - - def compute_loss(self, batch): - # normalize input - assert 'valid_mask' not in batch - nbatch = self.normalizer.normalize(batch) - obs = nbatch['obs'] - action = nbatch['action'] - - # handle different ways of passing observation - cond = None - trajectory = action - if self.obs_as_cond: - cond = obs[:,:self.n_obs_steps,:] - if self.pred_action_steps_only: - To = self.n_obs_steps - start = To - 1 - end = start + self.n_action_steps - trajectory = action[:,start:end] - else: - trajectory = torch.cat([action, obs], dim=-1) - - # generate impainting mask - if self.pred_action_steps_only: - condition_mask = torch.zeros_like(trajectory, dtype=torch.bool) - else: - condition_mask = self.mask_generator(trajectory.shape) - - # Sample noise that we'll add to the images - noise = torch.randn(trajectory.shape, device=trajectory.device) - bsz = trajectory.shape[0] - # Sample a random timestep for each image - timesteps = torch.randint( - 0, self.noise_scheduler.config.num_train_timesteps, - (bsz,), device=trajectory.device - ).long() - # Add noise to the clean images according to the noise magnitude at each timestep - # (this is the forward diffusion process) - noisy_trajectory = self.noise_scheduler.add_noise( - trajectory, noise, timesteps) - - # compute loss mask - loss_mask = ~condition_mask - - # apply conditioning - noisy_trajectory[condition_mask] = trajectory[condition_mask] - - # Predict the noise residual - pred = self.model(noisy_trajectory, timesteps, cond) - - pred_type = self.noise_scheduler.config.prediction_type - if pred_type == 'epsilon': - target = noise - elif pred_type == 'sample': - target = trajectory - else: - raise ValueError(f"Unsupported prediction type {pred_type}") - - loss = F.mse_loss(pred, target, reduction='none') - loss = loss * loss_mask.type(loss.dtype) - loss = reduce(loss, 'b ... -> b (...)', 'mean') - loss = loss.mean() - return loss diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_hybrid_image_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_hybrid_image_policy.py deleted file mode 100644 index cb2d84898..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_hybrid_image_policy.py +++ /dev/null @@ -1,351 +0,0 @@ -from typing import Dict -import math -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange, reduce -from diffusers.schedulers.scheduling_ddpm import DDPMScheduler - -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.policy.base_image_policy import BaseImagePolicy -from diffusion_policy.model.diffusion.conditional_unet1d import ConditionalUnet1D -from diffusion_policy.model.diffusion.mask_generator import LowdimMaskGenerator -from diffusion_policy.common.robomimic_config_util import get_robomimic_config -from robomimic.algo import algo_factory -from robomimic.algo.algo import PolicyAlgo -import robomimic.utils.obs_utils as ObsUtils -import robomimic.models.base_nets as rmbn -import diffusion_policy.model.vision.crop_randomizer as dmvc -from diffusion_policy.common.pytorch_util import dict_apply, replace_submodules - - -class DiffusionUnetHybridImagePolicy(BaseImagePolicy): - def __init__(self, - shape_meta: dict, - noise_scheduler: DDPMScheduler, - horizon, - n_action_steps, - n_obs_steps, - num_inference_steps=None, - obs_as_global_cond=True, - crop_shape=(76, 76), - diffusion_step_embed_dim=256, - down_dims=(256,512,1024), - kernel_size=5, - n_groups=8, - cond_predict_scale=True, - obs_encoder_group_norm=False, - eval_fixed_crop=False, - # parameters passed to step - **kwargs): - super().__init__() - - # parse shape_meta - action_shape = shape_meta['action']['shape'] - assert len(action_shape) == 1 - action_dim = action_shape[0] - obs_shape_meta = shape_meta['obs'] - obs_config = { - 'low_dim': [], - 'rgb': [], - 'depth': [], - 'scan': [] - } - obs_key_shapes = dict() - for key, attr in obs_shape_meta.items(): - shape = attr['shape'] - obs_key_shapes[key] = list(shape) - - type = attr.get('type', 'low_dim') - if type == 'rgb': - obs_config['rgb'].append(key) - elif type == 'low_dim': - obs_config['low_dim'].append(key) - else: - raise RuntimeError(f"Unsupported obs type: {type}") - - # get raw robomimic config - config = get_robomimic_config( - algo_name='bc_rnn', - hdf5_type='image', - task_name='square', - dataset_type='ph') - - with config.unlocked(): - # set config with shape_meta - config.observation.modalities.obs = obs_config - - if crop_shape is None: - for key, modality in config.observation.encoder.items(): - if modality.obs_randomizer_class == 'CropRandomizer': - modality['obs_randomizer_class'] = None - else: - # set random crop parameter - ch, cw = crop_shape - for key, modality in config.observation.encoder.items(): - if modality.obs_randomizer_class == 'CropRandomizer': - modality.obs_randomizer_kwargs.crop_height = ch - modality.obs_randomizer_kwargs.crop_width = cw - - # init global state - ObsUtils.initialize_obs_utils_with_config(config) - - # load model - policy: PolicyAlgo = algo_factory( - algo_name=config.algo_name, - config=config, - obs_key_shapes=obs_key_shapes, - ac_dim=action_dim, - device='cpu', - ) - - obs_encoder = policy.nets['policy'].nets['encoder'].nets['obs'] - - if obs_encoder_group_norm: - # replace batch norm with group norm - replace_submodules( - root_module=obs_encoder, - predicate=lambda x: isinstance(x, nn.BatchNorm2d), - func=lambda x: nn.GroupNorm( - num_groups=x.num_features//16, - num_channels=x.num_features) - ) - # obs_encoder.obs_nets['agentview_image'].nets[0].nets - - # obs_encoder.obs_randomizers['agentview_image'] - if eval_fixed_crop: - replace_submodules( - root_module=obs_encoder, - predicate=lambda x: isinstance(x, rmbn.CropRandomizer), - func=lambda x: dmvc.CropRandomizer( - input_shape=x.input_shape, - crop_height=x.crop_height, - crop_width=x.crop_width, - num_crops=x.num_crops, - pos_enc=x.pos_enc - ) - ) - - # create diffusion model - obs_feature_dim = obs_encoder.output_shape()[0] - input_dim = action_dim + obs_feature_dim - global_cond_dim = None - if obs_as_global_cond: - input_dim = action_dim - global_cond_dim = obs_feature_dim * n_obs_steps - - model = ConditionalUnet1D( - input_dim=input_dim, - local_cond_dim=None, - global_cond_dim=global_cond_dim, - diffusion_step_embed_dim=diffusion_step_embed_dim, - down_dims=down_dims, - kernel_size=kernel_size, - n_groups=n_groups, - cond_predict_scale=cond_predict_scale - ) - - self.obs_encoder = obs_encoder - self.model = model - self.noise_scheduler = noise_scheduler - self.mask_generator = LowdimMaskGenerator( - action_dim=action_dim, - obs_dim=0 if obs_as_global_cond else obs_feature_dim, - max_n_obs_steps=n_obs_steps, - fix_obs_steps=True, - action_visible=False - ) - self.normalizer = LinearNormalizer() - self.horizon = horizon - self.obs_feature_dim = obs_feature_dim - self.action_dim = action_dim - self.n_action_steps = n_action_steps - self.n_obs_steps = n_obs_steps - self.obs_as_global_cond = obs_as_global_cond - self.kwargs = kwargs - - if num_inference_steps is None: - num_inference_steps = noise_scheduler.config.num_train_timesteps - self.num_inference_steps = num_inference_steps - - print("Diffusion params: %e" % sum(p.numel() for p in self.model.parameters())) - print("Vision params: %e" % sum(p.numel() for p in self.obs_encoder.parameters())) - - # ========= inference ============ - def conditional_sample(self, - condition_data, condition_mask, - local_cond=None, global_cond=None, - generator=None, - # keyword arguments to scheduler.step - **kwargs - ): - model = self.model - scheduler = self.noise_scheduler - - trajectory = torch.randn( - size=condition_data.shape, - dtype=condition_data.dtype, - device=condition_data.device, - generator=generator) - - # set step values - scheduler.set_timesteps(self.num_inference_steps) - - for t in scheduler.timesteps: - # 1. apply conditioning - trajectory[condition_mask] = condition_data[condition_mask] - - # 2. predict model output - model_output = model(trajectory, t, - local_cond=local_cond, global_cond=global_cond) - - # 3. compute previous image: x_t -> x_t-1 - trajectory = scheduler.step( - model_output, t, trajectory, - generator=generator, - **kwargs - ).prev_sample - - # finally make sure conditioning is enforced - trajectory[condition_mask] = condition_data[condition_mask] - - return trajectory - - - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """ - obs_dict: must include "obs" key - result: must include "action" key - """ - assert 'past_action' not in obs_dict # not implemented yet - # normalize input - nobs = self.normalizer.normalize(obs_dict) - value = next(iter(nobs.values())) - B, To = value.shape[:2] - T = self.horizon - Da = self.action_dim - Do = self.obs_feature_dim - To = self.n_obs_steps - - # build input - device = self.device - dtype = self.dtype - - # handle different ways of passing observation - local_cond = None - global_cond = None - if self.obs_as_global_cond: - # condition through global feature - this_nobs = dict_apply(nobs, lambda x: x[:,:To,...].reshape(-1,*x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, Do - global_cond = nobs_features.reshape(B, -1) - # empty data for action - cond_data = torch.zeros(size=(B, T, Da), device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - else: - # condition through impainting - this_nobs = dict_apply(nobs, lambda x: x[:,:To,...].reshape(-1,*x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, To, Do - nobs_features = nobs_features.reshape(B, To, -1) - cond_data = torch.zeros(size=(B, T, Da+Do), device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - cond_data[:,:To,Da:] = nobs_features - cond_mask[:,:To,Da:] = True - - # run sampling - nsample = self.conditional_sample( - cond_data, - cond_mask, - local_cond=local_cond, - global_cond=global_cond, - **self.kwargs) - - # unnormalize prediction - naction_pred = nsample[...,:Da] - action_pred = self.normalizer['action'].unnormalize(naction_pred) - - # get action - start = To - 1 - end = start + self.n_action_steps - action = action_pred[:,start:end] - - result = { - 'action': action, - 'action_pred': action_pred - } - return result - - # ========= training ============ - def set_normalizer(self, normalizer: LinearNormalizer): - self.normalizer.load_state_dict(normalizer.state_dict()) - - def compute_loss(self, batch): - # normalize input - assert 'valid_mask' not in batch - nobs = self.normalizer.normalize(batch['obs']) - nactions = self.normalizer['action'].normalize(batch['action']) - batch_size = nactions.shape[0] - horizon = nactions.shape[1] - - # handle different ways of passing observation - local_cond = None - global_cond = None - trajectory = nactions - cond_data = trajectory - if self.obs_as_global_cond: - # reshape B, T, ... to B*T - this_nobs = dict_apply(nobs, - lambda x: x[:,:self.n_obs_steps,...].reshape(-1,*x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, Do - global_cond = nobs_features.reshape(batch_size, -1) - else: - # reshape B, T, ... to B*T - this_nobs = dict_apply(nobs, lambda x: x.reshape(-1, *x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, T, Do - nobs_features = nobs_features.reshape(batch_size, horizon, -1) - cond_data = torch.cat([nactions, nobs_features], dim=-1) - trajectory = cond_data.detach() - - # generate impainting mask - condition_mask = self.mask_generator(trajectory.shape) - - # Sample noise that we'll add to the images - noise = torch.randn(trajectory.shape, device=trajectory.device) - bsz = trajectory.shape[0] - # Sample a random timestep for each image - timesteps = torch.randint( - 0, self.noise_scheduler.config.num_train_timesteps, - (bsz,), device=trajectory.device - ).long() - # Add noise to the clean images according to the noise magnitude at each timestep - # (this is the forward diffusion process) - noisy_trajectory = self.noise_scheduler.add_noise( - trajectory, noise, timesteps) - - # compute loss mask - loss_mask = ~condition_mask - - # apply conditioning - noisy_trajectory[condition_mask] = cond_data[condition_mask] - - # Predict the noise residual - pred = self.model(noisy_trajectory, timesteps, - local_cond=local_cond, global_cond=global_cond) - - pred_type = self.noise_scheduler.config.prediction_type - if pred_type == 'epsilon': - target = noise - elif pred_type == 'sample': - target = trajectory - else: - raise ValueError(f"Unsupported prediction type {pred_type}") - - loss = F.mse_loss(pred, target, reduction='none') - loss = loss * loss_mask.type(loss.dtype) - loss = reduce(loss, 'b ... -> b (...)', 'mean') - loss = loss.mean() - return loss diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_image_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_image_policy.py deleted file mode 100644 index 2aaf876ea..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_image_policy.py +++ /dev/null @@ -1,259 +0,0 @@ -from typing import Dict -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange, reduce -from diffusers.schedulers.scheduling_ddpm import DDPMScheduler - -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.policy.base_image_policy import BaseImagePolicy -from diffusion_policy.model.diffusion.conditional_unet1d import ConditionalUnet1D -from diffusion_policy.model.diffusion.mask_generator import LowdimMaskGenerator -from diffusion_policy.model.vision.multi_image_obs_encoder import MultiImageObsEncoder -from diffusion_policy.common.pytorch_util import dict_apply - -class DiffusionUnetImagePolicy(BaseImagePolicy): - def __init__(self, - shape_meta: dict, - noise_scheduler: DDPMScheduler, - obs_encoder: MultiImageObsEncoder, - horizon, - n_action_steps, - n_obs_steps, - num_inference_steps=None, - obs_as_global_cond=True, - diffusion_step_embed_dim=256, - down_dims=(256,512,1024), - kernel_size=5, - n_groups=8, - cond_predict_scale=True, - # parameters passed to step - **kwargs): - super().__init__() - - # parse shapes - action_shape = shape_meta['action']['shape'] - assert len(action_shape) == 1 - action_dim = action_shape[0] - # get feature dim - obs_feature_dim = obs_encoder.output_shape()[0] - - # create diffusion model - input_dim = action_dim + obs_feature_dim - global_cond_dim = None - if obs_as_global_cond: - input_dim = action_dim - global_cond_dim = obs_feature_dim * n_obs_steps - - model = ConditionalUnet1D( - input_dim=input_dim, - local_cond_dim=None, - global_cond_dim=global_cond_dim, - diffusion_step_embed_dim=diffusion_step_embed_dim, - down_dims=down_dims, - kernel_size=kernel_size, - n_groups=n_groups, - cond_predict_scale=cond_predict_scale - ) - - self.obs_encoder = obs_encoder - self.model = model - self.noise_scheduler = noise_scheduler - self.mask_generator = LowdimMaskGenerator( - action_dim=action_dim, - obs_dim=0 if obs_as_global_cond else obs_feature_dim, - max_n_obs_steps=n_obs_steps, - fix_obs_steps=True, - action_visible=False - ) - self.normalizer = LinearNormalizer() - self.horizon = horizon - self.obs_feature_dim = obs_feature_dim - self.action_dim = action_dim - self.n_action_steps = n_action_steps - self.n_obs_steps = n_obs_steps - self.obs_as_global_cond = obs_as_global_cond - self.kwargs = kwargs - - if num_inference_steps is None: - num_inference_steps = noise_scheduler.config.num_train_timesteps - self.num_inference_steps = num_inference_steps - - # ========= inference ============ - def conditional_sample(self, - condition_data, condition_mask, - local_cond=None, global_cond=None, - generator=None, - # keyword arguments to scheduler.step - **kwargs - ): - model = self.model - scheduler = self.noise_scheduler - - trajectory = torch.randn( - size=condition_data.shape, - dtype=condition_data.dtype, - device=condition_data.device, - generator=generator) - - # set step values - scheduler.set_timesteps(self.num_inference_steps) - - for t in scheduler.timesteps: - # 1. apply conditioning - trajectory[condition_mask] = condition_data[condition_mask] - - # 2. predict model output - model_output = model(trajectory, t, - local_cond=local_cond, global_cond=global_cond) - - # 3. compute previous image: x_t -> x_t-1 - trajectory = scheduler.step( - model_output, t, trajectory, - generator=generator, - **kwargs - ).prev_sample - - # finally make sure conditioning is enforced - trajectory[condition_mask] = condition_data[condition_mask] - - return trajectory - - - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """ - obs_dict: must include "obs" key - result: must include "action" key - """ - assert 'past_action' not in obs_dict # not implemented yet - # normalize input - nobs = self.normalizer.normalize(obs_dict) - value = next(iter(nobs.values())) - B, To = value.shape[:2] - T = self.horizon - Da = self.action_dim - Do = self.obs_feature_dim - To = self.n_obs_steps - - # build input - device = self.device - dtype = self.dtype - - # handle different ways of passing observation - local_cond = None - global_cond = None - if self.obs_as_global_cond: - # condition through global feature - this_nobs = dict_apply(nobs, lambda x: x[:,:To,...].reshape(-1,*x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, Do - global_cond = nobs_features.reshape(B, -1) - # empty data for action - cond_data = torch.zeros(size=(B, T, Da), device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - else: - # condition through impainting - this_nobs = dict_apply(nobs, lambda x: x[:,:To,...].reshape(-1,*x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, T, Do - nobs_features = nobs_features.reshape(B, To, -1) - cond_data = torch.zeros(size=(B, T, Da+Do), device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - cond_data[:,:To,Da:] = nobs_features - cond_mask[:,:To,Da:] = True - - # run sampling - nsample = self.conditional_sample( - cond_data, - cond_mask, - local_cond=local_cond, - global_cond=global_cond, - **self.kwargs) - - # unnormalize prediction - naction_pred = nsample[...,:Da] - action_pred = self.normalizer['action'].unnormalize(naction_pred) - - # get action - start = To - 1 - end = start + self.n_action_steps - action = action_pred[:,start:end] - - result = { - 'action': action, - 'action_pred': action_pred - } - return result - - # ========= training ============ - def set_normalizer(self, normalizer: LinearNormalizer): - self.normalizer.load_state_dict(normalizer.state_dict()) - - def compute_loss(self, batch): - # normalize input - assert 'valid_mask' not in batch - nobs = self.normalizer.normalize(batch['obs']) - nactions = self.normalizer['action'].normalize(batch['action']) - batch_size = nactions.shape[0] - horizon = nactions.shape[1] - - # handle different ways of passing observation - local_cond = None - global_cond = None - trajectory = nactions - cond_data = trajectory - if self.obs_as_global_cond: - # reshape B, T, ... to B*T - this_nobs = dict_apply(nobs, - lambda x: x[:,:self.n_obs_steps,...].reshape(-1,*x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, Do - global_cond = nobs_features.reshape(batch_size, -1) - else: - # reshape B, T, ... to B*T - this_nobs = dict_apply(nobs, lambda x: x.reshape(-1, *x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, T, Do - nobs_features = nobs_features.reshape(batch_size, horizon, -1) - cond_data = torch.cat([nactions, nobs_features], dim=-1) - trajectory = cond_data.detach() - - # generate impainting mask - condition_mask = self.mask_generator(trajectory.shape) - - # Sample noise that we'll add to the images - noise = torch.randn(trajectory.shape, device=trajectory.device) - bsz = trajectory.shape[0] - # Sample a random timestep for each image - timesteps = torch.randint( - 0, self.noise_scheduler.config.num_train_timesteps, - (bsz,), device=trajectory.device - ).long() - # Add noise to the clean images according to the noise magnitude at each timestep - # (this is the forward diffusion process) - noisy_trajectory = self.noise_scheduler.add_noise( - trajectory, noise, timesteps) - - # compute loss mask - loss_mask = ~condition_mask - - # apply conditioning - noisy_trajectory[condition_mask] = cond_data[condition_mask] - - # Predict the noise residual - pred = self.model(noisy_trajectory, timesteps, - local_cond=local_cond, global_cond=global_cond) - - pred_type = self.noise_scheduler.config.prediction_type - if pred_type == 'epsilon': - target = noise - elif pred_type == 'sample': - target = trajectory - else: - raise ValueError(f"Unsupported prediction type {pred_type}") - - loss = F.mse_loss(pred, target, reduction='none') - loss = loss * loss_mask.type(loss.dtype) - loss = reduce(loss, 'b ... -> b (...)', 'mean') - loss = loss.mean() - return loss diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_lowdim_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_lowdim_policy.py deleted file mode 100644 index c01ed7a8c..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_lowdim_policy.py +++ /dev/null @@ -1,252 +0,0 @@ -from typing import Dict -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange, reduce -from diffusers.schedulers.scheduling_ddpm import DDPMScheduler - -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.policy.base_lowdim_policy import BaseLowdimPolicy -from diffusion_policy.model.diffusion.conditional_unet1d import ConditionalUnet1D -from diffusion_policy.model.diffusion.mask_generator import LowdimMaskGenerator - -class DiffusionUnetLowdimPolicy(BaseLowdimPolicy): - def __init__(self, - model: ConditionalUnet1D, - noise_scheduler: DDPMScheduler, - horizon, - obs_dim, - action_dim, - n_action_steps, - n_obs_steps, - num_inference_steps=None, - obs_as_local_cond=False, - obs_as_global_cond=False, - pred_action_steps_only=False, - oa_step_convention=False, - # parameters passed to step - **kwargs): - super().__init__() - assert not (obs_as_local_cond and obs_as_global_cond) - if pred_action_steps_only: - assert obs_as_global_cond - self.model = model - self.noise_scheduler = noise_scheduler - self.mask_generator = LowdimMaskGenerator( - action_dim=action_dim, - obs_dim=0 if (obs_as_local_cond or obs_as_global_cond) else obs_dim, - max_n_obs_steps=n_obs_steps, - fix_obs_steps=True, - action_visible=False - ) - self.normalizer = LinearNormalizer() - self.horizon = horizon - self.obs_dim = obs_dim - self.action_dim = action_dim - self.n_action_steps = n_action_steps - self.n_obs_steps = n_obs_steps - self.obs_as_local_cond = obs_as_local_cond - self.obs_as_global_cond = obs_as_global_cond - self.pred_action_steps_only = pred_action_steps_only - self.oa_step_convention = oa_step_convention - self.kwargs = kwargs - - if num_inference_steps is None: - num_inference_steps = noise_scheduler.config.num_train_timesteps - self.num_inference_steps = num_inference_steps - - # ========= inference ============ - def conditional_sample(self, - condition_data, condition_mask, - local_cond=None, global_cond=None, - generator=None, - # keyword arguments to scheduler.step - **kwargs - ): - model = self.model - scheduler = self.noise_scheduler - - trajectory = torch.randn( - size=condition_data.shape, - dtype=condition_data.dtype, - device=condition_data.device, - generator=generator) - - # set step values - scheduler.set_timesteps(self.num_inference_steps) - - for t in scheduler.timesteps: - # 1. apply conditioning - trajectory[condition_mask] = condition_data[condition_mask] - - # 2. predict model output - model_output = model(trajectory, t, - local_cond=local_cond, global_cond=global_cond) - - # 3. compute previous image: x_t -> x_t-1 - trajectory = scheduler.step( - model_output, t, trajectory, - generator=generator, - **kwargs - ).prev_sample - - # finally make sure conditioning is enforced - trajectory[condition_mask] = condition_data[condition_mask] - - return trajectory - - - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """ - obs_dict: must include "obs" key - result: must include "action" key - """ - - assert 'obs' in obs_dict - assert 'past_action' not in obs_dict # not implemented yet - nobs = self.normalizer['obs'].normalize(obs_dict['obs']) - B, _, Do = nobs.shape - To = self.n_obs_steps - assert Do == self.obs_dim - T = self.horizon - Da = self.action_dim - - # build input - device = self.device - dtype = self.dtype - - # handle different ways of passing observation - local_cond = None - global_cond = None - if self.obs_as_local_cond: - # condition through local feature - # all zero except first To timesteps - local_cond = torch.zeros(size=(B,T,Do), device=device, dtype=dtype) - local_cond[:,:To] = nobs[:,:To] - shape = (B, T, Da) - cond_data = torch.zeros(size=shape, device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - elif self.obs_as_global_cond: - # condition throught global feature - global_cond = nobs[:,:To].reshape(nobs.shape[0], -1) - shape = (B, T, Da) - if self.pred_action_steps_only: - shape = (B, self.n_action_steps, Da) - cond_data = torch.zeros(size=shape, device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - else: - # condition through impainting - shape = (B, T, Da+Do) - cond_data = torch.zeros(size=shape, device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - cond_data[:,:To,Da:] = nobs[:,:To] - cond_mask[:,:To,Da:] = True - - # run sampling - nsample = self.conditional_sample( - cond_data, - cond_mask, - local_cond=local_cond, - global_cond=global_cond, - **self.kwargs) - - # unnormalize prediction - naction_pred = nsample[...,:Da] - action_pred = self.normalizer['action'].unnormalize(naction_pred) - - # get action - if self.pred_action_steps_only: - action = action_pred - else: - start = To - if self.oa_step_convention: - start = To - 1 - end = start + self.n_action_steps - action = action_pred[:,start:end] - - result = { - 'action': action, - 'action_pred': action_pred - } - if not (self.obs_as_local_cond or self.obs_as_global_cond): - nobs_pred = nsample[...,Da:] - obs_pred = self.normalizer['obs'].unnormalize(nobs_pred) - action_obs_pred = obs_pred[:,start:end] - result['action_obs_pred'] = action_obs_pred - result['obs_pred'] = obs_pred - return result - - # ========= training ============ - def set_normalizer(self, normalizer: LinearNormalizer): - self.normalizer.load_state_dict(normalizer.state_dict()) - - def compute_loss(self, batch): - # normalize input - assert 'valid_mask' not in batch - nbatch = self.normalizer.normalize(batch) - obs = nbatch['obs'] - action = nbatch['action'] - - # handle different ways of passing observation - local_cond = None - global_cond = None - trajectory = action - if self.obs_as_local_cond: - # zero out observations after n_obs_steps - local_cond = obs - local_cond[:,self.n_obs_steps:,:] = 0 - elif self.obs_as_global_cond: - global_cond = obs[:,:self.n_obs_steps,:].reshape( - obs.shape[0], -1) - if self.pred_action_steps_only: - To = self.n_obs_steps - start = To - if self.oa_step_convention: - start = To - 1 - end = start + self.n_action_steps - trajectory = action[:,start:end] - else: - trajectory = torch.cat([action, obs], dim=-1) - - # generate impainting mask - if self.pred_action_steps_only: - condition_mask = torch.zeros_like(trajectory, dtype=torch.bool) - else: - condition_mask = self.mask_generator(trajectory.shape) - - # Sample noise that we'll add to the images - noise = torch.randn(trajectory.shape, device=trajectory.device) - bsz = trajectory.shape[0] - # Sample a random timestep for each image - timesteps = torch.randint( - 0, self.noise_scheduler.config.num_train_timesteps, - (bsz,), device=trajectory.device - ).long() - # Add noise to the clean images according to the noise magnitude at each timestep - # (this is the forward diffusion process) - noisy_trajectory = self.noise_scheduler.add_noise( - trajectory, noise, timesteps) - - # compute loss mask - loss_mask = ~condition_mask - - # apply conditioning - noisy_trajectory[condition_mask] = trajectory[condition_mask] - - # Predict the noise residual - pred = self.model(noisy_trajectory, timesteps, - local_cond=local_cond, global_cond=global_cond) - - pred_type = self.noise_scheduler.config.prediction_type - if pred_type == 'epsilon': - target = noise - elif pred_type == 'sample': - target = trajectory - else: - raise ValueError(f"Unsupported prediction type {pred_type}") - - loss = F.mse_loss(pred, target, reduction='none') - loss = loss * loss_mask.type(loss.dtype) - loss = reduce(loss, 'b ... -> b (...)', 'mean') - loss = loss.mean() - return loss diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_video_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_video_policy.py deleted file mode 100644 index eed968a4e..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/diffusion_unet_video_policy.py +++ /dev/null @@ -1,313 +0,0 @@ -from typing import Dict -import copy -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange, reduce -from diffusers.schedulers.scheduling_ddpm import DDPMScheduler - -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.policy.base_image_policy import BaseImagePolicy -from diffusion_policy.model.diffusion.conditional_unet1d import ConditionalUnet1D -from diffusion_policy.model.diffusion.mask_generator import LowdimMaskGenerator -from diffusion_policy.model.common.shape_util import get_output_shape -from diffusion_policy.model.obs_encoder.temporal_aggregator import TemporalAggregator - - -class DiffusionUnetVideoPolicy(BaseImagePolicy): - def __init__(self, - shape_meta: dict, - noise_scheduler: DDPMScheduler, - rgb_net: nn.Module, # (B,T,C,H,W) -> (B,Do) - horizon, - n_action_steps, - n_obs_steps, - num_inference_steps=None, - lowdim_as_global_cond=True, - diffusion_step_embed_dim=256, - down_dims=(256,512,1024), - kernel_size=5, - n_groups=8, - cond_predict_scale=True, - # parameters for TemporalAggregator - channel_mults=(1,1), - n_blocks_per_level=1, - ta_kernel_size=3, - ta_n_groups=8, - # parameters passed to step - **kwargs): - super().__init__() - - # parse shape_meta - action_shape = shape_meta['action']['shape'] - assert len(action_shape) == 1 - action_dim = action_shape[0] - obs_shape_meta = shape_meta['obs'] - - rgb_nets_map = nn.ModuleDict() - rgb_feature_dims = list() - lowdim_keys = list() - lowdim_input_dims = list() - - for key, attr in obs_shape_meta.items(): - shape = tuple(attr['shape']) - type = attr.get('type', 'lowdim') - if type == 'rgb': - # assign network for each rgb input - if len(rgb_nets_map) == 0: - net = rgb_net - else: - net = copy.deepcopy(rgb_net) - rgb_nets_map[key] = net - - # video input with n_obs_steps timesteps - shape = (n_obs_steps,) + shape - # compute output shape - output_shape = get_output_shape(shape, net) - assert(len(output_shape) == 1) - rgb_feature_dims.append(output_shape[0]) - elif type == 'lowdim': - lowdim_keys.append(key) - assert(len(shape) == 1) - lowdim_input_dims.append(shape[0]) - - # the order decides concatenation order - # dict preserves insertion order - # rgb and then lowdim - self.rgb_nets_map = rgb_nets_map - self.lowdim_keys = lowdim_keys - self.lowdim_net = None - - # compute dimensions for diffusion - rgb_feature_dim = sum(rgb_feature_dims) - lowdim_input_dim = sum(lowdim_input_dims) - global_cond_dim = rgb_feature_dim - input_dim = action_dim - if lowdim_as_global_cond: - lowdim_net = TemporalAggregator( - in_channels=lowdim_input_dim, - channel_mults=channel_mults, - n_blocks_per_level=n_blocks_per_level, - kernel_size=ta_kernel_size, - n_groups=ta_n_groups - ) - self.lowdim_net = lowdim_net - lowdim_feature_shape = get_output_shape( - (n_obs_steps, lowdim_input_dim), lowdim_net) - assert len(lowdim_feature_shape) == 1 - global_cond_dim += lowdim_feature_shape[0] - else: - input_dim += lowdim_input_dim - - model = ConditionalUnet1D( - input_dim=input_dim, - local_cond_dim=None, - global_cond_dim=global_cond_dim, - diffusion_step_embed_dim=diffusion_step_embed_dim, - down_dims=down_dims, - kernel_size=kernel_size, - n_groups=n_groups, - cond_predict_scale=cond_predict_scale - ) - - self.model = model - self.noise_scheduler = noise_scheduler - self.mask_generator = LowdimMaskGenerator( - action_dim=action_dim, - obs_dim=0 if lowdim_as_global_cond else lowdim_input_dim, - max_n_obs_steps=n_obs_steps, - fix_obs_steps=True, - action_visible=False - ) - self.normalizer = LinearNormalizer() - self.horizon = horizon - self.action_dim = action_dim - self.lowdim_input_dim = lowdim_input_dim - self.n_action_steps = n_action_steps - self.n_obs_steps = n_obs_steps - self.lowdim_as_global_cond = lowdim_as_global_cond - self.kwargs = kwargs - - if num_inference_steps is None: - num_inference_steps = noise_scheduler.config.num_train_timesteps - self.num_inference_steps = num_inference_steps - - # ========= inference ============ - def conditional_sample(self, - condition_data, condition_mask, - local_cond=None, global_cond=None, - generator=None, - # keyword arguments to scheduler.step - **kwargs - ): - model = self.model - scheduler = self.noise_scheduler - - trajectory = torch.randn( - size=condition_data.shape, - dtype=condition_data.dtype, - device=condition_data.device, - generator=generator) - - # set step values - scheduler.set_timesteps(self.num_inference_steps) - - for t in scheduler.timesteps: - # 1. apply conditioning - trajectory[condition_mask] = condition_data[condition_mask] - - # 2. predict model output - model_output = model(trajectory, t, - local_cond=local_cond, global_cond=global_cond) - - # 3. compute previous image: x_t -> x_t-1 - trajectory = scheduler.step( - model_output, t, trajectory, - generator=generator, - **kwargs - ).prev_sample - - # finally make sure conditioning is enforced - trajectory[condition_mask] = condition_data[condition_mask] - - return trajectory - - - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """ - obs_dict: must include "obs" key - result: must include "action" key - """ - assert 'past_action' not in obs_dict # not implemented yet - # normalize input - nobs = self.normalizer.normalize(obs_dict) - value = next(iter(nobs.values())) - B, To = value.shape[:2] - T = self.horizon - Da = self.action_dim - To = self.n_obs_steps - - # build input - device = self.device - dtype = self.dtype - - # run encoder first - # python 3.6+ dict preserves order - rgb_features_map = dict() - for key, net in self.rgb_nets_map.items(): - rgb_features_map[key] = net(nobs[key][:,:self.n_obs_steps]) - rgb_feature = torch.cat(list(rgb_features_map.values()), dim=-1) - - lowdim_input = torch.cat([nobs[k] for k in self.lowdim_keys], dim=-1) - - # handle different ways of passing lowdim - global_cond = None - cond_data = None - cond_mask = None - if self.lowdim_as_global_cond: - lowdim_feature = self.lowdim_net(lowdim_input[:,:To]) - global_cond = torch.cat([rgb_feature, lowdim_feature], dim=-1) - # empty data for action - cond_data = torch.zeros(size=(B, T, Da), device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - else: - global_cond = rgb_feature - cond_data = torch.zeros(size=(B, T, Da+self.lowdim_input_dim), device=device, dtype=dtype) - cond_mask = torch.zeros_like(cond_data, dtype=torch.bool) - cond_data[:,:To,Da:] = lowdim_input[:,:To] - cond_mask[:,:To,Da:] = True - - # run sampling - nsample = self.conditional_sample( - cond_data, - cond_mask, - local_cond=None, - global_cond=global_cond, - **self.kwargs) - - # unnormalize prediction - naction_pred = nsample[...,:Da] - action_pred = self.normalizer['action'].unnormalize(naction_pred) - - # get action - start = To - end = start + self.n_action_steps - action = action_pred[:,start:end] - - result = { - 'action': action, - 'action_pred': action_pred - } - return result - - # ========= training ============ - def set_normalizer(self, normalizer: LinearNormalizer): - self.normalizer.load_state_dict(normalizer.state_dict()) - - def compute_loss(self, batch): - # normalize input - assert 'valid_mask' not in batch - nobs = self.normalizer.normalize(batch['obs']) - nactions = self.normalizer['action'].normalize(batch['action']) - - # run encoder first - # python 3.6+ dict preserves order - rgb_features_map = dict() - for key, net in self.rgb_nets_map.items(): - rgb_features_map[key] = net(nobs[key][:,self.n_obs_steps:]) - rgb_feature = torch.cat(list(rgb_features_map.values()), dim=-1) - - lowdim_input = torch.cat([nobs[k] for k in self.lowdim_keys], axis=-1) - - # handle different ways of passing lowdim - global_cond = None - trajectory = None - cond_data = None - if self.lowdim_as_global_cond: - lowdim_feature = self.lowdim_net(lowdim_input[:,:self.n_obs_steps]) - global_cond = torch.cat([rgb_feature, lowdim_feature], dim=-1) - trajectory = nactions - cond_data = nactions - else: - global_cond = rgb_feature - trajectory = torch.cat([nactions, lowdim_input], dim=-1) - cond_data = trajectory - - # generate impainting mask - condition_mask = self.mask_generator(trajectory.shape) - - # Sample noise that we'll add to the images - noise = torch.randn(trajectory.shape, device=trajectory.device) - bsz = trajectory.shape[0] - # Sample a random timestep for each image - timesteps = torch.randint( - 0, self.noise_scheduler.config.num_train_timesteps, - (bsz,), device=trajectory.device - ).long() - # Add noise to the clean images according to the noise magnitude at each timestep - # (this is the forward diffusion process) - noisy_trajectory = self.noise_scheduler.add_noise( - trajectory, noise, timesteps) - - # compute loss mask - loss_mask = ~condition_mask - - # apply conditioning - noisy_trajectory[condition_mask] = cond_data[condition_mask] - - # Predict the noise residual - pred = self.model(noisy_trajectory, timesteps, - local_cond=None, global_cond=global_cond) - - if self.kwargs.get('predict_epsilon', True): - # default for most methods - target = noise - else: - # DDPM also has - target = trajectory - - loss = F.mse_loss(pred, target, reduction='none') - loss = loss * loss_mask.type(loss.dtype) - loss = reduce(loss, 'b ... -> b (...)', 'mean') - loss = loss.mean() - return loss diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/ibc_dfo_hybrid_image_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/ibc_dfo_hybrid_image_policy.py deleted file mode 100644 index c10e7deb2..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/ibc_dfo_hybrid_image_policy.py +++ /dev/null @@ -1,322 +0,0 @@ -from typing import Dict, Tuple -import torch -import torch.nn as nn -import torch.nn.functional as F -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.policy.base_image_policy import BaseImagePolicy -from diffusion_policy.common.robomimic_config_util import get_robomimic_config -from robomimic.algo import algo_factory -from robomimic.algo.algo import PolicyAlgo -import robomimic.utils.obs_utils as ObsUtils -import robomimic.models.base_nets as rmbn -import diffusion_policy.model.vision.crop_randomizer as dmvc -from diffusion_policy.common.pytorch_util import dict_apply, replace_submodules - - -class IbcDfoHybridImagePolicy(BaseImagePolicy): - def __init__(self, - shape_meta: dict, - horizon, - n_action_steps, - n_obs_steps, - dropout=0.1, - train_n_neg=128, - pred_n_iter=5, - pred_n_samples=16384, - kevin_inference=False, - andy_train=False, - obs_encoder_group_norm=True, - eval_fixed_crop=True, - crop_shape=(76, 76), - ): - super().__init__() - - # parse shape_meta - action_shape = shape_meta['action']['shape'] - assert len(action_shape) == 1 - action_dim = action_shape[0] - obs_shape_meta = shape_meta['obs'] - obs_config = { - 'low_dim': [], - 'rgb': [], - 'depth': [], - 'scan': [] - } - obs_key_shapes = dict() - for key, attr in obs_shape_meta.items(): - shape = attr['shape'] - obs_key_shapes[key] = list(shape) - - type = attr.get('type', 'low_dim') - if type == 'rgb': - obs_config['rgb'].append(key) - elif type == 'low_dim': - obs_config['low_dim'].append(key) - else: - raise RuntimeError(f"Unsupported obs type: {type}") - - # get raw robomimic config - config = get_robomimic_config( - algo_name='bc_rnn', - hdf5_type='image', - task_name='square', - dataset_type='ph') - - with config.unlocked(): - # set config with shape_meta - config.observation.modalities.obs = obs_config - - if crop_shape is None: - for key, modality in config.observation.encoder.items(): - if modality.obs_randomizer_class == 'CropRandomizer': - modality['obs_randomizer_class'] = None - else: - # set random crop parameter - ch, cw = crop_shape - for key, modality in config.observation.encoder.items(): - if modality.obs_randomizer_class == 'CropRandomizer': - modality.obs_randomizer_kwargs.crop_height = ch - modality.obs_randomizer_kwargs.crop_width = cw - - # init global state - ObsUtils.initialize_obs_utils_with_config(config) - - # load model - policy: PolicyAlgo = algo_factory( - algo_name=config.algo_name, - config=config, - obs_key_shapes=obs_key_shapes, - ac_dim=action_dim, - device='cpu', - ) - - self.obs_encoder = obs_encoder = policy.nets['policy'].nets['encoder'].nets['obs'] - - if obs_encoder_group_norm: - # replace batch norm with group norm - replace_submodules( - root_module=obs_encoder, - predicate=lambda x: isinstance(x, nn.BatchNorm2d), - func=lambda x: nn.GroupNorm( - num_groups=x.num_features//16, - num_channels=x.num_features) - ) - # obs_encoder.obs_nets['agentview_image'].nets[0].nets - - # obs_encoder.obs_randomizers['agentview_image'] - if eval_fixed_crop: - replace_submodules( - root_module=obs_encoder, - predicate=lambda x: isinstance(x, rmbn.CropRandomizer), - func=lambda x: dmvc.CropRandomizer( - input_shape=x.input_shape, - crop_height=x.crop_height, - crop_width=x.crop_width, - num_crops=x.num_crops, - pos_enc=x.pos_enc - ) - ) - - obs_feature_dim = obs_encoder.output_shape()[0] - in_action_channels = action_dim * n_action_steps - in_obs_channels = obs_feature_dim * n_obs_steps - in_channels = in_action_channels + in_obs_channels - mid_channels = 1024 - out_channels = 1 - - self.dense0 = nn.Linear(in_features=in_channels, out_features=mid_channels) - self.drop0 = nn.Dropout(dropout) - self.dense1 = nn.Linear(in_features=mid_channels, out_features=mid_channels) - self.drop1 = nn.Dropout(dropout) - self.dense2 = nn.Linear(in_features=mid_channels, out_features=mid_channels) - self.drop2 = nn.Dropout(dropout) - self.dense3 = nn.Linear(in_features=mid_channels, out_features=mid_channels) - self.drop3 = nn.Dropout(dropout) - self.dense4 = nn.Linear(in_features=mid_channels, out_features=out_channels) - - self.normalizer = LinearNormalizer() - - self.train_n_neg = train_n_neg - self.pred_n_iter = pred_n_iter - self.pred_n_samples = pred_n_samples - self.obs_feature_dim = obs_feature_dim - self.action_dim = action_dim - self.n_obs_steps = n_obs_steps - self.n_action_steps = n_action_steps - self.horizon = horizon - self.kevin_inference = kevin_inference - self.andy_train = andy_train - - def forward(self, obs, action): - B, N, Ta, Da = action.shape - B, To, Do = obs.shape - s = obs.reshape(B,1,-1).expand(-1,N,-1) - x = torch.cat([s, action.reshape(B,N,-1)], dim=-1).reshape(B*N,-1) - x = self.drop0(torch.relu(self.dense0(x))) - x = self.drop1(torch.relu(self.dense1(x))) - x = self.drop2(torch.relu(self.dense2(x))) - x = self.drop3(torch.relu(self.dense3(x))) - x = self.dense4(x) - x = x.reshape(B,N) - return x - - # ========= inference ============ - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """ - obs_dict: must include "obs" key - result: must include "action" key - """ - assert 'past_action' not in obs_dict # not implemented yet - # normalize input - nobs = self.normalizer.normalize(obs_dict) - value = next(iter(nobs.values())) - B, To = value.shape[:2] - T = self.horizon - Ta = self.n_action_steps - Da = self.action_dim - Do = self.obs_feature_dim - To = self.n_obs_steps - - # build input - device = self.device - dtype = self.dtype - - # encode obs - # reshape B, T, ... to B*T - this_nobs = dict_apply(nobs, - lambda x: x[:,:To,...].reshape(-1,*x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, To, Do - nobs_features = nobs_features.reshape(B,To,-1) - - # only take necessary obs - naction_stats = self.get_naction_stats() - - # first sample - action_dist = torch.distributions.Uniform( - low=naction_stats['min'], - high=naction_stats['max'] - ) - samples = action_dist.sample((B, self.pred_n_samples, Ta)).to( - dtype=dtype) - # (B, N, Ta, Da) - - if self.kevin_inference: - # kevin's implementation - noise_scale = 3e-2 - for i in range(self.pred_n_iter): - # Compute energies. - logits = self.forward(nobs_features, samples) - probs = F.softmax(logits, dim=-1) - - # Resample with replacement. - idxs = torch.multinomial(probs, self.pred_n_samples, replacement=True) - samples = samples[torch.arange(samples.size(0)).unsqueeze(-1), idxs] - - # Add noise and clip to target bounds. - samples = samples + torch.randn_like(samples) * noise_scale - samples = samples.clamp(min=naction_stats['min'], max=naction_stats['max']) - - # Return target with highest probability. - logits = self.forward(nobs_features, samples) - probs = F.softmax(logits, dim=-1) - best_idxs = probs.argmax(dim=-1) - acts_n = samples[torch.arange(samples.size(0)), best_idxs, :] - else: - # andy's implementation - zero = torch.tensor(0, device=self.device) - resample_std = torch.tensor(3e-2, device=self.device) - for i in range(self.pred_n_iter): - # Forward pass. - logits = self.forward(nobs_features, samples) # (B, N) - prob = torch.softmax(logits, dim=-1) - - if i < (self.pred_n_iter - 1): - idxs = torch.multinomial(prob, self.pred_n_samples, replacement=True) - samples = samples[torch.arange(samples.size(0)).unsqueeze(-1), idxs] - samples += torch.normal(zero, resample_std, size=samples.shape, device=self.device) - - # Return one sample per x in batch. - idxs = torch.multinomial(prob, num_samples=1, replacement=True) - acts_n = samples[torch.arange(samples.size(0)).unsqueeze(-1), idxs].squeeze(1) - - action = self.normalizer['action'].unnormalize(acts_n) - result = { - 'action': action - } - return result - - # ========= training ============ - def set_normalizer(self, normalizer: LinearNormalizer): - self.normalizer.load_state_dict(normalizer.state_dict()) - - def compute_loss(self, batch): - # normalize input - assert 'valid_mask' not in batch - nobs = self.normalizer.normalize(batch['obs']) - naction = self.normalizer['action'].normalize(batch['action']) - - # shapes - Do = self.obs_feature_dim - Da = self.action_dim - To = self.n_obs_steps - Ta = self.n_action_steps - T = self.horizon - B = naction.shape[0] - - # encode obs - # reshape B, T, ... to B*T - this_nobs = dict_apply(nobs, - lambda x: x[:,:To,...].reshape(-1,*x.shape[2:])) - nobs_features = self.obs_encoder(this_nobs) - # reshape back to B, To, Do - nobs_features = nobs_features.reshape(B,To,-1) - - start = To - 1 - end = start + Ta - this_action = naction[:,start:end] - - # Small additive noise to true positives. - this_action += torch.normal(mean=0, std=1e-4, - size=this_action.shape, - dtype=this_action.dtype, - device=this_action.device) - - # Sample negatives: (B, train_n_neg, Ta, Da) - naction_stats = self.get_naction_stats() - action_dist = torch.distributions.Uniform( - low=naction_stats['min'], - high=naction_stats['max'] - ) - samples = action_dist.sample((B, self.train_n_neg, Ta)).to( - dtype=this_action.dtype) - action_samples = torch.cat([ - this_action.unsqueeze(1), samples], dim=1) - # (B, train_n_neg+1, Ta, Da) - - if self.andy_train: - # Get onehot labels - labels = torch.zeros(action_samples.shape[:2], - dtype=this_action.dtype, device=this_action.device) - labels[:,0] = 1 - logits = self.forward(nobs_features, action_samples) - # (B, N) - logits = torch.log_softmax(logits, dim=-1) - loss = -torch.mean(torch.sum(logits * labels, axis=-1)) - else: - labels = torch.zeros((B,),dtype=torch.int64, device=this_action.device) - # training - logits = self.forward(nobs_features, action_samples) - loss = F.cross_entropy(logits, labels) - return loss - - def get_naction_stats(self): - Da = self.action_dim - naction_stats = self.normalizer['action'].get_output_stats() - repeated_stats = dict() - for key, value in naction_stats.items(): - assert len(value.shape) == 1 - n_repeats = Da // value.shape[0] - assert value.shape[0] * n_repeats == Da - repeated_stats[key] = value.repeat(n_repeats) - return repeated_stats diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/ibc_dfo_lowdim_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/ibc_dfo_lowdim_policy.py deleted file mode 100644 index 046f38415..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/ibc_dfo_lowdim_policy.py +++ /dev/null @@ -1,209 +0,0 @@ -from typing import Dict, Tuple -import torch -import torch.nn as nn -import torch.nn.functional as F -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.policy.base_lowdim_policy import BaseLowdimPolicy - -class IbcDfoLowdimPolicy(BaseLowdimPolicy): - def __init__(self, - horizon, - obs_dim, - action_dim, - n_action_steps, - n_obs_steps, - dropout=0.1, - train_n_neg=128, - pred_n_iter=5, - pred_n_samples=16384, - kevin_inference=False, - andy_train=False - ): - super().__init__() - - in_action_channels = action_dim * n_action_steps - in_obs_channels = obs_dim * n_obs_steps - in_channels = in_action_channels + in_obs_channels - mid_channels = 1024 - out_channels = 1 - - self.dense0 = nn.Linear(in_features=in_channels, out_features=mid_channels) - self.drop0 = nn.Dropout(dropout) - self.dense1 = nn.Linear(in_features=mid_channels, out_features=mid_channels) - self.drop1 = nn.Dropout(dropout) - self.dense2 = nn.Linear(in_features=mid_channels, out_features=mid_channels) - self.drop2 = nn.Dropout(dropout) - self.dense3 = nn.Linear(in_features=mid_channels, out_features=mid_channels) - self.drop3 = nn.Dropout(dropout) - self.dense4 = nn.Linear(in_features=mid_channels, out_features=out_channels) - - self.normalizer = LinearNormalizer() - - self.train_n_neg = train_n_neg - self.pred_n_iter = pred_n_iter - self.pred_n_samples = pred_n_samples - self.obs_dim = obs_dim - self.action_dim = action_dim - self.n_obs_steps = n_obs_steps - self.n_action_steps = n_action_steps - self.horizon = horizon - self.kevin_inference = kevin_inference - self.andy_train = andy_train - - def forward(self, obs, action): - B, N, Ta, Da = action.shape - B, To, Do = obs.shape - s = obs.reshape(B,1,-1).expand(-1,N,-1) - x = torch.cat([s, action.reshape(B,N,-1)], dim=-1).reshape(B*N,-1) - x = self.drop0(torch.relu(self.dense0(x))) - x = self.drop1(torch.relu(self.dense1(x))) - x = self.drop2(torch.relu(self.dense2(x))) - x = self.drop3(torch.relu(self.dense3(x))) - x = self.dense4(x) - x = x.reshape(B,N) - return x - - # ========= inference ============ - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """ - obs_dict: must include "obs" key - result: must include "action" key - """ - - assert 'obs' in obs_dict - assert 'past_action' not in obs_dict # not implemented yet - nobs = self.normalizer['obs'].normalize(obs_dict['obs']) - B, _, Do = nobs.shape - To = self.n_obs_steps - assert Do == self.obs_dim - T = self.horizon - Da = self.action_dim - Ta = self.n_action_steps - - # only take necessary obs - this_obs = nobs[:,:To] - naction_stats = self.get_naction_stats() - - # first sample - action_dist = torch.distributions.Uniform( - low=naction_stats['min'], - high=naction_stats['max'] - ) - samples = action_dist.sample((B, self.pred_n_samples, Ta)).to( - dtype=this_obs.dtype) - # (B, N, Ta, Da) - - if self.kevin_inference: - # kevin's implementation - noise_scale = 3e-2 - for i in range(self.pred_n_iter): - # Compute energies. - logits = self.forward(this_obs, samples) - probs = F.softmax(logits, dim=-1) - - # Resample with replacement. - idxs = torch.multinomial(probs, self.pred_n_samples, replacement=True) - samples = samples[torch.arange(samples.size(0)).unsqueeze(-1), idxs] - - # Add noise and clip to target bounds. - samples = samples + torch.randn_like(samples) * noise_scale - samples = samples.clamp(min=naction_stats['min'], max=naction_stats['max']) - - # Return target with highest probability. - logits = self.forward(this_obs, samples) - probs = F.softmax(logits, dim=-1) - best_idxs = probs.argmax(dim=-1) - acts_n = samples[torch.arange(samples.size(0)), best_idxs, :] - else: - # andy's implementation - zero = torch.tensor(0, device=self.device) - resample_std = torch.tensor(3e-2, device=self.device) - for i in range(self.pred_n_iter): - # Forward pass. - logits = self.forward(this_obs, samples) # (B, N) - prob = torch.softmax(logits, dim=-1) - - if i < (self.pred_n_iter - 1): - idxs = torch.multinomial(prob, self.pred_n_samples, replacement=True) - samples = samples[torch.arange(samples.size(0)).unsqueeze(-1), idxs] - samples += torch.normal(zero, resample_std, size=samples.shape, device=self.device) - - # Return one sample per x in batch. - idxs = torch.multinomial(prob, num_samples=1, replacement=True) - acts_n = samples[torch.arange(samples.size(0)).unsqueeze(-1), idxs].squeeze(1) - - action = self.normalizer['action'].unnormalize(acts_n) - result = { - 'action': action - } - return result - - # ========= training ============ - def set_normalizer(self, normalizer: LinearNormalizer): - self.normalizer.load_state_dict(normalizer.state_dict()) - - def compute_loss(self, batch): - # normalize input - assert 'valid_mask' not in batch - nbatch = self.normalizer.normalize(batch) - nobs = nbatch['obs'] - naction = nbatch['action'] - - # shapes - Do = self.obs_dim - Da = self.action_dim - To = self.n_obs_steps - Ta = self.n_action_steps - T = self.horizon - B = naction.shape[0] - - this_obs = nobs[:,:To] - start = To - 1 - end = start + Ta - this_action = naction[:,start:end] - - # Small additive noise to true positives. - this_action += torch.normal(mean=0, std=1e-4, - size=this_action.shape, - dtype=this_action.dtype, - device=this_action.device) - - # Sample negatives: (B, train_n_neg, Ta, Da) - naction_stats = self.get_naction_stats() - action_dist = torch.distributions.Uniform( - low=naction_stats['min'], - high=naction_stats['max'] - ) - samples = action_dist.sample((B, self.train_n_neg, Ta)).to( - dtype=this_action.dtype) - action_samples = torch.cat([ - this_action.unsqueeze(1), samples], dim=1) - # (B, train_n_neg+1, Ta, Da) - - if self.andy_train: - # Get onehot labels - labels = torch.zeros(action_samples.shape[:2], - dtype=this_action.dtype, device=this_action.device) - labels[:,0] = 1 - logits = self.forward(this_obs, action_samples) - # (B, N) - logits = torch.log_softmax(logits, dim=-1) - loss = -torch.mean(torch.sum(logits * labels, axis=-1)) - else: - labels = torch.zeros((B,),dtype=torch.int64, device=this_action.device) - # training - logits = self.forward(this_obs, action_samples) - loss = F.cross_entropy(logits, labels) - return loss - - - def get_naction_stats(self): - Da = self.action_dim - naction_stats = self.normalizer['action'].get_output_stats() - repeated_stats = dict() - for key, value in naction_stats.items(): - assert len(value.shape) == 1 - n_repeats = Da // value.shape[0] - assert value.shape[0] * n_repeats == Da - repeated_stats[key] = value.repeat(n_repeats) - return repeated_stats diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/robomimic_image_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/robomimic_image_policy.py deleted file mode 100644 index c2584f579..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/robomimic_image_policy.py +++ /dev/null @@ -1,142 +0,0 @@ -from typing import Dict -import torch -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.policy.base_image_policy import BaseImagePolicy -from diffusion_policy.common.pytorch_util import dict_apply - -from robomimic.algo import algo_factory -from robomimic.algo.algo import PolicyAlgo -import robomimic.utils.obs_utils as ObsUtils -from diffusion_policy.common.robomimic_config_util import get_robomimic_config - -class RobomimicImagePolicy(BaseImagePolicy): - def __init__(self, - shape_meta: dict, - algo_name='bc_rnn', - obs_type='image', - task_name='square', - dataset_type='ph', - crop_shape=(76,76) - ): - super().__init__() - - # parse shape_meta - action_shape = shape_meta['action']['shape'] - assert len(action_shape) == 1 - action_dim = action_shape[0] - obs_shape_meta = shape_meta['obs'] - obs_config = { - 'low_dim': [], - 'rgb': [], - 'depth': [], - 'scan': [] - } - obs_key_shapes = dict() - for key, attr in obs_shape_meta.items(): - shape = attr['shape'] - obs_key_shapes[key] = list(shape) - - type = attr.get('type', 'low_dim') - if type == 'rgb': - obs_config['rgb'].append(key) - elif type == 'low_dim': - obs_config['low_dim'].append(key) - else: - raise RuntimeError(f"Unsupported obs type: {type}") - - # get raw robomimic config - config = get_robomimic_config( - algo_name=algo_name, - hdf5_type=obs_type, - task_name=task_name, - dataset_type=dataset_type) - - - with config.unlocked(): - # set config with shape_meta - config.observation.modalities.obs = obs_config - - if crop_shape is None: - for key, modality in config.observation.encoder.items(): - if modality.obs_randomizer_class == 'CropRandomizer': - modality['obs_randomizer_class'] = None - else: - # set random crop parameter - ch, cw = crop_shape - for key, modality in config.observation.encoder.items(): - if modality.obs_randomizer_class == 'CropRandomizer': - modality.obs_randomizer_kwargs.crop_height = ch - modality.obs_randomizer_kwargs.crop_width = cw - - # init global state - ObsUtils.initialize_obs_utils_with_config(config) - - # load model - model: PolicyAlgo = algo_factory( - algo_name=config.algo_name, - config=config, - obs_key_shapes=obs_key_shapes, - ac_dim=action_dim, - device='cpu', - ) - - self.model = model - self.nets = model.nets - self.normalizer = LinearNormalizer() - self.config = config - - def to(self,*args,**kwargs): - device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs) - if device is not None: - self.model.device = device - super().to(*args,**kwargs) - - # =========== inference ============= - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - nobs_dict = self.normalizer(obs_dict) - robomimic_obs_dict = dict_apply(nobs_dict, lambda x: x[:,0,...]) - naction = self.model.get_action(robomimic_obs_dict) - action = self.normalizer['action'].unnormalize(naction) - # (B, Da) - result = { - 'action': action[:,None,:] # (B, 1, Da) - } - return result - - def reset(self): - self.model.reset() - - # =========== training ============== - def set_normalizer(self, normalizer: LinearNormalizer): - self.normalizer.load_state_dict(normalizer.state_dict()) - - def train_on_batch(self, batch, epoch, validate=False): - nobs = self.normalizer.normalize(batch['obs']) - nactions = self.normalizer['action'].normalize(batch['action']) - robomimic_batch = { - 'obs': nobs, - 'actions': nactions - } - input_batch = self.model.process_batch_for_training( - robomimic_batch) - info = self.model.train_on_batch( - batch=input_batch, epoch=epoch, validate=validate) - # keys: losses, predictions - return info - - def on_epoch_end(self, epoch): - self.model.on_epoch_end(epoch) - - def get_optimizer(self): - return self.model.optimizers['policy'] - - -def test(): - import os - from omegaconf import OmegaConf - cfg_path = os.path.expanduser('~/dev/diffusion_policy/diffusion_policy/config/task/lift_image.yaml') - cfg = OmegaConf.load(cfg_path) - shape_meta = cfg.shape_meta - - policy = RobomimicImagePolicy(shape_meta=shape_meta) - diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/robomimic_lowdim_policy.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/robomimic_lowdim_policy.py deleted file mode 100644 index 0c7b9a593..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/policy/robomimic_lowdim_policy.py +++ /dev/null @@ -1,87 +0,0 @@ -from typing import Dict -import torch -from diffusion_policy.model.common.normalizer import LinearNormalizer -from diffusion_policy.policy.base_lowdim_policy import BaseLowdimPolicy - -from robomimic.algo import algo_factory -from robomimic.algo.algo import PolicyAlgo -import robomimic.utils.obs_utils as ObsUtils -from diffusion_policy.common.robomimic_config_util import get_robomimic_config - -class RobomimicLowdimPolicy(BaseLowdimPolicy): - def __init__(self, - action_dim, - obs_dim, - algo_name='bc_rnn', - obs_type='low_dim', - task_name='square', - dataset_type='ph', - ): - super().__init__() - # key for robomimic obs input - # previously this is 'object', 'robot0_eef_pos' etc - obs_key = 'obs' - - config = get_robomimic_config( - algo_name=algo_name, - hdf5_type=obs_type, - task_name=task_name, - dataset_type=dataset_type) - with config.unlocked(): - config.observation.modalities.obs.low_dim = [obs_key] - - ObsUtils.initialize_obs_utils_with_config(config) - model: PolicyAlgo = algo_factory( - algo_name=config.algo_name, - config=config, - obs_key_shapes={obs_key: [obs_dim]}, - ac_dim=action_dim, - device='cpu', - ) - self.model = model - self.nets = model.nets - self.normalizer = LinearNormalizer() - self.obs_key = obs_key - self.config = config - - def to(self,*args,**kwargs): - device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs) - if device is not None: - self.model.device = device - super().to(*args,**kwargs) - - # =========== inference ============= - def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - obs = self.normalizer['obs'].normalize(obs_dict['obs']) - assert obs.shape[1] == 1 - robomimic_obs_dict = {self.obs_key: obs[:,0,:]} - naction = self.model.get_action(robomimic_obs_dict) - action = self.normalizer['action'].unnormalize(naction) - # (B, Da) - result = { - 'action': action[:,None,:] # (B, 1, Da) - } - return result - - def reset(self): - self.model.reset() - - # =========== training ============== - def set_normalizer(self, normalizer: LinearNormalizer): - self.normalizer.load_state_dict(normalizer.state_dict()) - - def train_on_batch(self, batch, epoch, validate=False): - nbatch = self.normalizer.normalize(batch) - robomimic_batch = { - 'obs': {self.obs_key: nbatch['obs']}, - 'actions': nbatch['action'] - } - input_batch = self.model.process_batch_for_training( - robomimic_batch) - info = self.model.train_on_batch( - batch=input_batch, epoch=epoch, validate=validate) - # keys: losses, predictions - return info - - def get_optimizer(self): - return self.model.optimizers['policy'] diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/keystroke_counter.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/keystroke_counter.py deleted file mode 100644 index 90147136d..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/keystroke_counter.py +++ /dev/null @@ -1,45 +0,0 @@ -from pynput.keyboard import Key, KeyCode, Listener -from collections import defaultdict -from threading import Lock - -class KeystrokeCounter(Listener): - def __init__(self): - self.key_count_map = defaultdict(lambda:0) - self.key_press_list = list() - self.lock = Lock() - super().__init__(on_press=self.on_press, on_release=self.on_release) - - def on_press(self, key): - with self.lock: - self.key_count_map[key] += 1 - self.key_press_list.append(key) - - def on_release(self, key): - pass - - def clear(self): - with self.lock: - self.key_count_map = defaultdict(lambda:0) - self.key_press_list = list() - - def __getitem__(self, key): - with self.lock: - return self.key_count_map[key] - - def get_press_events(self): - with self.lock: - events = list(self.key_press_list) - self.key_press_list = list() - return events - -if __name__ == '__main__': - import time - with KeystrokeCounter() as counter: - try: - while True: - print('Space:', counter[Key.space]) - print('q:', counter[KeyCode(char='q')]) - time.sleep(1/60) - except KeyboardInterrupt: - events = counter.get_press_events() - print(events) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/multi_camera_visualizer.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/multi_camera_visualizer.py deleted file mode 100644 index 84cbf9922..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/multi_camera_visualizer.py +++ /dev/null @@ -1,74 +0,0 @@ -import time -import multiprocessing as mp -import numpy as np -import cv2 -from threadpoolctl import threadpool_limits -from diffusion_policy.real_world.multi_realsense import MultiRealsense - -class MultiCameraVisualizer(mp.Process): - def __init__(self, - realsense: MultiRealsense, - row, col, - window_name='Multi Cam Vis', - vis_fps=60, - fill_value=0, - rgb_to_bgr=True - ): - super().__init__() - self.row = row - self.col = col - self.window_name = window_name - self.vis_fps = vis_fps - self.fill_value = fill_value - self.rgb_to_bgr=rgb_to_bgr - self.realsense = realsense - # shared variables - self.stop_event = mp.Event() - - def start(self, wait=False): - super().start() - - def stop(self, wait=False): - self.stop_event.set() - if wait: - self.stop_wait() - - def start_wait(self): - pass - - def stop_wait(self): - self.join() - - def run(self): - cv2.setNumThreads(1) - threadpool_limits(1) - channel_slice = slice(None) - if self.rgb_to_bgr: - channel_slice = slice(None,None,-1) - - vis_data = None - vis_img = None - while not self.stop_event.is_set(): - vis_data = self.realsense.get_vis(out=vis_data) - color = vis_data['color'] - N, H, W, C = color.shape - assert C == 3 - oh = H * self.row - ow = W * self.col - if vis_img is None: - vis_img = np.full((oh, ow, 3), - fill_value=self.fill_value, dtype=np.uint8) - for row in range(self.row): - for col in range(self.col): - idx = col + row * self.col - h_start = H * row - h_end = h_start + H - w_start = W * col - w_end = w_start + W - if idx < N: - # opencv uses bgr - vis_img[h_start:h_end,w_start:w_end - ] = color[idx,:,:,channel_slice] - cv2.imshow(self.window_name, vis_img) - cv2.pollKey() - time.sleep(1 / self.vis_fps) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/multi_realsense.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/multi_realsense.py deleted file mode 100644 index 3b29ad8db..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/multi_realsense.py +++ /dev/null @@ -1,224 +0,0 @@ -from typing import List, Optional, Union, Dict, Callable -import numbers -import time -import pathlib -from multiprocessing.managers import SharedMemoryManager -import numpy as np -import pyrealsense2 as rs -from diffusion_policy.real_world.single_realsense import SingleRealsense -from diffusion_policy.real_world.video_recorder import VideoRecorder - -class MultiRealsense: - def __init__(self, - serial_numbers: Optional[List[str]]=None, - shm_manager: Optional[SharedMemoryManager]=None, - resolution=(1280,720), - capture_fps=30, - put_fps=None, - put_downsample=True, - record_fps=None, - enable_color=True, - enable_depth=False, - enable_infrared=False, - get_max_k=30, - advanced_mode_config: Optional[Union[dict, List[dict]]]=None, - transform: Optional[Union[Callable[[Dict], Dict], List[Callable]]]=None, - vis_transform: Optional[Union[Callable[[Dict], Dict], List[Callable]]]=None, - recording_transform: Optional[Union[Callable[[Dict], Dict], List[Callable]]]=None, - video_recorder: Optional[Union[VideoRecorder, List[VideoRecorder]]]=None, - verbose=False - ): - if shm_manager is None: - shm_manager = SharedMemoryManager() - shm_manager.start() - if serial_numbers is None: - serial_numbers = SingleRealsense.get_connected_devices_serial() - n_cameras = len(serial_numbers) - - advanced_mode_config = repeat_to_list( - advanced_mode_config, n_cameras, dict) - transform = repeat_to_list( - transform, n_cameras, Callable) - vis_transform = repeat_to_list( - vis_transform, n_cameras, Callable) - recording_transform = repeat_to_list( - recording_transform, n_cameras, Callable) - - video_recorder = repeat_to_list( - video_recorder, n_cameras, VideoRecorder) - - cameras = dict() - for i, serial in enumerate(serial_numbers): - cameras[serial] = SingleRealsense( - shm_manager=shm_manager, - serial_number=serial, - resolution=resolution, - capture_fps=capture_fps, - put_fps=put_fps, - put_downsample=put_downsample, - record_fps=record_fps, - enable_color=enable_color, - enable_depth=enable_depth, - enable_infrared=enable_infrared, - get_max_k=get_max_k, - advanced_mode_config=advanced_mode_config[i], - transform=transform[i], - vis_transform=vis_transform[i], - recording_transform=recording_transform[i], - video_recorder=video_recorder[i], - verbose=verbose - ) - - self.cameras = cameras - self.shm_manager = shm_manager - - def __enter__(self): - self.start() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() - - @property - def n_cameras(self): - return len(self.cameras) - - @property - def is_ready(self): - is_ready = True - for camera in self.cameras.values(): - if not camera.is_ready: - is_ready = False - return is_ready - - def start(self, wait=True, put_start_time=None): - if put_start_time is None: - put_start_time = time.time() - for camera in self.cameras.values(): - camera.start(wait=False, put_start_time=put_start_time) - - if wait: - self.start_wait() - - def stop(self, wait=True): - for camera in self.cameras.values(): - camera.stop(wait=False) - - if wait: - self.stop_wait() - - def start_wait(self): - for camera in self.cameras.values(): - camera.start_wait() - - def stop_wait(self): - for camera in self.cameras.values(): - camera.join() - - def get(self, k=None, out=None) -> Dict[int, Dict[str, np.ndarray]]: - """ - Return order T,H,W,C - { - 0: { - 'rgb': (T,H,W,C), - 'timestamp': (T,) - }, - 1: ... - } - """ - if out is None: - out = dict() - for i, camera in enumerate(self.cameras.values()): - this_out = None - if i in out: - this_out = out[i] - this_out = camera.get(k=k, out=this_out) - out[i] = this_out - return out - - def get_vis(self, out=None): - results = list() - for i, camera in enumerate(self.cameras.values()): - this_out = None - if out is not None: - this_out = dict() - for key, v in out.items(): - # use the slicing trick to maintain the array - # when v is 1D - this_out[key] = v[i:i+1].reshape(v.shape[1:]) - this_out = camera.get_vis(out=this_out) - if out is None: - results.append(this_out) - if out is None: - out = dict() - for key in results[0].keys(): - out[key] = np.stack([x[key] for x in results]) - return out - - def set_color_option(self, option, value): - n_camera = len(self.cameras) - value = repeat_to_list(value, n_camera, numbers.Number) - for i, camera in enumerate(self.cameras.values()): - camera.set_color_option(option, value[i]) - - def set_exposure(self, exposure=None, gain=None): - """ - exposure: (1, 10000) 100us unit. (0.1 ms, 1/10000s) - gain: (0, 128) - """ - - if exposure is None and gain is None: - # auto exposure - self.set_color_option(rs.option.enable_auto_exposure, 1.0) - else: - # manual exposure - self.set_color_option(rs.option.enable_auto_exposure, 0.0) - if exposure is not None: - self.set_color_option(rs.option.exposure, exposure) - if gain is not None: - self.set_color_option(rs.option.gain, gain) - - def set_white_balance(self, white_balance=None): - if white_balance is None: - self.set_color_option(rs.option.enable_auto_white_balance, 1.0) - else: - self.set_color_option(rs.option.enable_auto_white_balance, 0.0) - self.set_color_option(rs.option.white_balance, white_balance) - - def get_intrinsics(self): - return np.array([c.get_intrinsics() for c in self.cameras.values()]) - - def get_depth_scale(self): - return np.array([c.get_depth_scale() for c in self.cameras.values()]) - - def start_recording(self, video_path: Union[str, List[str]], start_time: float): - if isinstance(video_path, str): - # directory - video_dir = pathlib.Path(video_path) - assert video_dir.parent.is_dir() - video_dir.mkdir(parents=True, exist_ok=True) - video_path = list() - for i in range(self.n_cameras): - video_path.append( - str(video_dir.joinpath(f'{i}.mp4').absolute())) - assert len(video_path) == self.n_cameras - - for i, camera in enumerate(self.cameras.values()): - camera.start_recording(video_path[i], start_time) - - def stop_recording(self): - for i, camera in enumerate(self.cameras.values()): - camera.stop_recording() - - def restart_put(self, start_time): - for camera in self.cameras.values(): - camera.restart_put(start_time) - - -def repeat_to_list(x, n: int, cls): - if x is None: - x = [None] * n - if isinstance(x, cls): - x = [x] * n - assert len(x) == n - return x diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/real_data_conversion.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/real_data_conversion.py deleted file mode 100644 index a75e4ccf9..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/real_data_conversion.py +++ /dev/null @@ -1,194 +0,0 @@ -from typing import Sequence, Tuple, Dict, Optional, Union -import os -import pathlib -import numpy as np -import av -import zarr -import numcodecs -import multiprocessing -import concurrent.futures -from tqdm import tqdm -from diffusion_policy.common.replay_buffer import ReplayBuffer, get_optimal_chunks -from diffusion_policy.common.cv2_util import get_image_transform -from diffusion_policy.real_world.video_recorder import read_video -from diffusion_policy.codecs.imagecodecs_numcodecs import ( - register_codecs, - Jpeg2k -) -register_codecs() - - -def real_data_to_replay_buffer( - dataset_path: str, - out_store: Optional[zarr.ABSStore]=None, - out_resolutions: Union[None, tuple, Dict[str,tuple]]=None, # (width, height) - lowdim_keys: Optional[Sequence[str]]=None, - image_keys: Optional[Sequence[str]]=None, - lowdim_compressor: Optional[numcodecs.abc.Codec]=None, - image_compressor: Optional[numcodecs.abc.Codec]=None, - n_decoding_threads: int=multiprocessing.cpu_count(), - n_encoding_threads: int=multiprocessing.cpu_count(), - max_inflight_tasks: int=multiprocessing.cpu_count()*5, - verify_read: bool=True - ) -> ReplayBuffer: - """ - It is recommended to use before calling this function - to avoid CPU oversubscription - cv2.setNumThreads(1) - threadpoolctl.threadpool_limits(1) - - out_resolution: - if None: - use video resolution - if (width, height) e.g. (1280, 720) - if dict: - camera_0: (1280, 720) - image_keys: ['camera_0', 'camera_1'] - """ - if out_store is None: - out_store = zarr.MemoryStore() - if n_decoding_threads <= 0: - n_decoding_threads = multiprocessing.cpu_count() - if n_encoding_threads <= 0: - n_encoding_threads = multiprocessing.cpu_count() - if image_compressor is None: - image_compressor = Jpeg2k(level=50) - - # verify input - input = pathlib.Path(os.path.expanduser(dataset_path)) - in_zarr_path = input.joinpath('replay_buffer.zarr') - in_video_dir = input.joinpath('videos') - assert in_zarr_path.is_dir() - assert in_video_dir.is_dir() - - in_replay_buffer = ReplayBuffer.create_from_path(str(in_zarr_path.absolute()), mode='r') - - # save lowdim data to single chunk - chunks_map = dict() - compressor_map = dict() - for key, value in in_replay_buffer.data.items(): - chunks_map[key] = value.shape - compressor_map[key] = lowdim_compressor - - print('Loading lowdim data') - out_replay_buffer = ReplayBuffer.copy_from_store( - src_store=in_replay_buffer.root.store, - store=out_store, - keys=lowdim_keys, - chunks=chunks_map, - compressors=compressor_map - ) - - # worker function - def put_img(zarr_arr, zarr_idx, img): - try: - zarr_arr[zarr_idx] = img - # make sure we can successfully decode - if verify_read: - _ = zarr_arr[zarr_idx] - return True - except Exception as e: - return False - - - n_cameras = 0 - camera_idxs = set() - if image_keys is not None: - n_cameras = len(image_keys) - camera_idxs = set(int(x.split('_')[-1]) for x in image_keys) - else: - # estimate number of cameras - episode_video_dir = in_video_dir.joinpath(str(0)) - episode_video_paths = sorted(episode_video_dir.glob('*.mp4'), key=lambda x: int(x.stem)) - camera_idxs = set(int(x.stem) for x in episode_video_paths) - n_cameras = len(episode_video_paths) - - n_steps = in_replay_buffer.n_steps - episode_starts = in_replay_buffer.episode_ends[:] - in_replay_buffer.episode_lengths[:] - episode_lengths = in_replay_buffer.episode_lengths - timestamps = in_replay_buffer['timestamp'][:] - dt = timestamps[1] - timestamps[0] - - with tqdm(total=n_steps*n_cameras, desc="Loading image data", mininterval=1.0) as pbar: - # one chunk per thread, therefore no synchronization needed - with concurrent.futures.ThreadPoolExecutor(max_workers=n_encoding_threads) as executor: - futures = set() - for episode_idx, episode_length in enumerate(episode_lengths): - episode_video_dir = in_video_dir.joinpath(str(episode_idx)) - episode_start = episode_starts[episode_idx] - - episode_video_paths = sorted(episode_video_dir.glob('*.mp4'), key=lambda x: int(x.stem)) - this_camera_idxs = set(int(x.stem) for x in episode_video_paths) - if image_keys is None: - for i in this_camera_idxs - camera_idxs: - print(f"Unexpected camera {i} at episode {episode_idx}") - for i in camera_idxs - this_camera_idxs: - print(f"Missing camera {i} at episode {episode_idx}") - if image_keys is not None: - raise RuntimeError(f"Missing camera {i} at episode {episode_idx}") - - for video_path in episode_video_paths: - camera_idx = int(video_path.stem) - if image_keys is not None: - # if image_keys provided, skip not used cameras - if camera_idx not in camera_idxs: - continue - - # read resolution - with av.open(str(video_path.absolute())) as container: - video = container.streams.video[0] - vcc = video.codec_context - this_res = (vcc.width, vcc.height) - in_img_res = this_res - - arr_name = f'camera_{camera_idx}' - # figure out save resolution - out_img_res = in_img_res - if isinstance(out_resolutions, dict): - if arr_name in out_resolutions: - out_img_res = tuple(out_resolutions[arr_name]) - elif out_resolutions is not None: - out_img_res = tuple(out_resolutions) - - # allocate array - if arr_name not in out_replay_buffer: - ow, oh = out_img_res - _ = out_replay_buffer.data.require_dataset( - name=arr_name, - shape=(n_steps,oh,ow,3), - chunks=(1,oh,ow,3), - compressor=image_compressor, - dtype=np.uint8 - ) - arr = out_replay_buffer[arr_name] - - image_tf = get_image_transform( - input_res=in_img_res, output_res=out_img_res, bgr_to_rgb=False) - for step_idx, frame in enumerate(read_video( - video_path=str(video_path), - dt=dt, - img_transform=image_tf, - thread_type='FRAME', - thread_count=n_decoding_threads - )): - if len(futures) >= max_inflight_tasks: - # limit number of inflight tasks - completed, futures = concurrent.futures.wait(futures, - return_when=concurrent.futures.FIRST_COMPLETED) - for f in completed: - if not f.result(): - raise RuntimeError('Failed to encode image!') - pbar.update(len(completed)) - - global_idx = episode_start + step_idx - futures.add(executor.submit(put_img, arr, global_idx, frame)) - - if step_idx == (episode_length - 1): - break - completed, futures = concurrent.futures.wait(futures) - for f in completed: - if not f.result(): - raise RuntimeError('Failed to encode image!') - pbar.update(len(completed)) - return out_replay_buffer - diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/real_env.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/real_env.py deleted file mode 100644 index b731205f1..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/real_env.py +++ /dev/null @@ -1,435 +0,0 @@ -from typing import Optional -import pathlib -import numpy as np -import time -import shutil -import math -from multiprocessing.managers import SharedMemoryManager -from diffusion_policy.real_world.rtde_interpolation_controller import RTDEInterpolationController -from diffusion_policy.real_world.multi_realsense import MultiRealsense, SingleRealsense -from diffusion_policy.real_world.video_recorder import VideoRecorder -from diffusion_policy.common.timestamp_accumulator import ( - TimestampObsAccumulator, - TimestampActionAccumulator, - align_timestamps -) -from diffusion_policy.real_world.multi_camera_visualizer import MultiCameraVisualizer -from diffusion_policy.common.replay_buffer import ReplayBuffer -from diffusion_policy.common.cv2_util import ( - get_image_transform, optimal_row_cols) - -DEFAULT_OBS_KEY_MAP = { - # robot - 'ActualTCPPose': 'robot_eef_pose', - 'ActualTCPSpeed': 'robot_eef_pose_vel', - 'ActualQ': 'robot_joint', - 'ActualQd': 'robot_joint_vel', - # timestamps - 'step_idx': 'step_idx', - 'timestamp': 'timestamp' -} - -class RealEnv: - def __init__(self, - # required params - output_dir, - robot_ip, - # env params - frequency=10, - n_obs_steps=2, - # obs - obs_image_resolution=(640,480), - max_obs_buffer_size=30, - camera_serial_numbers=None, - obs_key_map=DEFAULT_OBS_KEY_MAP, - obs_float32=False, - # action - max_pos_speed=0.25, - max_rot_speed=0.6, - # robot - tcp_offset=0.13, - init_joints=False, - # video capture params - video_capture_fps=30, - video_capture_resolution=(1280,720), - # saving params - record_raw_video=True, - thread_per_video=2, - video_crf=21, - # vis params - enable_multi_cam_vis=True, - multi_cam_vis_resolution=(1280,720), - # shared memory - shm_manager=None - ): - assert frequency <= video_capture_fps - output_dir = pathlib.Path(output_dir) - assert output_dir.parent.is_dir() - video_dir = output_dir.joinpath('videos') - video_dir.mkdir(parents=True, exist_ok=True) - zarr_path = str(output_dir.joinpath('replay_buffer.zarr').absolute()) - replay_buffer = ReplayBuffer.create_from_path( - zarr_path=zarr_path, mode='a') - - if shm_manager is None: - shm_manager = SharedMemoryManager() - shm_manager.start() - if camera_serial_numbers is None: - camera_serial_numbers = SingleRealsense.get_connected_devices_serial() - - color_tf = get_image_transform( - input_res=video_capture_resolution, - output_res=obs_image_resolution, - # obs output rgb - bgr_to_rgb=True) - color_transform = color_tf - if obs_float32: - color_transform = lambda x: color_tf(x).astype(np.float32) / 255 - - def transform(data): - data['color'] = color_transform(data['color']) - return data - - rw, rh, col, row = optimal_row_cols( - n_cameras=len(camera_serial_numbers), - in_wh_ratio=obs_image_resolution[0]/obs_image_resolution[1], - max_resolution=multi_cam_vis_resolution - ) - vis_color_transform = get_image_transform( - input_res=video_capture_resolution, - output_res=(rw,rh), - bgr_to_rgb=False - ) - def vis_transform(data): - data['color'] = vis_color_transform(data['color']) - return data - - recording_transfrom = None - recording_fps = video_capture_fps - recording_pix_fmt = 'bgr24' - if not record_raw_video: - recording_transfrom = transform - recording_fps = frequency - recording_pix_fmt = 'rgb24' - - video_recorder = VideoRecorder.create_h264( - fps=recording_fps, - codec='h264', - input_pix_fmt=recording_pix_fmt, - crf=video_crf, - thread_type='FRAME', - thread_count=thread_per_video) - - realsense = MultiRealsense( - serial_numbers=camera_serial_numbers, - shm_manager=shm_manager, - resolution=video_capture_resolution, - capture_fps=video_capture_fps, - put_fps=video_capture_fps, - # send every frame immediately after arrival - # ignores put_fps - put_downsample=False, - record_fps=recording_fps, - enable_color=True, - enable_depth=False, - enable_infrared=False, - get_max_k=max_obs_buffer_size, - transform=transform, - vis_transform=vis_transform, - recording_transform=recording_transfrom, - video_recorder=video_recorder, - verbose=False - ) - - multi_cam_vis = None - if enable_multi_cam_vis: - multi_cam_vis = MultiCameraVisualizer( - realsense=realsense, - row=row, - col=col, - rgb_to_bgr=False - ) - - cube_diag = np.linalg.norm([1,1,1]) - j_init = np.array([0,-90,-90,-90,90,0]) / 180 * np.pi - if not init_joints: - j_init = None - - robot = RTDEInterpolationController( - shm_manager=shm_manager, - robot_ip=robot_ip, - frequency=125, # UR5 CB3 RTDE - lookahead_time=0.1, - gain=300, - max_pos_speed=max_pos_speed*cube_diag, - max_rot_speed=max_rot_speed*cube_diag, - launch_timeout=3, - tcp_offset_pose=[0,0,tcp_offset,0,0,0], - payload_mass=None, - payload_cog=None, - joints_init=j_init, - joints_init_speed=1.05, - soft_real_time=False, - verbose=False, - receive_keys=None, - get_max_k=max_obs_buffer_size - ) - self.realsense = realsense - self.robot = robot - self.multi_cam_vis = multi_cam_vis - self.video_capture_fps = video_capture_fps - self.frequency = frequency - self.n_obs_steps = n_obs_steps - self.max_obs_buffer_size = max_obs_buffer_size - self.max_pos_speed = max_pos_speed - self.max_rot_speed = max_rot_speed - self.obs_key_map = obs_key_map - # recording - self.output_dir = output_dir - self.video_dir = video_dir - self.replay_buffer = replay_buffer - # temp memory buffers - self.last_realsense_data = None - # recording buffers - self.obs_accumulator = None - self.action_accumulator = None - self.stage_accumulator = None - - self.start_time = None - - # ======== start-stop API ============= - @property - def is_ready(self): - return self.realsense.is_ready and self.robot.is_ready - - def start(self, wait=True): - self.realsense.start(wait=False) - self.robot.start(wait=False) - if self.multi_cam_vis is not None: - self.multi_cam_vis.start(wait=False) - if wait: - self.start_wait() - - def stop(self, wait=True): - self.end_episode() - if self.multi_cam_vis is not None: - self.multi_cam_vis.stop(wait=False) - self.robot.stop(wait=False) - self.realsense.stop(wait=False) - if wait: - self.stop_wait() - - def start_wait(self): - self.realsense.start_wait() - self.robot.start_wait() - if self.multi_cam_vis is not None: - self.multi_cam_vis.start_wait() - - def stop_wait(self): - self.robot.stop_wait() - self.realsense.stop_wait() - if self.multi_cam_vis is not None: - self.multi_cam_vis.stop_wait() - - # ========= context manager =========== - def __enter__(self): - self.start() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() - - # ========= async env API =========== - def get_obs(self) -> dict: - "observation dict" - assert self.is_ready - - # get data - # 30 Hz, camera_receive_timestamp - k = math.ceil(self.n_obs_steps * (self.video_capture_fps / self.frequency)) - self.last_realsense_data = self.realsense.get( - k=k, - out=self.last_realsense_data) - - # 125 hz, robot_receive_timestamp - last_robot_data = self.robot.get_all_state() - # both have more than n_obs_steps data - - # align camera obs timestamps - dt = 1 / self.frequency - last_timestamp = np.max([x['timestamp'][-1] for x in self.last_realsense_data.values()]) - obs_align_timestamps = last_timestamp - (np.arange(self.n_obs_steps)[::-1] * dt) - - camera_obs = dict() - for camera_idx, value in self.last_realsense_data.items(): - this_timestamps = value['timestamp'] - this_idxs = list() - for t in obs_align_timestamps: - is_before_idxs = np.nonzero(this_timestamps < t)[0] - this_idx = 0 - if len(is_before_idxs) > 0: - this_idx = is_before_idxs[-1] - this_idxs.append(this_idx) - # remap key - camera_obs[f'camera_{camera_idx}'] = value['color'][this_idxs] - - # align robot obs - robot_timestamps = last_robot_data['robot_receive_timestamp'] - this_timestamps = robot_timestamps - this_idxs = list() - for t in obs_align_timestamps: - is_before_idxs = np.nonzero(this_timestamps < t)[0] - this_idx = 0 - if len(is_before_idxs) > 0: - this_idx = is_before_idxs[-1] - this_idxs.append(this_idx) - - robot_obs_raw = dict() - for k, v in last_robot_data.items(): - if k in self.obs_key_map: - robot_obs_raw[self.obs_key_map[k]] = v - - robot_obs = dict() - for k, v in robot_obs_raw.items(): - robot_obs[k] = v[this_idxs] - - # accumulate obs - if self.obs_accumulator is not None: - self.obs_accumulator.put( - robot_obs_raw, - robot_timestamps - ) - - # return obs - obs_data = dict(camera_obs) - obs_data.update(robot_obs) - obs_data['timestamp'] = obs_align_timestamps - return obs_data - - def exec_actions(self, - actions: np.ndarray, - timestamps: np.ndarray, - stages: Optional[np.ndarray]=None): - assert self.is_ready - if not isinstance(actions, np.ndarray): - actions = np.array(actions) - if not isinstance(timestamps, np.ndarray): - timestamps = np.array(timestamps) - if stages is None: - stages = np.zeros_like(timestamps, dtype=np.int64) - elif not isinstance(stages, np.ndarray): - stages = np.array(stages, dtype=np.int64) - - # convert action to pose - receive_time = time.time() - is_new = timestamps > receive_time - new_actions = actions[is_new] - new_timestamps = timestamps[is_new] - new_stages = stages[is_new] - - # schedule waypoints - for i in range(len(new_actions)): - self.robot.schedule_waypoint( - pose=new_actions[i], - target_time=new_timestamps[i] - ) - - # record actions - if self.action_accumulator is not None: - self.action_accumulator.put( - new_actions, - new_timestamps - ) - if self.stage_accumulator is not None: - self.stage_accumulator.put( - new_stages, - new_timestamps - ) - - def get_robot_state(self): - return self.robot.get_state() - - # recording API - def start_episode(self, start_time=None): - "Start recording and return first obs" - if start_time is None: - start_time = time.time() - self.start_time = start_time - - assert self.is_ready - - # prepare recording stuff - episode_id = self.replay_buffer.n_episodes - this_video_dir = self.video_dir.joinpath(str(episode_id)) - this_video_dir.mkdir(parents=True, exist_ok=True) - n_cameras = self.realsense.n_cameras - video_paths = list() - for i in range(n_cameras): - video_paths.append( - str(this_video_dir.joinpath(f'{i}.mp4').absolute())) - - # start recording on realsense - self.realsense.restart_put(start_time=start_time) - self.realsense.start_recording(video_path=video_paths, start_time=start_time) - - # create accumulators - self.obs_accumulator = TimestampObsAccumulator( - start_time=start_time, - dt=1/self.frequency - ) - self.action_accumulator = TimestampActionAccumulator( - start_time=start_time, - dt=1/self.frequency - ) - self.stage_accumulator = TimestampActionAccumulator( - start_time=start_time, - dt=1/self.frequency - ) - print(f'Episode {episode_id} started!') - - def end_episode(self): - "Stop recording" - assert self.is_ready - - # stop video recorder - self.realsense.stop_recording() - - if self.obs_accumulator is not None: - # recording - assert self.action_accumulator is not None - assert self.stage_accumulator is not None - - # Since the only way to accumulate obs and action is by calling - # get_obs and exec_actions, which will be in the same thread. - # We don't need to worry new data come in here. - obs_data = self.obs_accumulator.data - obs_timestamps = self.obs_accumulator.timestamps - - actions = self.action_accumulator.actions - action_timestamps = self.action_accumulator.timestamps - stages = self.stage_accumulator.actions - n_steps = min(len(obs_timestamps), len(action_timestamps)) - if n_steps > 0: - episode = dict() - episode['timestamp'] = obs_timestamps[:n_steps] - episode['action'] = actions[:n_steps] - episode['stage'] = stages[:n_steps] - for key, value in obs_data.items(): - episode[key] = value[:n_steps] - self.replay_buffer.add_episode(episode, compressors='disk') - episode_id = self.replay_buffer.n_episodes - 1 - print(f'Episode {episode_id} saved!') - - self.obs_accumulator = None - self.action_accumulator = None - self.stage_accumulator = None - - def drop_episode(self): - self.end_episode() - self.replay_buffer.drop_episode() - episode_id = self.replay_buffer.n_episodes - this_video_dir = self.video_dir.joinpath(str(episode_id)) - if this_video_dir.exists(): - shutil.rmtree(str(this_video_dir)) - print(f'Episode {episode_id} dropped!') - diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/real_inference_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/real_inference_util.py deleted file mode 100644 index 4a262fa95..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/real_inference_util.py +++ /dev/null @@ -1,52 +0,0 @@ -from typing import Dict, Callable, Tuple -import numpy as np -from diffusion_policy.common.cv2_util import get_image_transform - -def get_real_obs_dict( - env_obs: Dict[str, np.ndarray], - shape_meta: dict, - ) -> Dict[str, np.ndarray]: - obs_dict_np = dict() - obs_shape_meta = shape_meta['obs'] - for key, attr in obs_shape_meta.items(): - type = attr.get('type', 'low_dim') - shape = attr.get('shape') - if type == 'rgb': - this_imgs_in = env_obs[key] - t,hi,wi,ci = this_imgs_in.shape - co,ho,wo = shape - assert ci == co - out_imgs = this_imgs_in - if (ho != hi) or (wo != wi) or (this_imgs_in.dtype == np.uint8): - tf = get_image_transform( - input_res=(wi,hi), - output_res=(wo,ho), - bgr_to_rgb=False) - out_imgs = np.stack([tf(x) for x in this_imgs_in]) - if this_imgs_in.dtype == np.uint8: - out_imgs = out_imgs.astype(np.float32) / 255 - # THWC to TCHW - obs_dict_np[key] = np.moveaxis(out_imgs,-1,1) - elif type == 'low_dim': - this_data_in = env_obs[key] - if 'pose' in key and shape == (2,): - # take X,Y coordinates - this_data_in = this_data_in[...,[0,1]] - obs_dict_np[key] = this_data_in - return obs_dict_np - - -def get_real_obs_resolution( - shape_meta: dict - ) -> Tuple[int, int]: - out_res = None - obs_shape_meta = shape_meta['obs'] - for key, attr in obs_shape_meta.items(): - type = attr.get('type', 'low_dim') - shape = attr.get('shape') - if type == 'rgb': - co,ho,wo = shape - if out_res is None: - out_res = (wo, ho) - assert out_res == (wo, ho) - return out_res diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/realsense_config/415_high_accuracy_mode.json b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/realsense_config/415_high_accuracy_mode.json deleted file mode 100644 index a3c00b4d2..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/realsense_config/415_high_accuracy_mode.json +++ /dev/null @@ -1,95 +0,0 @@ -{ - "aux-param-autoexposure-setpoint": "400", - "aux-param-colorcorrection1": "0.461914", - "aux-param-colorcorrection10": "-0.553711", - "aux-param-colorcorrection11": "-0.553711", - "aux-param-colorcorrection12": "0.0458984", - "aux-param-colorcorrection2": "0.540039", - "aux-param-colorcorrection3": "0.540039", - "aux-param-colorcorrection4": "0.208008", - "aux-param-colorcorrection5": "-0.332031", - "aux-param-colorcorrection6": "-0.212891", - "aux-param-colorcorrection7": "-0.212891", - "aux-param-colorcorrection8": "0.68457", - "aux-param-colorcorrection9": "0.930664", - "aux-param-depthclampmax": "65535", - "aux-param-depthclampmin": "0", - "aux-param-disparityshift": "0", - "controls-autoexposure-auto": "True", - "controls-autoexposure-manual": "33000", - "controls-color-autoexposure-auto": "True", - "controls-color-autoexposure-manual": "100", - "controls-color-backlight-compensation": "0", - "controls-color-brightness": "0", - "controls-color-contrast": "50", - "controls-color-gain": "100", - "controls-color-gamma": "300", - "controls-color-hue": "0", - "controls-color-power-line-frequency": "3", - "controls-color-saturation": "64", - "controls-color-sharpness": "50", - "controls-color-white-balance-auto": "True", - "controls-color-white-balance-manual": "4600", - "controls-depth-gain": "16", - "controls-depth-white-balance-auto": "False", - "controls-laserpower": "150", - "controls-laserstate": "on", - "ignoreSAD": "0", - "param-amplitude-factor": "0", - "param-autoexposure-setpoint": "400", - "param-censusenablereg-udiameter": "9", - "param-censusenablereg-vdiameter": "3", - "param-censususize": "9", - "param-censusvsize": "3", - "param-depthclampmax": "65535", - "param-depthclampmin": "0", - "param-depthunits": "1000", - "param-disableraucolor": "0", - "param-disablesadcolor": "0", - "param-disablesadnormalize": "0", - "param-disablesloleftcolor": "0", - "param-disableslorightcolor": "1", - "param-disparitymode": "0", - "param-disparityshift": "0", - "param-lambdaad": "751", - "param-lambdacensus": "6", - "param-leftrightthreshold": "10", - "param-maxscorethreshb": "2893", - "param-medianthreshold": "796", - "param-minscorethresha": "4", - "param-neighborthresh": "108", - "param-raumine": "6", - "param-rauminn": "3", - "param-rauminnssum": "7", - "param-raumins": "2", - "param-rauminw": "2", - "param-rauminwesum": "12", - "param-regioncolorthresholdb": "0.785714", - "param-regioncolorthresholdg": "0.565558", - "param-regioncolorthresholdr": "0.985323", - "param-regionshrinku": "3", - "param-regionshrinkv": "0", - "param-robbinsmonrodecrement": "25", - "param-robbinsmonroincrement": "2", - "param-rsmdiffthreshold": "1.65625", - "param-rsmrauslodiffthreshold": "0.71875", - "param-rsmremovethreshold": "0.809524", - "param-scanlineedgetaub": "13", - "param-scanlineedgetaug": "15", - "param-scanlineedgetaur": "30", - "param-scanlinep1": "155", - "param-scanlinep1onediscon": "160", - "param-scanlinep1twodiscon": "59", - "param-scanlinep2": "190", - "param-scanlinep2onediscon": "507", - "param-scanlinep2twodiscon": "493", - "param-secondpeakdelta": "647", - "param-texturecountthresh": "0", - "param-texturedifferencethresh": "1722", - "param-usersm": "1", - "param-zunits": "1000", - "stream-depth-format": "Z16", - "stream-fps": "30", - "stream-height": "480", - "stream-width": "640" -} diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/realsense_config/435_high_accuracy_mode.json b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/realsense_config/435_high_accuracy_mode.json deleted file mode 100644 index 8b88e5f13..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/realsense_config/435_high_accuracy_mode.json +++ /dev/null @@ -1,94 +0,0 @@ -{ - "aux-param-autoexposure-setpoint": "1536", - "aux-param-colorcorrection1": "0.298828", - "aux-param-colorcorrection10": "-0", - "aux-param-colorcorrection11": "-0", - "aux-param-colorcorrection12": "-0", - "aux-param-colorcorrection2": "0.293945", - "aux-param-colorcorrection3": "0.293945", - "aux-param-colorcorrection4": "0.114258", - "aux-param-colorcorrection5": "-0", - "aux-param-colorcorrection6": "-0", - "aux-param-colorcorrection7": "-0", - "aux-param-colorcorrection8": "-0", - "aux-param-colorcorrection9": "-0", - "aux-param-depthclampmax": "65536", - "aux-param-depthclampmin": "0", - "aux-param-disparityshift": "0", - "controls-autoexposure-auto": "True", - "controls-autoexposure-manual": "8500", - "controls-color-autoexposure-auto": "True", - "controls-color-autoexposure-manual": "100", - "controls-color-backlight-compensation": "0", - "controls-color-brightness": "0", - "controls-color-contrast": "50", - "controls-color-gain": "100", - "controls-color-gamma": "300", - "controls-color-hue": "0", - "controls-color-power-line-frequency": "3", - "controls-color-saturation": "64", - "controls-color-sharpness": "50", - "controls-color-white-balance-auto": "True", - "controls-color-white-balance-manual": "4600", - "controls-depth-gain": "16", - "controls-laserpower": "150", - "controls-laserstate": "on", - "ignoreSAD": "0", - "param-amplitude-factor": "0", - "param-autoexposure-setpoint": "1536", - "param-censusenablereg-udiameter": "9", - "param-censusenablereg-vdiameter": "9", - "param-censususize": "9", - "param-censusvsize": "9", - "param-depthclampmax": "65536", - "param-depthclampmin": "0", - "param-depthunits": "1000", - "param-disableraucolor": "0", - "param-disablesadcolor": "0", - "param-disablesadnormalize": "0", - "param-disablesloleftcolor": "0", - "param-disableslorightcolor": "1", - "param-disparitymode": "0", - "param-disparityshift": "0", - "param-lambdaad": "751", - "param-lambdacensus": "6", - "param-leftrightthreshold": "10", - "param-maxscorethreshb": "2893", - "param-medianthreshold": "796", - "param-minscorethresha": "4", - "param-neighborthresh": "108", - "param-raumine": "6", - "param-rauminn": "3", - "param-rauminnssum": "7", - "param-raumins": "2", - "param-rauminw": "2", - "param-rauminwesum": "12", - "param-regioncolorthresholdb": "0.785714", - "param-regioncolorthresholdg": "0.565558", - "param-regioncolorthresholdr": "0.985323", - "param-regionshrinku": "3", - "param-regionshrinkv": "0", - "param-robbinsmonrodecrement": "25", - "param-robbinsmonroincrement": "2", - "param-rsmdiffthreshold": "1.65625", - "param-rsmrauslodiffthreshold": "0.71875", - "param-rsmremovethreshold": "0.809524", - "param-scanlineedgetaub": "13", - "param-scanlineedgetaug": "15", - "param-scanlineedgetaur": "30", - "param-scanlinep1": "155", - "param-scanlinep1onediscon": "160", - "param-scanlinep1twodiscon": "59", - "param-scanlinep2": "190", - "param-scanlinep2onediscon": "507", - "param-scanlinep2twodiscon": "493", - "param-secondpeakdelta": "647", - "param-texturecountthresh": "0", - "param-texturedifferencethresh": "1722", - "param-usersm": "1", - "param-zunits": "1000", - "stream-depth-format": "Z16", - "stream-fps": "30", - "stream-height": "480", - "stream-width": "848" -} \ No newline at end of file diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/rtde_interpolation_controller.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/rtde_interpolation_controller.py deleted file mode 100644 index af27b2ed4..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/rtde_interpolation_controller.py +++ /dev/null @@ -1,361 +0,0 @@ -import os -import time -import enum -import multiprocessing as mp -from multiprocessing.managers import SharedMemoryManager -import scipy.interpolate as si -import scipy.spatial.transform as st -import numpy as np -from rtde_control import RTDEControlInterface -from rtde_receive import RTDEReceiveInterface -from diffusion_policy.shared_memory.shared_memory_queue import ( - SharedMemoryQueue, Empty) -from diffusion_policy.shared_memory.shared_memory_ring_buffer import SharedMemoryRingBuffer -from diffusion_policy.common.pose_trajectory_interpolator import PoseTrajectoryInterpolator - -class Command(enum.Enum): - STOP = 0 - SERVOL = 1 - SCHEDULE_WAYPOINT = 2 - - -class RTDEInterpolationController(mp.Process): - """ - To ensure sending command to the robot with predictable latency - this controller need its separate process (due to python GIL) - """ - - - def __init__(self, - shm_manager: SharedMemoryManager, - robot_ip, - frequency=125, - lookahead_time=0.1, - gain=300, - max_pos_speed=0.25, # 5% of max speed - max_rot_speed=0.16, # 5% of max speed - launch_timeout=3, - tcp_offset_pose=None, - payload_mass=None, - payload_cog=None, - joints_init=None, - joints_init_speed=1.05, - soft_real_time=False, - verbose=False, - receive_keys=None, - get_max_k=128, - ): - """ - frequency: CB2=125, UR3e=500 - lookahead_time: [0.03, 0.2]s smoothens the trajectory with this lookahead time - gain: [100, 2000] proportional gain for following target position - max_pos_speed: m/s - max_rot_speed: rad/s - tcp_offset_pose: 6d pose - payload_mass: float - payload_cog: 3d position, center of gravity - soft_real_time: enables round-robin scheduling and real-time priority - requires running scripts/rtprio_setup.sh before hand. - - """ - # verify - assert 0 < frequency <= 500 - assert 0.03 <= lookahead_time <= 0.2 - assert 100 <= gain <= 2000 - assert 0 < max_pos_speed - assert 0 < max_rot_speed - if tcp_offset_pose is not None: - tcp_offset_pose = np.array(tcp_offset_pose) - assert tcp_offset_pose.shape == (6,) - if payload_mass is not None: - assert 0 <= payload_mass <= 5 - if payload_cog is not None: - payload_cog = np.array(payload_cog) - assert payload_cog.shape == (3,) - assert payload_mass is not None - if joints_init is not None: - joints_init = np.array(joints_init) - assert joints_init.shape == (6,) - - super().__init__(name="RTDEPositionalController") - self.robot_ip = robot_ip - self.frequency = frequency - self.lookahead_time = lookahead_time - self.gain = gain - self.max_pos_speed = max_pos_speed - self.max_rot_speed = max_rot_speed - self.launch_timeout = launch_timeout - self.tcp_offset_pose = tcp_offset_pose - self.payload_mass = payload_mass - self.payload_cog = payload_cog - self.joints_init = joints_init - self.joints_init_speed = joints_init_speed - self.soft_real_time = soft_real_time - self.verbose = verbose - - # build input queue - example = { - 'cmd': Command.SERVOL.value, - 'target_pose': np.zeros((6,), dtype=np.float64), - 'duration': 0.0, - 'target_time': 0.0 - } - input_queue = SharedMemoryQueue.create_from_examples( - shm_manager=shm_manager, - examples=example, - buffer_size=256 - ) - - # build ring buffer - if receive_keys is None: - receive_keys = [ - 'ActualTCPPose', - 'ActualTCPSpeed', - 'ActualQ', - 'ActualQd', - - 'TargetTCPPose', - 'TargetTCPSpeed', - 'TargetQ', - 'TargetQd' - ] - rtde_r = RTDEReceiveInterface(hostname=robot_ip) - example = dict() - for key in receive_keys: - example[key] = np.array(getattr(rtde_r, 'get'+key)()) - example['robot_receive_timestamp'] = time.time() - ring_buffer = SharedMemoryRingBuffer.create_from_examples( - shm_manager=shm_manager, - examples=example, - get_max_k=get_max_k, - get_time_budget=0.2, - put_desired_frequency=frequency - ) - - self.ready_event = mp.Event() - self.input_queue = input_queue - self.ring_buffer = ring_buffer - self.receive_keys = receive_keys - - # ========= launch method =========== - def start(self, wait=True): - super().start() - if wait: - self.start_wait() - if self.verbose: - print(f"[RTDEPositionalController] Controller process spawned at {self.pid}") - - def stop(self, wait=True): - message = { - 'cmd': Command.STOP.value - } - self.input_queue.put(message) - if wait: - self.stop_wait() - - def start_wait(self): - self.ready_event.wait(self.launch_timeout) - assert self.is_alive() - - def stop_wait(self): - self.join() - - @property - def is_ready(self): - return self.ready_event.is_set() - - # ========= context manager =========== - def __enter__(self): - self.start() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() - - # ========= command methods ============ - def servoL(self, pose, duration=0.1): - """ - duration: desired time to reach pose - """ - assert self.is_alive() - assert(duration >= (1/self.frequency)) - pose = np.array(pose) - assert pose.shape == (6,) - - message = { - 'cmd': Command.SERVOL.value, - 'target_pose': pose, - 'duration': duration - } - self.input_queue.put(message) - - def schedule_waypoint(self, pose, target_time): - assert target_time > time.time() - pose = np.array(pose) - assert pose.shape == (6,) - - message = { - 'cmd': Command.SCHEDULE_WAYPOINT.value, - 'target_pose': pose, - 'target_time': target_time - } - self.input_queue.put(message) - - # ========= receive APIs ============= - def get_state(self, k=None, out=None): - if k is None: - return self.ring_buffer.get(out=out) - else: - return self.ring_buffer.get_last_k(k=k,out=out) - - def get_all_state(self): - return self.ring_buffer.get_all() - - # ========= main loop in process ============ - def run(self): - # enable soft real-time - if self.soft_real_time: - os.sched_setscheduler( - 0, os.SCHED_RR, os.sched_param(20)) - - # start rtde - robot_ip = self.robot_ip - rtde_c = RTDEControlInterface(hostname=robot_ip) - rtde_r = RTDEReceiveInterface(hostname=robot_ip) - - try: - if self.verbose: - print(f"[RTDEPositionalController] Connect to robot: {robot_ip}") - - # set parameters - if self.tcp_offset_pose is not None: - rtde_c.setTcp(self.tcp_offset_pose) - if self.payload_mass is not None: - if self.payload_cog is not None: - assert rtde_c.setPayload(self.payload_mass, self.payload_cog) - else: - assert rtde_c.setPayload(self.payload_mass) - - # init pose - if self.joints_init is not None: - assert rtde_c.moveJ(self.joints_init, self.joints_init_speed, 1.4) - - # main loop - dt = 1. / self.frequency - curr_pose = rtde_r.getActualTCPPose() - # use monotonic time to make sure the control loop never go backward - curr_t = time.monotonic() - last_waypoint_time = curr_t - pose_interp = PoseTrajectoryInterpolator( - times=[curr_t], - poses=[curr_pose] - ) - - iter_idx = 0 - keep_running = True - while keep_running: - # start control iteration - t_start = rtde_c.initPeriod() - - # send command to robot - t_now = time.monotonic() - # diff = t_now - pose_interp.times[-1] - # if diff > 0: - # print('extrapolate', diff) - pose_command = pose_interp(t_now) - vel = 0.5 - acc = 0.5 - assert rtde_c.servoL(pose_command, - vel, acc, # dummy, not used by ur5 - dt, - self.lookahead_time, - self.gain) - - # update robot state - state = dict() - for key in self.receive_keys: - state[key] = np.array(getattr(rtde_r, 'get'+key)()) - state['robot_receive_timestamp'] = time.time() - self.ring_buffer.put(state) - - # fetch command from queue - try: - commands = self.input_queue.get_all() - n_cmd = len(commands['cmd']) - except Empty: - n_cmd = 0 - - # execute commands - for i in range(n_cmd): - command = dict() - for key, value in commands.items(): - command[key] = value[i] - cmd = command['cmd'] - - if cmd == Command.STOP.value: - keep_running = False - # stop immediately, ignore later commands - break - elif cmd == Command.SERVOL.value: - # since curr_pose always lag behind curr_target_pose - # if we start the next interpolation with curr_pose - # the command robot receive will have discontinouity - # and cause jittery robot behavior. - target_pose = command['target_pose'] - duration = float(command['duration']) - curr_time = t_now + dt - t_insert = curr_time + duration - pose_interp = pose_interp.drive_to_waypoint( - pose=target_pose, - time=t_insert, - curr_time=curr_time, - max_pos_speed=self.max_pos_speed, - max_rot_speed=self.max_rot_speed - ) - last_waypoint_time = t_insert - if self.verbose: - print("[RTDEPositionalController] New pose target:{} duration:{}s".format( - target_pose, duration)) - elif cmd == Command.SCHEDULE_WAYPOINT.value: - target_pose = command['target_pose'] - target_time = float(command['target_time']) - # translate global time to monotonic time - target_time = time.monotonic() - time.time() + target_time - curr_time = t_now + dt - pose_interp = pose_interp.schedule_waypoint( - pose=target_pose, - time=target_time, - max_pos_speed=self.max_pos_speed, - max_rot_speed=self.max_rot_speed, - curr_time=curr_time, - last_waypoint_time=last_waypoint_time - ) - last_waypoint_time = target_time - else: - keep_running = False - break - - # regulate frequency - rtde_c.waitPeriod(t_start) - - # first loop successful, ready to receive command - if iter_idx == 0: - self.ready_event.set() - iter_idx += 1 - - if self.verbose: - print(f"[RTDEPositionalController] Actual frequency {1/(time.perf_counter() - t_start)}") - - finally: - # manditory cleanup - # decelerate - rtde_c.servoStop() - - # terminate - rtde_c.stopScript() - rtde_c.disconnect() - rtde_r.disconnect() - self.ready_event.set() - - if self.verbose: - print(f"[RTDEPositionalController] Disconnected from robot: {robot_ip}") diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/single_realsense.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/single_realsense.py deleted file mode 100644 index 7a8443b96..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/single_realsense.py +++ /dev/null @@ -1,480 +0,0 @@ -from typing import Optional, Callable, Dict -import os -import enum -import time -import json -import numpy as np -import pyrealsense2 as rs -import multiprocessing as mp -import cv2 -from threadpoolctl import threadpool_limits -from multiprocessing.managers import SharedMemoryManager -from diffusion_policy.common.timestamp_accumulator import get_accumulate_timestamp_idxs -from diffusion_policy.shared_memory.shared_ndarray import SharedNDArray -from diffusion_policy.shared_memory.shared_memory_ring_buffer import SharedMemoryRingBuffer -from diffusion_policy.shared_memory.shared_memory_queue import SharedMemoryQueue, Full, Empty -from diffusion_policy.real_world.video_recorder import VideoRecorder - -class Command(enum.Enum): - SET_COLOR_OPTION = 0 - SET_DEPTH_OPTION = 1 - START_RECORDING = 2 - STOP_RECORDING = 3 - RESTART_PUT = 4 - -class SingleRealsense(mp.Process): - MAX_PATH_LENGTH = 4096 # linux path has a limit of 4096 bytes - - def __init__( - self, - shm_manager: SharedMemoryManager, - serial_number, - resolution=(1280,720), - capture_fps=30, - put_fps=None, - put_downsample=True, - record_fps=None, - enable_color=True, - enable_depth=False, - enable_infrared=False, - get_max_k=30, - advanced_mode_config=None, - transform: Optional[Callable[[Dict], Dict]] = None, - vis_transform: Optional[Callable[[Dict], Dict]] = None, - recording_transform: Optional[Callable[[Dict], Dict]] = None, - video_recorder: Optional[VideoRecorder] = None, - verbose=False - ): - super().__init__() - - if put_fps is None: - put_fps = capture_fps - if record_fps is None: - record_fps = capture_fps - - # create ring buffer - resolution = tuple(resolution) - shape = resolution[::-1] - examples = dict() - if enable_color: - examples['color'] = np.empty( - shape=shape+(3,), dtype=np.uint8) - if enable_depth: - examples['depth'] = np.empty( - shape=shape, dtype=np.uint16) - if enable_infrared: - examples['infrared'] = np.empty( - shape=shape, dtype=np.uint8) - examples['camera_capture_timestamp'] = 0.0 - examples['camera_receive_timestamp'] = 0.0 - examples['timestamp'] = 0.0 - examples['step_idx'] = 0 - - vis_ring_buffer = SharedMemoryRingBuffer.create_from_examples( - shm_manager=shm_manager, - examples=examples if vis_transform is None - else vis_transform(dict(examples)), - get_max_k=1, - get_time_budget=0.2, - put_desired_frequency=capture_fps - ) - - ring_buffer = SharedMemoryRingBuffer.create_from_examples( - shm_manager=shm_manager, - examples=examples if transform is None - else transform(dict(examples)), - get_max_k=get_max_k, - get_time_budget=0.2, - put_desired_frequency=put_fps - ) - - # create command queue - examples = { - 'cmd': Command.SET_COLOR_OPTION.value, - 'option_enum': rs.option.exposure.value, - 'option_value': 0.0, - 'video_path': np.array('a'*self.MAX_PATH_LENGTH), - 'recording_start_time': 0.0, - 'put_start_time': 0.0 - } - - command_queue = SharedMemoryQueue.create_from_examples( - shm_manager=shm_manager, - examples=examples, - buffer_size=128 - ) - - # create shared array for intrinsics - intrinsics_array = SharedNDArray.create_from_shape( - mem_mgr=shm_manager, - shape=(7,), - dtype=np.float64) - intrinsics_array.get()[:] = 0 - - # create video recorder - if video_recorder is None: - # realsense uses bgr24 pixel format - # default thread_type to FRAEM - # i.e. each frame uses one core - # instead of all cores working on all frames. - # this prevents CPU over-subpscription and - # improves performance significantly - video_recorder = VideoRecorder.create_h264( - fps=record_fps, - codec='h264', - input_pix_fmt='bgr24', - crf=18, - thread_type='FRAME', - thread_count=1) - - # copied variables - self.serial_number = serial_number - self.resolution = resolution - self.capture_fps = capture_fps - self.put_fps = put_fps - self.put_downsample = put_downsample - self.record_fps = record_fps - self.enable_color = enable_color - self.enable_depth = enable_depth - self.enable_infrared = enable_infrared - self.advanced_mode_config = advanced_mode_config - self.transform = transform - self.vis_transform = vis_transform - self.recording_transform = recording_transform - self.video_recorder = video_recorder - self.verbose = verbose - self.put_start_time = None - - # shared variables - self.stop_event = mp.Event() - self.ready_event = mp.Event() - self.ring_buffer = ring_buffer - self.vis_ring_buffer = vis_ring_buffer - self.command_queue = command_queue - self.intrinsics_array = intrinsics_array - - @staticmethod - def get_connected_devices_serial(): - serials = list() - for d in rs.context().devices: - if d.get_info(rs.camera_info.name).lower() != 'platform camera': - serial = d.get_info(rs.camera_info.serial_number) - product_line = d.get_info(rs.camera_info.product_line) - if product_line == 'D400': - # only works with D400 series - serials.append(serial) - serials = sorted(serials) - return serials - - # ========= context manager =========== - def __enter__(self): - self.start() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() - - # ========= user API =========== - def start(self, wait=True, put_start_time=None): - self.put_start_time = put_start_time - super().start() - if wait: - self.start_wait() - - def stop(self, wait=True): - self.stop_event.set() - if wait: - self.end_wait() - - def start_wait(self): - self.ready_event.wait() - - def end_wait(self): - self.join() - - @property - def is_ready(self): - return self.ready_event.is_set() - - def get(self, k=None, out=None): - if k is None: - return self.ring_buffer.get(out=out) - else: - return self.ring_buffer.get_last_k(k, out=out) - - def get_vis(self, out=None): - return self.vis_ring_buffer.get(out=out) - - # ========= user API =========== - def set_color_option(self, option: rs.option, value: float): - self.command_queue.put({ - 'cmd': Command.SET_COLOR_OPTION.value, - 'option_enum': option.value, - 'option_value': value - }) - - def set_exposure(self, exposure=None, gain=None): - """ - exposure: (1, 10000) 100us unit. (0.1 ms, 1/10000s) - gain: (0, 128) - """ - - if exposure is None and gain is None: - # auto exposure - self.set_color_option(rs.option.enable_auto_exposure, 1.0) - else: - # manual exposure - self.set_color_option(rs.option.enable_auto_exposure, 0.0) - if exposure is not None: - self.set_color_option(rs.option.exposure, exposure) - if gain is not None: - self.set_color_option(rs.option.gain, gain) - - def set_white_balance(self, white_balance=None): - if white_balance is None: - self.set_color_option(rs.option.enable_auto_white_balance, 1.0) - else: - self.set_color_option(rs.option.enable_auto_white_balance, 0.0) - self.set_color_option(rs.option.white_balance, white_balance) - - def get_intrinsics(self): - assert self.ready_event.is_set() - fx, fy, ppx, ppy = self.intrinsics_array.get()[:4] - mat = np.eye(3) - mat[0,0] = fx - mat[1,1] = fy - mat[0,2] = ppx - mat[1,2] = ppy - return mat - - def get_depth_scale(self): - assert self.ready_event.is_set() - scale = self.intrinsics_array.get()[-1] - return scale - - def start_recording(self, video_path: str, start_time: float=-1): - assert self.enable_color - - path_len = len(video_path.encode('utf-8')) - if path_len > self.MAX_PATH_LENGTH: - raise RuntimeError('video_path too long.') - self.command_queue.put({ - 'cmd': Command.START_RECORDING.value, - 'video_path': video_path, - 'recording_start_time': start_time - }) - - def stop_recording(self): - self.command_queue.put({ - 'cmd': Command.STOP_RECORDING.value - }) - - def restart_put(self, start_time): - self.command_queue.put({ - 'cmd': Command.RESTART_PUT.value, - 'put_start_time': start_time - }) - - # ========= interval API =========== - def run(self): - # limit threads - threadpool_limits(1) - cv2.setNumThreads(1) - - w, h = self.resolution - fps = self.capture_fps - align = rs.align(rs.stream.color) - # Enable the streams from all the intel realsense devices - rs_config = rs.config() - if self.enable_color: - rs_config.enable_stream(rs.stream.color, - w, h, rs.format.bgr8, fps) - if self.enable_depth: - rs_config.enable_stream(rs.stream.depth, - w, h, rs.format.z16, fps) - if self.enable_infrared: - rs_config.enable_stream(rs.stream.infrared, - w, h, rs.format.y8, fps) - - try: - rs_config.enable_device(self.serial_number) - - # start pipeline - pipeline = rs.pipeline() - pipeline_profile = pipeline.start(rs_config) - - # report global time - # https://github.com/IntelRealSense/librealsense/pull/3909 - d = pipeline_profile.get_device().first_color_sensor() - d.set_option(rs.option.global_time_enabled, 1) - - # setup advanced mode - if self.advanced_mode_config is not None: - json_text = json.dumps(self.advanced_mode_config) - device = pipeline_profile.get_device() - advanced_mode = rs.rs400_advanced_mode(device) - advanced_mode.load_json(json_text) - - # get - color_stream = pipeline_profile.get_stream(rs.stream.color) - intr = color_stream.as_video_stream_profile().get_intrinsics() - order = ['fx', 'fy', 'ppx', 'ppy', 'height', 'width'] - for i, name in enumerate(order): - self.intrinsics_array.get()[i] = getattr(intr, name) - - if self.enable_depth: - depth_sensor = pipeline_profile.get_device().first_depth_sensor() - depth_scale = depth_sensor.get_depth_scale() - self.intrinsics_array.get()[-1] = depth_scale - - # one-time setup (intrinsics etc, ignore for now) - if self.verbose: - print(f'[SingleRealsense {self.serial_number}] Main loop started.') - - # put frequency regulation - put_idx = None - put_start_time = self.put_start_time - if put_start_time is None: - put_start_time = time.time() - - iter_idx = 0 - t_start = time.time() - while not self.stop_event.is_set(): - # wait for frames to come in - frameset = pipeline.wait_for_frames() - receive_time = time.time() - # align frames to color - frameset = align.process(frameset) - - # grab data - data = dict() - data['camera_receive_timestamp'] = receive_time - # realsense report in ms - data['camera_capture_timestamp'] = frameset.get_timestamp() / 1000 - if self.enable_color: - color_frame = frameset.get_color_frame() - data['color'] = np.asarray(color_frame.get_data()) - t = color_frame.get_timestamp() / 1000 - data['camera_capture_timestamp'] = t - # print('device', time.time() - t) - # print(color_frame.get_frame_timestamp_domain()) - if self.enable_depth: - data['depth'] = np.asarray( - frameset.get_depth_frame().get_data()) - if self.enable_infrared: - data['infrared'] = np.asarray( - frameset.get_infrared_frame().get_data()) - - # apply transform - put_data = data - if self.transform is not None: - put_data = self.transform(dict(data)) - - if self.put_downsample: - # put frequency regulation - local_idxs, global_idxs, put_idx \ - = get_accumulate_timestamp_idxs( - timestamps=[receive_time], - start_time=put_start_time, - dt=1/self.put_fps, - # this is non in first iteration - # and then replaced with a concrete number - next_global_idx=put_idx, - # continue to pump frames even if not started. - # start_time is simply used to align timestamps. - allow_negative=True - ) - - for step_idx in global_idxs: - put_data['step_idx'] = step_idx - # put_data['timestamp'] = put_start_time + step_idx / self.put_fps - put_data['timestamp'] = receive_time - # print(step_idx, data['timestamp']) - self.ring_buffer.put(put_data, wait=False) - else: - step_idx = int((receive_time - put_start_time) * self.put_fps) - put_data['step_idx'] = step_idx - put_data['timestamp'] = receive_time - self.ring_buffer.put(put_data, wait=False) - - # signal ready - if iter_idx == 0: - self.ready_event.set() - - # put to vis - vis_data = data - if self.vis_transform == self.transform: - vis_data = put_data - elif self.vis_transform is not None: - vis_data = self.vis_transform(dict(data)) - self.vis_ring_buffer.put(vis_data, wait=False) - - # record frame - rec_data = data - if self.recording_transform == self.transform: - rec_data = put_data - elif self.recording_transform is not None: - rec_data = self.recording_transform(dict(data)) - - if self.video_recorder.is_ready(): - self.video_recorder.write_frame(rec_data['color'], - frame_time=receive_time) - - # perf - t_end = time.time() - duration = t_end - t_start - frequency = np.round(1 / duration, 1) - t_start = t_end - if self.verbose: - print(f'[SingleRealsense {self.serial_number}] FPS {frequency}') - - # fetch command from queue - try: - commands = self.command_queue.get_all() - n_cmd = len(commands['cmd']) - except Empty: - n_cmd = 0 - - # execute commands - for i in range(n_cmd): - command = dict() - for key, value in commands.items(): - command[key] = value[i] - cmd = command['cmd'] - if cmd == Command.SET_COLOR_OPTION.value: - sensor = pipeline_profile.get_device().first_color_sensor() - option = rs.option(command['option_enum']) - value = float(command['option_value']) - sensor.set_option(option, value) - # print('auto', sensor.get_option(rs.option.enable_auto_exposure)) - # print('exposure', sensor.get_option(rs.option.exposure)) - # print('gain', sensor.get_option(rs.option.gain)) - elif cmd == Command.SET_DEPTH_OPTION.value: - sensor = pipeline_profile.get_device().first_depth_sensor() - option = rs.option(command['option_enum']) - value = float(command['option_value']) - sensor.set_option(option, value) - elif cmd == Command.START_RECORDING.value: - video_path = str(command['video_path']) - start_time = command['recording_start_time'] - if start_time < 0: - start_time = None - self.video_recorder.start(video_path, start_time=start_time) - elif cmd == Command.STOP_RECORDING.value: - self.video_recorder.stop() - # stop need to flush all in-flight frames to disk, which might take longer than dt. - # soft-reset put to drop frames to prevent ring buffer overflow. - put_idx = None - elif cmd == Command.RESTART_PUT.value: - put_idx = None - put_start_time = command['put_start_time'] - # self.ring_buffer.clear() - - iter_idx += 1 - finally: - self.video_recorder.stop() - rs_config.disable_all_streams() - self.ready_event.set() - - if self.verbose: - print(f'[SingleRealsense {self.serial_number}] Exiting worker process.') diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/spacemouse.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/spacemouse.py deleted file mode 100644 index c13a50555..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/spacemouse.py +++ /dev/null @@ -1,108 +0,0 @@ -from spnav import spnav_open, spnav_poll_event, spnav_close, SpnavMotionEvent, SpnavButtonEvent -from threading import Thread, Event -from collections import defaultdict -import numpy as np -import time - - -class Spacemouse(Thread): - def __init__(self, max_value=500, deadzone=(0,0,0,0,0,0), dtype=np.float32): - """ - Continuously listen to 3D connection space naviagtor events - and update the latest state. - - max_value: {300, 500} 300 for wired version and 500 for wireless - deadzone: [0,1], number or tuple, axis with value lower than this value will stay at 0 - - front - z - ^ _ - | (O) space mouse - | - *----->x right - y - """ - if np.issubdtype(type(deadzone), np.number): - deadzone = np.full(6, fill_value=deadzone, dtype=dtype) - else: - deadzone = np.array(deadzone, dtype=dtype) - assert (deadzone >= 0).all() - - super().__init__() - self.stop_event = Event() - self.max_value = max_value - self.dtype = dtype - self.deadzone = deadzone - self.motion_event = SpnavMotionEvent([0,0,0], [0,0,0], 0) - self.button_state = defaultdict(lambda: False) - self.tx_zup_spnav = np.array([ - [0,0,-1], - [1,0,0], - [0,1,0] - ], dtype=dtype) - - def get_motion_state(self): - me = self.motion_event - state = np.array(me.translation + me.rotation, - dtype=self.dtype) / self.max_value - is_dead = (-self.deadzone < state) & (state < self.deadzone) - state[is_dead] = 0 - return state - - def get_motion_state_transformed(self): - """ - Return in right-handed coordinate - z - *------>y right - | _ - | (O) space mouse - v - x - back - - """ - state = self.get_motion_state() - tf_state = np.zeros_like(state) - tf_state[:3] = self.tx_zup_spnav @ state[:3] - tf_state[3:] = self.tx_zup_spnav @ state[3:] - return tf_state - - def is_button_pressed(self, button_id): - return self.button_state[button_id] - - def stop(self): - self.stop_event.set() - self.join() - - def __enter__(self): - self.start() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() - - def run(self): - spnav_open() - try: - while not self.stop_event.is_set(): - event = spnav_poll_event() - if isinstance(event, SpnavMotionEvent): - self.motion_event = event - elif isinstance(event, SpnavButtonEvent): - self.button_state[event.bnum] = event.press - else: - time.sleep(1/200) - finally: - spnav_close() - - -def test(): - with Spacemouse(deadzone=0.3) as sm: - for i in range(2000): - # print(sm.get_motion_state()) - print(sm.get_motion_state_transformed()) - print(sm.is_button_pressed(0)) - time.sleep(1/100) - -if __name__ == '__main__': - test() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/spacemouse_shared_memory.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/spacemouse_shared_memory.py deleted file mode 100644 index 06102fdb2..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/spacemouse_shared_memory.py +++ /dev/null @@ -1,160 +0,0 @@ -import multiprocessing as mp -import numpy as np -import time -from spnav import spnav_open, spnav_poll_event, spnav_close, SpnavMotionEvent, SpnavButtonEvent -from diffusion_policy.shared_memory.shared_memory_ring_buffer import SharedMemoryRingBuffer - -class Spacemouse(mp.Process): - def __init__(self, - shm_manager, - get_max_k=30, - frequency=200, - max_value=500, - deadzone=(0,0,0,0,0,0), - dtype=np.float32, - n_buttons=2, - ): - """ - Continuously listen to 3D connection space naviagtor events - and update the latest state. - - max_value: {300, 500} 300 for wired version and 500 for wireless - deadzone: [0,1], number or tuple, axis with value lower than this value will stay at 0 - - front - z - ^ _ - | (O) space mouse - | - *----->x right - y - """ - super().__init__() - if np.issubdtype(type(deadzone), np.number): - deadzone = np.full(6, fill_value=deadzone, dtype=dtype) - else: - deadzone = np.array(deadzone, dtype=dtype) - assert (deadzone >= 0).all() - - # copied variables - self.frequency = frequency - self.max_value = max_value - self.dtype = dtype - self.deadzone = deadzone - self.n_buttons = n_buttons - # self.motion_event = SpnavMotionEvent([0,0,0], [0,0,0], 0) - # self.button_state = defaultdict(lambda: False) - self.tx_zup_spnav = np.array([ - [0,0,-1], - [1,0,0], - [0,1,0] - ], dtype=dtype) - - example = { - # 3 translation, 3 rotation, 1 period - 'motion_event': np.zeros((7,), dtype=np.int64), - # left and right button - 'button_state': np.zeros((n_buttons,), dtype=bool), - 'receive_timestamp': time.time() - } - ring_buffer = SharedMemoryRingBuffer.create_from_examples( - shm_manager=shm_manager, - examples=example, - get_max_k=get_max_k, - get_time_budget=0.2, - put_desired_frequency=frequency - ) - - # shared variables - self.ready_event = mp.Event() - self.stop_event = mp.Event() - self.ring_buffer = ring_buffer - - # ======= get state APIs ========== - - def get_motion_state(self): - state = self.ring_buffer.get() - state = np.array(state['motion_event'][:6], - dtype=self.dtype) / self.max_value - is_dead = (-self.deadzone < state) & (state < self.deadzone) - state[is_dead] = 0 - return state - - def get_motion_state_transformed(self): - """ - Return in right-handed coordinate - z - *------>y right - | _ - | (O) space mouse - v - x - back - - """ - state = self.get_motion_state() - tf_state = np.zeros_like(state) - tf_state[:3] = self.tx_zup_spnav @ state[:3] - tf_state[3:] = self.tx_zup_spnav @ state[3:] - return tf_state - - def get_button_state(self): - state = self.ring_buffer.get() - return state['button_state'] - - def is_button_pressed(self, button_id): - return self.get_button_state()[button_id] - - #========== start stop API =========== - - def start(self, wait=True): - super().start() - if wait: - self.ready_event.wait() - - def stop(self, wait=True): - self.stop_event.set() - if wait: - self.join() - - def __enter__(self): - self.start() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() - - # ========= main loop ========== - def run(self): - spnav_open() - try: - motion_event = np.zeros((7,), dtype=np.int64) - button_state = np.zeros((self.n_buttons,), dtype=bool) - # send one message immediately so client can start reading - self.ring_buffer.put({ - 'motion_event': motion_event, - 'button_state': button_state, - 'receive_timestamp': time.time() - }) - self.ready_event.set() - - while not self.stop_event.is_set(): - event = spnav_poll_event() - receive_timestamp = time.time() - if isinstance(event, SpnavMotionEvent): - motion_event[:3] = event.translation - motion_event[3:6] = event.rotation - motion_event[6] = event.period - elif isinstance(event, SpnavButtonEvent): - button_state[event.bnum] = event.press - else: - # finish integrating this round of events - # before sending over - self.ring_buffer.put({ - 'motion_event': motion_event, - 'button_state': button_state, - 'receive_timestamp': receive_timestamp - }) - time.sleep(1/self.frequency) - finally: - spnav_close() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/video_recorder.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/video_recorder.py deleted file mode 100644 index 434605f82..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/real_world/video_recorder.py +++ /dev/null @@ -1,161 +0,0 @@ -from typing import Optional, Callable, Generator -import numpy as np -import av -from diffusion_policy.common.timestamp_accumulator import get_accumulate_timestamp_idxs - -def read_video( - video_path: str, dt: float, - video_start_time: float=0.0, - start_time: float=0.0, - img_transform: Optional[Callable[[np.ndarray], np.ndarray]]=None, - thread_type: str="AUTO", - thread_count: int=0, - max_pad_frames: int=10 - ) -> Generator[np.ndarray, None, None]: - frame = None - with av.open(video_path) as container: - stream = container.streams.video[0] - stream.thread_type = thread_type - stream.thread_count = thread_count - next_global_idx = 0 - for frame_idx, frame in enumerate(container.decode(stream)): - # The presentation time in seconds for this frame. - since_start = frame.time - frame_time = video_start_time + since_start - local_idxs, global_idxs, next_global_idx \ - = get_accumulate_timestamp_idxs( - # only one timestamp - timestamps=[frame_time], - start_time=start_time, - dt=dt, - next_global_idx=next_global_idx - ) - if len(global_idxs) > 0: - array = frame.to_ndarray(format='rgb24') - img = array - if img_transform is not None: - img = img_transform(array) - for global_idx in global_idxs: - yield img - # repeat last frame max_pad_frames times - array = frame.to_ndarray(format='rgb24') - img = array - if img_transform is not None: - img = img_transform(array) - for i in range(max_pad_frames): - yield img - -class VideoRecorder: - def __init__(self, - fps, - codec, - input_pix_fmt, - # options for codec - **kwargs - ): - """ - input_pix_fmt: rgb24, bgr24 see https://github.com/PyAV-Org/PyAV/blob/bc4eedd5fc474e0f25b22102b2771fe5a42bb1c7/av/video/frame.pyx#L352 - """ - - self.fps = fps - self.codec = codec - self.input_pix_fmt = input_pix_fmt - self.kwargs = kwargs - # runtime set - self._reset_state() - - def _reset_state(self): - self.container = None - self.stream = None - self.shape = None - self.dtype = None - self.start_time = None - self.next_global_idx = 0 - - @classmethod - def create_h264(cls, - fps, - codec='h264', - input_pix_fmt='rgb24', - output_pix_fmt='yuv420p', - crf=18, - profile='high', - **kwargs - ): - obj = cls( - fps=fps, - codec=codec, - input_pix_fmt=input_pix_fmt, - pix_fmt=output_pix_fmt, - options={ - 'crf': str(crf), - 'profile': profile - }, - **kwargs - ) - return obj - - - def __del__(self): - self.stop() - - def is_ready(self): - return self.stream is not None - - def start(self, file_path, start_time=None): - if self.is_ready(): - # if still recording, stop first and start anew. - self.stop() - - self.container = av.open(file_path, mode='w') - self.stream = self.container.add_stream(self.codec, rate=self.fps) - codec_context = self.stream.codec_context - for k, v in self.kwargs.items(): - setattr(codec_context, k, v) - self.start_time = start_time - - def write_frame(self, img: np.ndarray, frame_time=None): - if not self.is_ready(): - raise RuntimeError('Must run start() before writing!') - - n_repeats = 1 - if self.start_time is not None: - local_idxs, global_idxs, self.next_global_idx \ - = get_accumulate_timestamp_idxs( - # only one timestamp - timestamps=[frame_time], - start_time=self.start_time, - dt=1/self.fps, - next_global_idx=self.next_global_idx - ) - # number of appearance means repeats - n_repeats = len(local_idxs) - - if self.shape is None: - self.shape = img.shape - self.dtype = img.dtype - h,w,c = img.shape - self.stream.width = w - self.stream.height = h - assert img.shape == self.shape - assert img.dtype == self.dtype - - frame = av.VideoFrame.from_ndarray( - img, format=self.input_pix_fmt) - for i in range(n_repeats): - for packet in self.stream.encode(frame): - self.container.mux(packet) - - def stop(self): - if not self.is_ready(): - return - - # Flush stream - for packet in self.stream.encode(): - self.container.mux(packet) - - # Close the file - self.container.close() - - # reset runtime parameters - self._reset_state() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/bet_blockpush_conversion.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/bet_blockpush_conversion.py deleted file mode 100644 index c23b48762..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/bet_blockpush_conversion.py +++ /dev/null @@ -1,46 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - - -import os -import click -import pathlib -import numpy as np -from diffusion_policy.common.replay_buffer import ReplayBuffer - -@click.command() -@click.option('-i', '--input', required=True, help='input dir contains npy files') -@click.option('-o', '--output', required=True, help='output zarr path') -@click.option('--abs_action', is_flag=True, default=False) -def main(input, output, abs_action): - data_directory = pathlib.Path(input) - observations = np.load( - data_directory / "multimodal_push_observations.npy" - ) - actions = np.load(data_directory / "multimodal_push_actions.npy") - masks = np.load(data_directory / "multimodal_push_masks.npy") - - buffer = ReplayBuffer.create_empty_numpy() - for i in range(len(masks)): - eps_len = int(masks[i].sum()) - obs = observations[i,:eps_len].astype(np.float32) - action = actions[i,:eps_len].astype(np.float32) - if abs_action: - prev_eef_target = obs[:,8:10] - next_eef_target = prev_eef_target + action - action = next_eef_target - data = { - 'obs': obs, - 'action': action - } - buffer.add_episode(data) - - buffer.save_to_path(zarr_path=output, chunk_length=-1) - -if __name__ == '__main__': - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/blockpush_abs_conversion.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/blockpush_abs_conversion.py deleted file mode 100644 index 6c11e4245..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/blockpush_abs_conversion.py +++ /dev/null @@ -1,29 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - -import os -import click -import pathlib -from diffusion_policy.common.replay_buffer import ReplayBuffer - - -@click.command() -@click.option('-i', '--input', required=True) -@click.option('-o', '--output', required=True) -@click.option('-t', '--target_eef_idx', default=8, type=int) -def main(input, output, target_eef_idx): - buffer = ReplayBuffer.copy_from_path(input) - obs = buffer['obs'] - action = buffer['action'] - prev_eef_target = obs[:,target_eef_idx:target_eef_idx+action.shape[1]] - next_eef_target = prev_eef_target + action - action[:] = next_eef_target - buffer.save_to_path(zarr_path=output, chunk_length=-1) - -if __name__ == '__main__': - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/episode_lengths.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/episode_lengths.py deleted file mode 100644 index 91b2ddaae..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/episode_lengths.py +++ /dev/null @@ -1,29 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - -import click -import numpy as np -import json -from diffusion_policy.common.replay_buffer import ReplayBuffer - -@click.command() -@click.option('--input', '-i', required=True) -@click.option('--dt', default=0.1, type=float) -def main(input, dt): - buffer = ReplayBuffer.create_from_path(input) - lengths = buffer.episode_lengths - durations = lengths * dt - result = { - 'duration/mean': np.mean(durations) - } - - text = json.dumps(result, indent=2) - print(text) - -if __name__ == '__main__': - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/generate_bet_blockpush.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/generate_bet_blockpush.py deleted file mode 100644 index e27122797..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/generate_bet_blockpush.py +++ /dev/null @@ -1,64 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - - -import os -import click -import pathlib -import numpy as np -from tqdm import tqdm -from diffusion_policy.common.replay_buffer import ReplayBuffer -from tf_agents.environments.wrappers import TimeLimit -from tf_agents.environments.gym_wrapper import GymWrapper -from tf_agents.trajectories.time_step import StepType -from diffusion_policy.env.block_pushing.block_pushing_multimodal import BlockPushMultimodal -from diffusion_policy.env.block_pushing.block_pushing import BlockPush -from diffusion_policy.env.block_pushing.oracles.multimodal_push_oracle import MultimodalOrientedPushOracle - -@click.command() -@click.option('-o', '--output', required=True) -@click.option('-n', '--n_episodes', default=1000) -@click.option('-c', '--chunk_length', default=-1) -def main(output, n_episodes, chunk_length): - - buffer = ReplayBuffer.create_empty_numpy() - env = TimeLimit(GymWrapper(BlockPushMultimodal()), duration=350) - for i in tqdm(range(n_episodes)): - print(i) - obs_history = list() - action_history = list() - - env.seed(i) - policy = MultimodalOrientedPushOracle(env) - time_step = env.reset() - policy_state = policy.get_initial_state(1) - while True: - action_step = policy.action(time_step, policy_state) - obs = np.concatenate(list(time_step.observation.values()), axis=-1) - action = action_step.action - obs_history.append(obs) - action_history.append(action) - - if time_step.step_type == 2: - break - - # state = env.wrapped_env().gym.get_pybullet_state() - time_step = env.step(action) - obs_history = np.array(obs_history) - action_history = np.array(action_history) - - episode = { - 'obs': obs_history, - 'action': action_history - } - buffer.add_episode(episode) - - buffer.save_to_path(output, chunk_length=chunk_length) - -if __name__ == '__main__': - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/real_dataset_conversion.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/real_dataset_conversion.py deleted file mode 100644 index 3e4b0f9c0..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/real_dataset_conversion.py +++ /dev/null @@ -1,60 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - -import os -import click -import pathlib -import zarr -import cv2 -import threadpoolctl -from diffusion_policy.real_world.real_data_conversion import real_data_to_replay_buffer - -@click.command() -@click.option('--input', '-i', required=True) -@click.option('--output', '-o', default=None) -@click.option('--resolution', '-r', default='640x480') -@click.option('--n_decoding_threads', '-nd', default=-1, type=int) -@click.option('--n_encoding_threads', '-ne', default=-1, type=int) -def main(input, output, resolution, n_decoding_threads, n_encoding_threads): - out_resolution = tuple(int(x) for x in resolution.split('x')) - input = pathlib.Path(os.path.expanduser(input)) - in_zarr_path = input.joinpath('replay_buffer.zarr') - in_video_dir = input.joinpath('videos') - assert in_zarr_path.is_dir() - assert in_video_dir.is_dir() - if output is None: - output = input.joinpath(resolution + '.zarr.zip') - else: - output = pathlib.Path(os.path.expanduser(output)) - - if output.exists(): - click.confirm('Output path already exists! Overrite?', abort=True) - - cv2.setNumThreads(1) - with threadpoolctl.threadpool_limits(1): - replay_buffer = real_data_to_replay_buffer( - dataset_path=str(input), - out_resolutions=out_resolution, - n_decoding_threads=n_decoding_threads, - n_encoding_threads=n_encoding_threads - ) - - print('Saving to disk') - if output.suffix == '.zip': - with zarr.ZipStore(output) as zip_store: - replay_buffer.save_to_store( - store=zip_store - ) - else: - with zarr.DirectoryStore(output) as store: - replay_buffer.save_to_store( - store=store - ) - -if __name__ == '__main__': - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/real_pusht_metrics.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/real_pusht_metrics.py deleted file mode 100644 index cdb05ef71..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/real_pusht_metrics.py +++ /dev/null @@ -1,151 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - -import os -import click -import av -import cv2 -import collections -import multiprocessing as mp -import numpy as np -from tqdm import tqdm -import threadpoolctl -from matplotlib import pyplot as plt -import json - -def get_t_mask(img, hsv_ranges=None): - if hsv_ranges is None: - hsv_ranges = [ - [0,255], - [130,216], - [150,230] - ] - hsv_img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV) - mask = np.ones(img.shape[:2], dtype=bool) - for c in range(len(hsv_ranges)): - l, h = hsv_ranges[c] - mask &= (l <= hsv_img[...,c]) - mask &= (hsv_img[...,c] <= h) - return mask - -def get_mask_metrics(target_mask, mask): - total = np.sum(target_mask) - i = np.sum(target_mask & mask) - u = np.sum(target_mask | mask) - iou = i / u - coverage = i / total - result = { - 'iou': iou, - 'coverage': coverage - } - return result - -def get_video_metrics(video_path, target_mask, use_tqdm=True): - threadpoolctl.threadpool_limits(1) - cv2.setNumThreads(1) - - metrics = collections.defaultdict(list) - with av.open(video_path) as container: - stream = container.streams.video[0] - iterator = None - if use_tqdm: - iterator = tqdm(container.decode(stream), total=stream.frames) - else: - iterator = container.decode(stream) - for frame in iterator: - img = frame.to_ndarray(format='rgb24') - mask = get_t_mask(img) - metric = get_mask_metrics( - target_mask=target_mask, mask=mask) - for k, v in metric.items(): - metrics[k].append(v) - return metrics - -def worker(x): - return get_video_metrics(*x) - -@click.command() -@click.option( - '--reference', '-r', required=True, - help="Reference video whose last frame will define goal.") -@click.option( - '--input', '-i', required=True, - help='Dataset path to evaluate.') -@click.option( - '--camera_idx', '-ci', default=0, type=int, - help="Camera index to compute metrics") -@click.option('--n_workers', '-n', default=20, type=int) -def main(reference, input, camera_idx, n_workers): - # read last frame of the reference video to get target mask - last_frame = None - with av.open(reference) as container: - stream = container.streams.video[0] - for frame in tqdm( - container.decode(stream), - total=stream.frames): - last_frame = frame - - last_img = last_frame.to_ndarray(format='rgb24') - target_mask = get_t_mask(last_img) - - # path = '/home/ubuntu/dev/diffusion_policy/data/pusht_real/eval_20230109/diffusion_hybrid_ep136/videos/4/0.mp4' - # last_frame = None - # with av.open(path) as container: - # stream = container.streams.video[0] - # for frame in tqdm( - # container.decode(stream), - # total=stream.frames): - # last_frame = frame - # img = last_frame.to_ndarray(format='rgb24') - # mask = get_t_mask(img) - - # get metrics for each episode - episode_video_path_map = dict() - input_dir = pathlib.Path(input) - input_video_dir = input_dir.joinpath('videos') - for vid_dir in input_video_dir.glob("*/"): - episode_idx = int(vid_dir.stem) - video_path = vid_dir.joinpath(f'{camera_idx}.mp4') - if video_path.exists(): - episode_video_path_map[episode_idx] = str(video_path.absolute()) - - episode_idxs = sorted(episode_video_path_map.keys()) - print(f"Found video for following episodes: {episode_idxs}") - - # run - with mp.Pool(n_workers) as pool: - args = list() - for idx in episode_idxs: - args.append((episode_video_path_map[idx], target_mask)) - results = pool.map(worker, args) - episode_metric_map = dict() - for idx, result in zip(episode_idxs, results): - episode_metric_map[idx] = result - - # aggregate metrics - agg_map = collections.defaultdict(list) - for idx, metric in episode_metric_map.items(): - for key, value in metric.items(): - agg_map['max/'+key].append(np.max(value)) - agg_map['last/'+key].append(value[-1]) - - final_metric = dict() - for key, value in agg_map.items(): - final_metric[key] = np.mean(value) - - # save metrics - print('Saving metrics!') - with input_dir.joinpath('metrics_agg.json').open('w') as f: - json.dump(final_metric, f, sort_keys=True, indent=2) - - with input_dir.joinpath('metrics_raw.json').open('w') as f: - json.dump(episode_metric_map, f, sort_keys=True, indent=2) - print('Done!') - -if __name__ == '__main__': - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/real_pusht_successrate.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/real_pusht_successrate.py deleted file mode 100644 index a5900faad..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/real_pusht_successrate.py +++ /dev/null @@ -1,69 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - -import os -import click -import collections -import numpy as np -from tqdm import tqdm -import json - -@click.command() -@click.option( - '--reference', '-r', required=True, - help='Reference metrics_raw.json from demonstration dataset.' -) -@click.option( - '--input', '-i', required=True, - help='Data search path' -) -def main(reference, input): - # compute the min last metric for demo metrics - demo_metrics = json.load(open(reference, 'r')) - demo_min_metrics = collections.defaultdict(lambda:float('inf')) - for episode_idx, metrics in demo_metrics.items(): - for key, value in metrics.items(): - last_value = value[-1] - demo_min_metrics[key] = min(demo_min_metrics[key], last_value) - print(demo_min_metrics) - - # find all metric - name = 'metrics_raw.json' - search_dir = pathlib.Path(input) - success_rate_map = dict() - for json_path in search_dir.glob('**/'+name): - rel_path = json_path.relative_to(search_dir) - rel_name = str(rel_path.parent) - this_metrics = json.load(json_path.open('r')) - metric_success_idxs = collections.defaultdict(list) - metric_failure_idxs = collections.defaultdict(list) - for episode_idx, metrics in this_metrics.items(): - for key, value in metrics.items(): - last_value = value[-1] - # print(episode_idx, key, last_value) - demo_min = demo_min_metrics[key] - if last_value >= demo_min: - # success - metric_success_idxs[key].append(episode_idx) - else: - metric_failure_idxs[key].append(episode_idx) - # in case of no success - _ = metric_success_idxs[key] - _ = metric_failure_idxs[key] - metric_success_rate = dict() - n_episodes = len(this_metrics) - for key, value in metric_success_idxs.items(): - metric_success_rate[key] = len(value) / n_episodes - # metric_success_rate['failured_idxs'] = metric_failure_idxs - success_rate_map[rel_name] = metric_success_rate - - text = json.dumps(success_rate_map, indent=2) - print(text) - -if __name__ == '__main__': - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/robomimic_dataset_action_comparison.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/robomimic_dataset_action_comparison.py deleted file mode 100644 index 43114de4b..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/robomimic_dataset_action_comparison.py +++ /dev/null @@ -1,51 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - -import os -import click -import pathlib -import h5py -import numpy as np -from tqdm import tqdm -from scipy.spatial.transform import Rotation - -def read_all_actions(hdf5_file, metric_skip_steps=1): - n_demos = len(hdf5_file['data']) - all_actions = list() - for i in tqdm(range(n_demos)): - actions = hdf5_file[f'data/demo_{i}/actions'][:] - all_actions.append(actions[metric_skip_steps:]) - all_actions = np.concatenate(all_actions, axis=0) - return all_actions - - -@click.command() -@click.option('-i', '--input', required=True, help='input hdf5 path') -@click.option('-o', '--output', required=True, help='output hdf5 path. Parent directory must exist') -def main(input, output): - # process inputs - input = pathlib.Path(input).expanduser() - assert input.is_file() - output = pathlib.Path(output).expanduser() - assert output.is_file() - - input_file = h5py.File(str(input), 'r') - output_file = h5py.File(str(output), 'r') - - input_all_actions = read_all_actions(input_file) - output_all_actions = read_all_actions(output_file) - pos_dist = np.linalg.norm(input_all_actions[:,:3] - output_all_actions[:,:3], axis=-1) - rot_dist = (Rotation.from_rotvec(input_all_actions[:,3:6] - ) * Rotation.from_rotvec(output_all_actions[:,3:6]).inv() - ).magnitude() - - print(f'max pos dist: {pos_dist.max()}') - print(f'max rot dist: {rot_dist.max()}') - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/robomimic_dataset_conversion.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/robomimic_dataset_conversion.py deleted file mode 100644 index 5496d55cd..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/scripts/robomimic_dataset_conversion.py +++ /dev/null @@ -1,103 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - -import multiprocessing -import os -import shutil -import click -import pathlib -import h5py -from tqdm import tqdm -import collections -import pickle -from diffusion_policy.common.robomimic_util import RobomimicAbsoluteActionConverter - -def worker(x): - path, idx, do_eval = x - converter = RobomimicAbsoluteActionConverter(path) - if do_eval: - abs_actions, info = converter.convert_and_eval_idx(idx) - else: - abs_actions = converter.convert_idx(idx) - info = dict() - return abs_actions, info - -@click.command() -@click.option('-i', '--input', required=True, help='input hdf5 path') -@click.option('-o', '--output', required=True, help='output hdf5 path. Parent directory must exist') -@click.option('-e', '--eval_dir', default=None, help='directory to output evaluation metrics') -@click.option('-n', '--num_workers', default=None, type=int) -def main(input, output, eval_dir, num_workers): - # process inputs - input = pathlib.Path(input).expanduser() - assert input.is_file() - output = pathlib.Path(output).expanduser() - assert output.parent.is_dir() - assert not output.is_dir() - - 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(input) - - # run - with multiprocessing.Pool(num_workers) as pool: - results = pool.map(worker, [(input, i, do_eval) for i in range(len(converter))]) - - # save output - print('Copying hdf5') - shutil.copy(str(input), str(output)) - - # modify action - with h5py.File(output, 'r+') as out_file: - for i in tqdm(range(len(converter)), desc="Writing to output"): - abs_actions, info = results[i] - demo = out_file[f'data/demo_{i}'] - demo['actions'][:] = 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__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_memory_queue.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_memory_queue.py deleted file mode 100644 index a4f316386..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_memory_queue.py +++ /dev/null @@ -1,187 +0,0 @@ -from typing import Dict, List, Union -import numbers -from queue import (Empty, Full) -from multiprocessing.managers import SharedMemoryManager -import numpy as np -from diffusion_policy.shared_memory.shared_memory_util import ArraySpec, SharedAtomicCounter -from diffusion_policy.shared_memory.shared_ndarray import SharedNDArray - - -class SharedMemoryQueue: - """ - A Lock-Free FIFO Shared Memory Data Structure. - Stores a sequence of dict of numpy arrays. - """ - - def __init__(self, - shm_manager: SharedMemoryManager, - array_specs: List[ArraySpec], - buffer_size: int - ): - - # create atomic counter - write_counter = SharedAtomicCounter(shm_manager) - read_counter = SharedAtomicCounter(shm_manager) - - # allocate shared memory - shared_arrays = dict() - for spec in array_specs: - key = spec.name - assert key not in shared_arrays - array = SharedNDArray.create_from_shape( - mem_mgr=shm_manager, - shape=(buffer_size,) + tuple(spec.shape), - dtype=spec.dtype) - shared_arrays[key] = array - - self.buffer_size = buffer_size - self.array_specs = array_specs - self.write_counter = write_counter - self.read_counter = read_counter - self.shared_arrays = shared_arrays - - @classmethod - def create_from_examples(cls, - shm_manager: SharedMemoryManager, - examples: Dict[str, Union[np.ndarray, numbers.Number]], - buffer_size: int - ): - specs = list() - for key, value in examples.items(): - shape = None - dtype = None - if isinstance(value, np.ndarray): - shape = value.shape - dtype = value.dtype - assert dtype != np.dtype('O') - elif isinstance(value, numbers.Number): - shape = tuple() - dtype = np.dtype(type(value)) - else: - raise TypeError(f'Unsupported type {type(value)}') - - spec = ArraySpec( - name=key, - shape=shape, - dtype=dtype - ) - specs.append(spec) - - obj = cls( - shm_manager=shm_manager, - array_specs=specs, - buffer_size=buffer_size - ) - return obj - - def qsize(self): - read_count = self.read_counter.load() - write_count = self.write_counter.load() - n_data = write_count - read_count - return n_data - - def empty(self): - n_data = self.qsize() - return n_data <= 0 - - def clear(self): - self.read_counter.store(self.write_counter.load()) - - def put(self, data: Dict[str, Union[np.ndarray, numbers.Number]]): - read_count = self.read_counter.load() - write_count = self.write_counter.load() - n_data = write_count - read_count - if n_data >= self.buffer_size: - raise Full() - - next_idx = write_count % self.buffer_size - - # write to shared memory - for key, value in data.items(): - arr: np.ndarray - arr = self.shared_arrays[key].get() - if isinstance(value, np.ndarray): - arr[next_idx] = value - else: - arr[next_idx] = np.array(value, dtype=arr.dtype) - - # update idx - self.write_counter.add(1) - - def get(self, out=None) -> Dict[str, np.ndarray]: - write_count = self.write_counter.load() - read_count = self.read_counter.load() - n_data = write_count - read_count - if n_data <= 0: - raise Empty() - - if out is None: - out = self._allocate_empty() - - next_idx = read_count % self.buffer_size - for key, value in self.shared_arrays.items(): - arr = value.get() - np.copyto(out[key], arr[next_idx]) - - # update idx - self.read_counter.add(1) - return out - - def get_k(self, k, out=None) -> Dict[str, np.ndarray]: - write_count = self.write_counter.load() - read_count = self.read_counter.load() - n_data = write_count - read_count - if n_data <= 0: - raise Empty() - assert k <= n_data - - out = self._get_k_impl(k, read_count, out=out) - self.read_counter.add(k) - return out - - def get_all(self, out=None) -> Dict[str, np.ndarray]: - write_count = self.write_counter.load() - read_count = self.read_counter.load() - n_data = write_count - read_count - if n_data <= 0: - raise Empty() - - out = self._get_k_impl(n_data, read_count, out=out) - self.read_counter.add(n_data) - return out - - def _get_k_impl(self, k, read_count, out=None) -> Dict[str, np.ndarray]: - if out is None: - out = self._allocate_empty(k) - - curr_idx = read_count % self.buffer_size - for key, value in self.shared_arrays.items(): - arr = value.get() - target = out[key] - - start = curr_idx - end = min(start + k, self.buffer_size) - target_start = 0 - target_end = (end - start) - target[target_start: target_end] = arr[start:end] - - remainder = k - (end - start) - if remainder > 0: - # wrap around - start = 0 - end = start + remainder - target_start = target_end - target_end = k - target[target_start: target_end] = arr[start:end] - - return out - - def _allocate_empty(self, k=None): - result = dict() - for spec in self.array_specs: - shape = spec.shape - if k is not None: - shape = (k,) + shape - result[spec.name] = np.empty( - shape=shape, dtype=spec.dtype) - return result diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_memory_ring_buffer.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_memory_ring_buffer.py deleted file mode 100644 index 23feef0b1..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_memory_ring_buffer.py +++ /dev/null @@ -1,219 +0,0 @@ -from typing import Dict, List, Union - -from queue import Empty -import numbers -import time -from multiprocessing.managers import SharedMemoryManager -import numpy as np - -from diffusion_policy.shared_memory.shared_ndarray import SharedNDArray -from diffusion_policy.shared_memory.shared_memory_util import ArraySpec, SharedAtomicCounter - -class SharedMemoryRingBuffer: - """ - A Lock-Free FILO Shared Memory Data Structure. - Stores a sequence of dict of numpy arrays. - """ - - def __init__(self, - shm_manager: SharedMemoryManager, - array_specs: List[ArraySpec], - get_max_k: int, - get_time_budget: float, - put_desired_frequency: float, - safety_margin: float=1.5 - ): - """ - shm_manager: Manages the life cycle of share memories - across processes. Remember to run .start() before passing. - array_specs: Name, shape and type of arrays for a single time step. - get_max_k: The maxmum number of items can be queried at once. - get_time_budget: The maxmum amount of time spent copying data from - shared memory to local memory. Increase this number for larger arrays. - put_desired_frequency: The maximum frequency that .put() can be called. - This influces the buffer size. - """ - - # create atomic counter - counter = SharedAtomicCounter(shm_manager) - - # compute buffer size - # At any given moment, the past get_max_k items should never - # be touched (to be read freely). Assuming the reading is reading - # these k items, which takes maximum of get_time_budget seconds, - # we need enough empty slots to make sure put_desired_frequency Hz - # of put can be sustaied. - buffer_size = int(np.ceil( - put_desired_frequency * get_time_budget - * safety_margin)) + get_max_k - - # allocate shared memory - shared_arrays = dict() - for spec in array_specs: - key = spec.name - assert key not in shared_arrays - array = SharedNDArray.create_from_shape( - mem_mgr=shm_manager, - shape=(buffer_size,) + tuple(spec.shape), - dtype=spec.dtype) - shared_arrays[key] = array - - # allocate timestamp array - timestamp_array = SharedNDArray.create_from_shape( - mem_mgr=shm_manager, - shape=(buffer_size,), - dtype=np.float64) - timestamp_array.get()[:] = -np.inf - - self.buffer_size = buffer_size - self.array_specs = array_specs - self.counter = counter - self.shared_arrays = shared_arrays - self.timestamp_array = timestamp_array - self.get_time_budget = get_time_budget - self.get_max_k = get_max_k - self.put_desired_frequency = put_desired_frequency - - - @property - def count(self): - return self.counter.load() - - @classmethod - def create_from_examples(cls, - shm_manager: SharedMemoryManager, - examples: Dict[str, Union[np.ndarray, numbers.Number]], - get_max_k: int=32, - get_time_budget: float=0.01, - put_desired_frequency: float=60 - ): - specs = list() - for key, value in examples.items(): - shape = None - dtype = None - if isinstance(value, np.ndarray): - shape = value.shape - dtype = value.dtype - assert dtype != np.dtype('O') - elif isinstance(value, numbers.Number): - shape = tuple() - dtype = np.dtype(type(value)) - else: - raise TypeError(f'Unsupported type {type(value)}') - - spec = ArraySpec( - name=key, - shape=shape, - dtype=dtype - ) - specs.append(spec) - - obj = cls( - shm_manager=shm_manager, - array_specs=specs, - get_max_k=get_max_k, - get_time_budget=get_time_budget, - put_desired_frequency=put_desired_frequency - ) - return obj - - def clear(self): - self.counter.store(0) - - def put(self, data: Dict[str, Union[np.ndarray, numbers.Number]], wait: bool=True): - count = self.counter.load() - next_idx = count % self.buffer_size - # Make sure the next self.get_max_k elements in the ring buffer have at least - # self.get_time_budget seconds untouched after written, so that - # get_last_k can safely read k elements from any count location. - # Sanity check: when get_max_k == 1, the element pointed by next_idx - # should be rewritten at minimum self.get_time_budget seconds later. - timestamp_lookahead_idx = (next_idx + self.get_max_k - 1) % self.buffer_size - old_timestamp = self.timestamp_array.get()[timestamp_lookahead_idx] - t = time.monotonic() - if (t - old_timestamp) < self.get_time_budget: - deltat = t - old_timestamp - if wait: - # sleep the remaining time to be safe - time.sleep(self.get_time_budget - deltat) - else: - # throw an error - past_iters = self.buffer_size - self.get_max_k - hz = past_iters / deltat - raise TimeoutError( - 'Put executed too fast {}items/{:.4f}s ~= {}Hz'.format( - past_iters, deltat,hz)) - - # write to shared memory - for key, value in data.items(): - arr: np.ndarray - arr = self.shared_arrays[key].get() - if isinstance(value, np.ndarray): - arr[next_idx] = value - else: - arr[next_idx] = np.array(value, dtype=arr.dtype) - - # update timestamp - self.timestamp_array.get()[next_idx] = time.monotonic() - self.counter.add(1) - - def _allocate_empty(self, k=None): - result = dict() - for spec in self.array_specs: - shape = spec.shape - if k is not None: - shape = (k,) + shape - result[spec.name] = np.empty( - shape=shape, dtype=spec.dtype) - return result - - def get(self, out=None) -> Dict[str, np.ndarray]: - if out is None: - out = self._allocate_empty() - start_time = time.monotonic() - count = self.counter.load() - curr_idx = (count - 1) % self.buffer_size - for key, value in self.shared_arrays.items(): - arr = value.get() - np.copyto(out[key], arr[curr_idx]) - end_time = time.monotonic() - dt = end_time - start_time - if dt > self.get_time_budget: - raise TimeoutError(f'Get time out {dt} vs {self.get_time_budget}') - return out - - def get_last_k(self, k:int, out=None) -> Dict[str, np.ndarray]: - assert k <= self.get_max_k - if out is None: - out = self._allocate_empty(k) - start_time = time.monotonic() - count = self.counter.load() - assert k <= count - curr_idx = (count - 1) % self.buffer_size - for key, value in self.shared_arrays.items(): - arr = value.get() - target = out[key] - - end = curr_idx + 1 - start = max(0, end - k) - target_end = k - target_start = target_end - (end - start) - target[target_start: target_end] = arr[start:end] - - remainder = k - (end - start) - if remainder > 0: - # wrap around - end = self.buffer_size - start = end - remainder - target_start = 0 - target_end = end - start - target[target_start: target_end] = arr[start:end] - end_time = time.monotonic() - dt = end_time - start_time - if dt > self.get_time_budget: - raise TimeoutError(f'Get time out {dt} vs {self.get_time_budget}') - return out - - def get_all(self) -> Dict[str, np.ndarray]: - k = min(self.count, self.get_max_k) - return self.get_last_k(k=k) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_memory_util.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_memory_util.py deleted file mode 100644 index 24ee7df35..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_memory_util.py +++ /dev/null @@ -1,39 +0,0 @@ -from typing import Tuple -from dataclasses import dataclass -import numpy as np -from multiprocessing.managers import SharedMemoryManager -from atomics import atomicview, MemoryOrder, UINT - -@dataclass -class ArraySpec: - name: str - shape: Tuple[int] - dtype: np.dtype - - -class SharedAtomicCounter: - def __init__(self, - shm_manager: SharedMemoryManager, - size :int=8 # 64bit int - ): - shm = shm_manager.SharedMemory(size=size) - self.shm = shm - self.size = size - self.store(0) # initialize - - @property - def buf(self): - return self.shm.buf[:self.size] - - def load(self) -> int: - with atomicview(buffer=self.buf, atype=UINT) as a: - value = a.load(order=MemoryOrder.ACQUIRE) - return value - - def store(self, value: int): - with atomicview(buffer=self.buf, atype=UINT) as a: - a.store(value, order=MemoryOrder.RELEASE) - - def add(self, value: int): - with atomicview(buffer=self.buf, atype=UINT) as a: - a.add(value, order=MemoryOrder.ACQ_REL) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_ndarray.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_ndarray.py deleted file mode 100644 index cbe5c9f48..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/shared_memory/shared_ndarray.py +++ /dev/null @@ -1,167 +0,0 @@ -from __future__ import annotations - -import multiprocessing -import multiprocessing.synchronize -from multiprocessing.managers import SharedMemoryManager -from multiprocessing.shared_memory import SharedMemory -from typing import Any, TYPE_CHECKING, Generic, Optional, Tuple, TypeVar, Union - -import numpy as np -import numpy.typing as npt -from diffusion_policy.common.nested_dict_util import (nested_dict_check, nested_dict_map) - - -SharedMemoryLike = Union[str, SharedMemory] # shared memory or name of shared memory -SharedT = TypeVar("SharedT", bound=np.generic) - - -class SharedNDArray(Generic[SharedT]): - """Class to keep track of and retrieve the data in a shared array - Attributes - ---------- - shm - SharedMemory object containing the data of the array - shape - Shape of the NumPy array - dtype - Type of the NumPy array. Anything that may be passed to the `dtype=` argument in `np.ndarray`. - lock - (Optional) multiprocessing.Lock to manage access to the SharedNDArray. This is only created if - lock=True is passed to the constructor, otherwise it is set to `None`. - A SharedNDArray object may be created either directly with a preallocated shared memory object plus the - dtype and shape of the numpy array it represents: - >>> from multiprocessing.shared_memory import SharedMemory - >>> import numpy as np - >>> from shared_ndarray2 import SharedNDArray - >>> x = np.array([1, 2, 3]) - >>> shm = SharedMemory(name="x", create=True, size=x.nbytes) - >>> arr = SharedNDArray(shm, x.shape, x.dtype) - >>> arr[:] = x[:] # copy x into the array - >>> print(arr[:]) - [1 2 3] - >>> shm.close() - >>> shm.unlink() - Or using a SharedMemoryManager either from an existing array or from arbitrary shape and nbytes: - >>> from multiprocessing.managers import SharedMemoryManager - >>> mem_mgr = SharedMemoryManager() - >>> mem_mgr.start() # Better yet, use SharedMemoryManager context manager - >>> arr = SharedNDArray.from_shape(mem_mgr, x.shape, x.dtype) - >>> arr[:] = x[:] # copy x into the array - >>> print(arr[:]) - [1 2 3] - >>> # -or in one step- - >>> arr = SharedNDArray.from_array(mem_mgr, x) - >>> print(arr[:]) - [1 2 3] - `SharedNDArray` does not subclass numpy.ndarray but rather generates an ndarray on-the-fly in get(), - which is used in __getitem__ and __setitem__. Thus to access the data and/or use any ndarray methods - get() or __getitem__ or __setitem__ must be used - >>> arr.max() # ERROR: SharedNDArray has no `max` method. - Traceback (most recent call last): - .... - AttributeError: SharedNDArray object has no attribute 'max'. To access NumPy ndarray object use .get() method. - >>> arr.get().max() # (or arr[:].max()) OK: This gets an ndarray on which we can operate - 3 - >>> y = np.zeros(3) - >>> y[:] = arr # ERROR: Cannot broadcast-assign a SharedNDArray to ndarray `y` - Traceback (most recent call last): - ... - ValueError: setting an array element with a sequence. - >>> y[:] = arr[:] # OK: This gets an ndarray that can be copied element-wise to `y` - >>> mem_mgr.shutdown() - """ - - shm: SharedMemory - # shape: Tuple[int, ...] # is a property - dtype: np.dtype - lock: Optional[multiprocessing.synchronize.Lock] - - def __init__( - self, shm: SharedMemoryLike, shape: Tuple[int, ...], dtype: npt.DTypeLike): - """Initialize a SharedNDArray object from existing shared memory, object shape, and dtype. - To initialize a SharedNDArray object from a memory manager and data or shape, use the `from_array() - or `from_shape()` classmethods. - Parameters - ---------- - shm - `multiprocessing.shared_memory.SharedMemory` object or name for connecting to an existing block - of shared memory (using SharedMemory constructor) - shape - Shape of the NumPy array to be represented in the shared memory - dtype - Data type for the NumPy array to be represented in shared memory. Any valid argument for - `np.dtype` may be used as it will be converted to an actual `dtype` object. - lock : bool, optional - If True, create a multiprocessing.Lock object accessible with the `.lock` attribute, by default - False. If passing the `SharedNDArray` as an argument to a `multiprocessing.Pool` function this - should not be used -- see this comment to a Stack Overflow question about `multiprocessing.Lock`: - https://stackoverflow.com/questions/25557686/python-sharing-a-lock-between-processes#comment72803059_25558333 - Raises - ------ - ValueError - The SharedMemory size (number of bytes) does not match the product of the shape and dtype - itemsize. - """ - if isinstance(shm, str): - shm = SharedMemory(name=shm, create=False) - dtype = np.dtype(dtype) # Try to convert to dtype - assert shm.size >= (dtype.itemsize * np.prod(shape)) - self.shm = shm - self.dtype = dtype - self._shape: Tuple[int, ...] = shape - - def __repr__(self): - # Like numpy's ndarray repr - cls_name = self.__class__.__name__ - nspaces = len(cls_name) + 1 - array_repr = str(self.get()) - array_repr = array_repr.replace("\n", "\n" + " " * nspaces) - return f"{cls_name}({array_repr}, dtype={self.dtype})" - - @classmethod - def create_from_array( - cls, mem_mgr: SharedMemoryManager, arr: npt.NDArray[SharedT] - ) -> SharedNDArray[SharedT]: - """Create a SharedNDArray from a SharedMemoryManager and an existing numpy array. - Parameters - ---------- - mem_mgr - Running `multiprocessing.managers.SharedMemoryManager` instance from which to create the - SharedMemory for the SharedNDArray - arr - NumPy `ndarray` object to copy into the created SharedNDArray upon initialization. - """ - # Simply use from_shape() to create the SharedNDArray and copy the data into it. - shared_arr = cls.create_from_shape(mem_mgr, arr.shape, arr.dtype) - shared_arr.get()[:] = arr[:] - return shared_arr - - @classmethod - def create_from_shape( - cls, mem_mgr: SharedMemoryManager, shape: Tuple, dtype: npt.DTypeLike) -> SharedNDArray: - """Create a SharedNDArray directly from a SharedMemoryManager - Parameters - ---------- - mem_mgr - SharedMemoryManager instance that has been started - shape - Shape of the array - dtype - Data type for the NumPy array to be represented in shared memory. Any valid argument for - `np.dtype` may be used as it will be converted to an actual `dtype` object. - """ - dtype = np.dtype(dtype) # Convert to dtype if possible - shm = mem_mgr.SharedMemory(np.prod(shape) * dtype.itemsize) - return cls(shm=shm, shape=shape, dtype=dtype) - - @property - def shape(self) -> Tuple[int, ...]: - return self._shape - - - def get(self) -> npt.NDArray[SharedT]: - """Get a numpy array with access to the shared memory""" - return np.ndarray(self.shape, dtype=self.dtype, buffer=self.shm.buf) - - def __del__(self): - self.shm.close() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/base_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/base_workspace.py deleted file mode 100644 index 1352404a3..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/base_workspace.py +++ /dev/null @@ -1,145 +0,0 @@ -from typing import Optional -import os -import pathlib -import hydra -import copy -from hydra.core.hydra_config import HydraConfig -from omegaconf import OmegaConf -import dill -import torch -import threading - - -class BaseWorkspace: - include_keys = tuple() - exclude_keys = tuple() - - def __init__(self, cfg: OmegaConf, output_dir: Optional[str]=None): - self.cfg = cfg - self._output_dir = output_dir - self._saving_thread = None - - @property - def output_dir(self): - output_dir = self._output_dir - if output_dir is None: - output_dir = HydraConfig.get().runtime.output_dir - return output_dir - - def run(self): - """ - Create any resource shouldn't be serialized as local variables - """ - pass - - def save_checkpoint(self, path=None, tag='latest', - exclude_keys=None, - include_keys=None, - use_thread=True): - if path is None: - path = pathlib.Path(self.output_dir).joinpath('checkpoints', f'{tag}.ckpt') - else: - path = pathlib.Path(path) - if exclude_keys is None: - exclude_keys = tuple(self.exclude_keys) - if include_keys is None: - include_keys = tuple(self.include_keys) + ('_output_dir',) - - path.parent.mkdir(parents=False, exist_ok=True) - payload = { - 'cfg': self.cfg, - 'state_dicts': dict(), - 'pickles': dict() - } - - for key, value in self.__dict__.items(): - if hasattr(value, 'state_dict') and hasattr(value, 'load_state_dict'): - # modules, optimizers and samplers etc - if key not in exclude_keys: - if use_thread: - payload['state_dicts'][key] = _copy_to_cpu(value.state_dict()) - else: - payload['state_dicts'][key] = value.state_dict() - elif key in include_keys: - payload['pickles'][key] = dill.dumps(value) - if use_thread: - self._saving_thread = threading.Thread( - target=lambda : torch.save(payload, path.open('wb'), pickle_module=dill)) - self._saving_thread.start() - else: - torch.save(payload, path.open('wb'), pickle_module=dill) - return str(path.absolute()) - - def get_checkpoint_path(self, tag='latest'): - return pathlib.Path(self.output_dir).joinpath('checkpoints', f'{tag}.ckpt') - - def load_payload(self, payload, exclude_keys=None, include_keys=None, **kwargs): - if exclude_keys is None: - exclude_keys = tuple() - if include_keys is None: - include_keys = payload['pickles'].keys() - - for key, value in payload['state_dicts'].items(): - if key not in exclude_keys: - self.__dict__[key].load_state_dict(value, **kwargs) - for key in include_keys: - if key in payload['pickles']: - self.__dict__[key] = dill.loads(payload['pickles'][key]) - - def load_checkpoint(self, path=None, tag='latest', - exclude_keys=None, - include_keys=None, - **kwargs): - if path is None: - path = self.get_checkpoint_path(tag=tag) - else: - path = pathlib.Path(path) - payload = torch.load(path.open('rb'), pickle_module=dill, **kwargs) - self.load_payload(payload, - exclude_keys=exclude_keys, - include_keys=include_keys) - return payload - - @classmethod - def create_from_checkpoint(cls, path, - exclude_keys=None, - include_keys=None, - **kwargs): - payload = torch.load(open(path, 'rb'), pickle_module=dill) - instance = cls(payload['cfg']) - instance.load_payload( - payload=payload, - exclude_keys=exclude_keys, - include_keys=include_keys, - **kwargs) - return instance - - def save_snapshot(self, tag='latest'): - """ - Quick loading and saving for reserach, saves full state of the workspace. - - However, loading a snapshot assumes the code stays exactly the same. - Use save_checkpoint for long-term storage. - """ - path = pathlib.Path(self.output_dir).joinpath('snapshots', f'{tag}.pkl') - path.parent.mkdir(parents=False, exist_ok=True) - torch.save(self, path.open('wb'), pickle_module=dill) - return str(path.absolute()) - - @classmethod - def create_from_snapshot(cls, path): - return torch.load(open(path, 'rb'), pickle_module=dill) - - -def _copy_to_cpu(x): - if isinstance(x, torch.Tensor): - return x.detach().to('cpu') - elif isinstance(x, dict): - result = dict() - for k, v in x.items(): - result[k] = _copy_to_cpu(v) - return result - elif isinstance(x, list): - return [_copy_to_cpu(k) for k in x] - else: - return copy.deepcopy(x) diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_bet_lowdim_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_bet_lowdim_workspace.py deleted file mode 100644 index c8909e1b0..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_bet_lowdim_workspace.py +++ /dev/null @@ -1,293 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - os.chdir(ROOT_DIR) - -import os -import hydra -import torch -from omegaconf import OmegaConf -import pathlib -from torch.utils.data import DataLoader -import copy -import random -import wandb -import tqdm -import numpy as np -import shutil - -from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.bet_lowdim_policy import BETLowdimPolicy -from diffusion_policy.dataset.base_dataset import BaseLowdimDataset -from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner -from diffusion_policy.common.checkpoint_util import TopKCheckpointManager -from diffusion_policy.model.common.normalizer import ( - LinearNormalizer, - SingleFieldLinearNormalizer -) -from diffusion_policy.common.json_logger import JsonLogger -from diffusers.training_utils import EMAModel - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -# %% -class TrainBETLowdimWorkspace(BaseWorkspace): - include_keys = ['global_step', 'epoch'] - - def __init__(self, cfg: OmegaConf): - super().__init__(cfg) - - # set seed - seed = cfg.training.seed - torch.manual_seed(seed) - np.random.seed(seed) - random.seed(seed) - - # configure model - self.policy: BETLowdimPolicy - self.policy = hydra.utils.instantiate(cfg.policy) - - # configure training state - self.optimizer = self.policy.get_optimizer(**cfg.optimizer) - - self.global_step = 0 - self.epoch = 0 - - - def run(self): - cfg = copy.deepcopy(self.cfg) - OmegaConf.resolve(cfg) - - # resume training - if cfg.training.resume: - lastest_ckpt_path = self.get_checkpoint_path() - if lastest_ckpt_path.is_file(): - print(f"Resuming from checkpoint {lastest_ckpt_path}") - self.load_checkpoint(path=lastest_ckpt_path) - - # configure dataset - dataset: BaseLowdimDataset - dataset = hydra.utils.instantiate(cfg.task.dataset) - assert isinstance(dataset, BaseLowdimDataset) - train_dataloader = DataLoader(dataset, **cfg.dataloader) - - # configure validation dataset - val_dataset = dataset.get_validation_dataset() - val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader) - - # set normalizer - normalizer = None - if cfg.training.enable_normalizer: - normalizer = dataset.get_normalizer() - else: - normalizer = LinearNormalizer() - normalizer['action'] = SingleFieldLinearNormalizer.create_identity() - normalizer['obs'] = SingleFieldLinearNormalizer.create_identity() - - self.policy.set_normalizer(normalizer) - - # fit action_ae (K-Means) - self.policy.fit_action_ae( - normalizer['action'].normalize( - dataset.get_all_actions())) - - # configure env runner - env_runner: BaseLowdimRunner - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=self.output_dir) - assert isinstance(env_runner, BaseLowdimRunner) - - # configure logging - wandb_run = wandb.init( - dir=str(self.output_dir), - config=OmegaConf.to_container(cfg, resolve=True), - **cfg.logging - ) - wandb.config.update( - { - "output_dir": self.output_dir, - } - ) - - # configure checkpoint - topk_manager = TopKCheckpointManager( - save_dir=os.path.join(self.output_dir, 'checkpoints'), - **cfg.checkpoint.topk - ) - - # device transfer - device = torch.device(cfg.training.device) - self.policy.to(device) - optimizer_to(self.optimizer, device) - - # save batch for sampling - train_sampling_batch = None - - if cfg.training.debug: - cfg.training.num_epochs = 2 - cfg.training.max_train_steps = 3 - cfg.training.max_val_steps = 3 - cfg.training.rollout_every = 1 - cfg.training.checkpoint_every = 1 - cfg.training.val_every = 1 - cfg.training.sample_every = 1 - - # training loop - log_path = os.path.join(self.output_dir, 'logs.json.txt') - with JsonLogger(log_path) as json_logger: - for local_epoch_idx in range(cfg.training.num_epochs): - step_log = dict() - # ========= train for this epoch ========== - train_losses = list() - with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - # device transfer - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - if train_sampling_batch is None: - train_sampling_batch = batch - - # compute loss - raw_loss, loss_components = self.policy.compute_loss(batch) - loss = raw_loss / cfg.training.gradient_accumulate_every - loss.backward() - - # clip grad norm - torch.nn.utils.clip_grad_norm_( - self.policy.state_prior.parameters(), cfg.training.grad_norm_clip - ) - - # step optimizer - if self.global_step % cfg.training.gradient_accumulate_every == 0: - self.optimizer.step() - self.optimizer.zero_grad(set_to_none=True) - - # logging - raw_loss_cpu = raw_loss.item() - tepoch.set_postfix(loss=raw_loss_cpu, refresh=False) - train_losses.append(raw_loss_cpu) - step_log = { - 'train_loss': raw_loss_cpu, - 'global_step': self.global_step, - 'epoch': self.epoch, - 'train_loss_offset': loss_components['offset'].item(), - 'train_loss_class': loss_components['class'].item() - } - - is_last_batch = (batch_idx == (len(train_dataloader)-1)) - if not is_last_batch: - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - - if (cfg.training.max_train_steps is not None) \ - and batch_idx >= (cfg.training.max_train_steps-1): - break - - # at the end of each epoch - # replace train_loss with epoch average - train_loss = np.mean(train_losses) - step_log['train_loss'] = train_loss - - # ========= eval for this epoch ========== - self.policy.eval() - - # run rollout - if (self.epoch % cfg.training.rollout_every) == 0: - runner_log = env_runner.run(self.policy) - # log all - step_log.update(runner_log) - - # run validation - if (self.epoch % cfg.training.val_every) == 0: - with torch.no_grad(): - val_losses = list() - with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - raw_loss, loss_components = self.policy.compute_loss(batch) - val_losses.append(raw_loss) - if (cfg.training.max_val_steps is not None) \ - and batch_idx >= (cfg.training.max_val_steps-1): - break - if len(val_losses) > 0: - val_loss = torch.mean(torch.tensor(val_losses)).item() - # log epoch average validation loss - step_log['val_loss'] = val_loss - - # run sample on a training batch - if (self.epoch % cfg.training.sample_every) == 0: - with torch.no_grad(): - # sample trajectory from training set, and evaluate difference - batch = train_sampling_batch - obs_dict = {'obs': batch['obs']} - gt_action = batch['action'] - - result = self.policy.predict_action(obs_dict) - if cfg.pred_action_steps_only: - pred_action = result['action'] - start = cfg.n_obs_steps - 1 - end = start + cfg.n_action_steps - gt_action = gt_action[:,start:end] - else: - pred_action = result['action_pred'] - mse = torch.nn.functional.mse_loss(pred_action, gt_action) - # log - step_log['train_action_mse_error'] = mse.item() - # release RAM - del batch - del obs_dict - del gt_action - del result - del pred_action - del mse - - # checkpoint - if (self.epoch % cfg.training.checkpoint_every) == 0: - # checkpointing - if cfg.checkpoint.save_last_ckpt: - self.save_checkpoint() - if cfg.checkpoint.save_last_snapshot: - self.save_snapshot() - - # sanitize metric names - metric_dict = dict() - for key, value in step_log.items(): - new_key = key.replace('/', '_') - metric_dict[new_key] = value - - # We can't copy the last checkpoint here - # since save_checkpoint uses threads. - # therefore at this point the file might have been empty! - topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict) - - if topk_ckpt_path is not None: - self.save_checkpoint(path=topk_ckpt_path) - # ========= eval end for this epoch ========== - self.policy.train() - - # end of epoch - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - self.epoch += 1 - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), - config_name=pathlib.Path(__file__).stem) -def main(cfg): - workspace = TrainBETLowdimWorkspace(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_transformer_hybrid_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_transformer_hybrid_workspace.py deleted file mode 100644 index 38b3be135..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_transformer_hybrid_workspace.py +++ /dev/null @@ -1,295 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - os.chdir(ROOT_DIR) - -import os -import hydra -import torch -from omegaconf import OmegaConf -import pathlib -from torch.utils.data import DataLoader -import copy -import random -import wandb -import tqdm -import numpy as np -import shutil -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.diffusion_transformer_hybrid_image_policy import DiffusionTransformerHybridImagePolicy -from diffusion_policy.dataset.base_dataset import BaseImageDataset -from diffusion_policy.env_runner.base_image_runner import BaseImageRunner -from diffusion_policy.common.checkpoint_util import TopKCheckpointManager -from diffusion_policy.common.json_logger import JsonLogger -from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to -from diffusion_policy.model.diffusion.ema_model import EMAModel -from diffusion_policy.model.common.lr_scheduler import get_scheduler - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -class TrainDiffusionTransformerHybridWorkspace(BaseWorkspace): - include_keys = ['global_step', 'epoch'] - - def __init__(self, cfg: OmegaConf, output_dir=None): - super().__init__(cfg, output_dir=output_dir) - - # set seed - seed = cfg.training.seed - torch.manual_seed(seed) - np.random.seed(seed) - random.seed(seed) - - # configure model - self.model: DiffusionTransformerHybridImagePolicy = hydra.utils.instantiate(cfg.policy) - - self.ema_model: DiffusionTransformerHybridImagePolicy = None - if cfg.training.use_ema: - self.ema_model = copy.deepcopy(self.model) - - # configure training state - self.optimizer = self.model.get_optimizer(**cfg.optimizer) - - # configure training state - self.global_step = 0 - self.epoch = 0 - - def run(self): - cfg = copy.deepcopy(self.cfg) - - # resume training - if cfg.training.resume: - lastest_ckpt_path = self.get_checkpoint_path() - if lastest_ckpt_path.is_file(): - print(f"Resuming from checkpoint {lastest_ckpt_path}") - self.load_checkpoint(path=lastest_ckpt_path) - - # configure dataset - dataset: BaseImageDataset - dataset = hydra.utils.instantiate(cfg.task.dataset) - assert isinstance(dataset, BaseImageDataset) - train_dataloader = DataLoader(dataset, **cfg.dataloader) - normalizer = dataset.get_normalizer() - - # configure validation dataset - val_dataset = dataset.get_validation_dataset() - val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader) - - self.model.set_normalizer(normalizer) - if cfg.training.use_ema: - self.ema_model.set_normalizer(normalizer) - - # configure lr scheduler - lr_scheduler = get_scheduler( - cfg.training.lr_scheduler, - optimizer=self.optimizer, - num_warmup_steps=cfg.training.lr_warmup_steps, - num_training_steps=( - len(train_dataloader) * cfg.training.num_epochs) \ - // cfg.training.gradient_accumulate_every, - # pytorch assumes stepping LRScheduler every epoch - # however huggingface diffusers steps it every batch - last_epoch=self.global_step-1 - ) - - # configure ema - ema: EMAModel = None - if cfg.training.use_ema: - ema = hydra.utils.instantiate( - cfg.ema, - model=self.ema_model) - - # configure env - env_runner: BaseImageRunner - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=self.output_dir) - assert isinstance(env_runner, BaseImageRunner) - - # configure logging - wandb_run = wandb.init( - dir=str(self.output_dir), - config=OmegaConf.to_container(cfg, resolve=True), - **cfg.logging - ) - wandb.config.update( - { - "output_dir": self.output_dir, - } - ) - - # configure checkpoint - topk_manager = TopKCheckpointManager( - save_dir=os.path.join(self.output_dir, 'checkpoints'), - **cfg.checkpoint.topk - ) - - # device transfer - device = torch.device(cfg.training.device) - self.model.to(device) - if self.ema_model is not None: - self.ema_model.to(device) - optimizer_to(self.optimizer, device) - - # save batch for sampling - train_sampling_batch = None - - if cfg.training.debug: - cfg.training.num_epochs = 2 - cfg.training.max_train_steps = 3 - cfg.training.max_val_steps = 3 - cfg.training.rollout_every = 1 - cfg.training.checkpoint_every = 1 - cfg.training.val_every = 1 - cfg.training.sample_every = 1 - - # training loop - log_path = os.path.join(self.output_dir, 'logs.json.txt') - with JsonLogger(log_path) as json_logger: - for local_epoch_idx in range(cfg.training.num_epochs): - step_log = dict() - # ========= train for this epoch ========== - train_losses = list() - with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - # device transfer - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - if train_sampling_batch is None: - train_sampling_batch = batch - - # compute loss - raw_loss = self.model.compute_loss(batch) - loss = raw_loss / cfg.training.gradient_accumulate_every - loss.backward() - - # step optimizer - if self.global_step % cfg.training.gradient_accumulate_every == 0: - self.optimizer.step() - self.optimizer.zero_grad() - lr_scheduler.step() - - # update ema - if cfg.training.use_ema: - ema.step(self.model) - - # logging - raw_loss_cpu = raw_loss.item() - tepoch.set_postfix(loss=raw_loss_cpu, refresh=False) - train_losses.append(raw_loss_cpu) - step_log = { - 'train_loss': raw_loss_cpu, - 'global_step': self.global_step, - 'epoch': self.epoch, - 'lr': lr_scheduler.get_last_lr()[0] - } - - is_last_batch = (batch_idx == (len(train_dataloader)-1)) - if not is_last_batch: - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - - if (cfg.training.max_train_steps is not None) \ - and batch_idx >= (cfg.training.max_train_steps-1): - break - - # at the end of each epoch - # replace train_loss with epoch average - train_loss = np.mean(train_losses) - step_log['train_loss'] = train_loss - - # ========= eval for this epoch ========== - policy = self.model - if cfg.training.use_ema: - policy = self.ema_model - policy.eval() - - # run rollout - if (self.epoch % cfg.training.rollout_every) == 0: - runner_log = env_runner.run(policy) - # log all - step_log.update(runner_log) - - # run validation - if (self.epoch % cfg.training.val_every) == 0: - with torch.no_grad(): - val_losses = list() - with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - loss = self.model.compute_loss(batch) - val_losses.append(loss) - if (cfg.training.max_val_steps is not None) \ - and batch_idx >= (cfg.training.max_val_steps-1): - break - if len(val_losses) > 0: - val_loss = torch.mean(torch.tensor(val_losses)).item() - # log epoch average validation loss - step_log['val_loss'] = val_loss - - # run diffusion sampling on a training batch - if (self.epoch % cfg.training.sample_every) == 0: - with torch.no_grad(): - # sample trajectory from training set, and evaluate difference - batch = dict_apply(train_sampling_batch, lambda x: x.to(device, non_blocking=True)) - obs_dict = batch['obs'] - gt_action = batch['action'] - - result = policy.predict_action(obs_dict) - pred_action = result['action_pred'] - mse = torch.nn.functional.mse_loss(pred_action, gt_action) - step_log['train_action_mse_error'] = mse.item() - del batch - del obs_dict - del gt_action - del result - del pred_action - del mse - - # checkpoint - if (self.epoch % cfg.training.checkpoint_every) == 0: - # checkpointing - if cfg.checkpoint.save_last_ckpt: - self.save_checkpoint() - if cfg.checkpoint.save_last_snapshot: - self.save_snapshot() - - # sanitize metric names - metric_dict = dict() - for key, value in step_log.items(): - new_key = key.replace('/', '_') - metric_dict[new_key] = value - - # We can't copy the last checkpoint here - # since save_checkpoint uses threads. - # therefore at this point the file might have been empty! - topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict) - - if topk_ckpt_path is not None: - self.save_checkpoint(path=topk_ckpt_path) - # ========= eval end for this epoch ========== - policy.train() - - # end of epoch - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - self.epoch += 1 - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), - config_name=pathlib.Path(__file__).stem) -def main(cfg): - workspace = TrainDiffusionTransformerHybridWorkspace(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_transformer_lowdim_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_transformer_lowdim_workspace.py deleted file mode 100644 index cbfa54b5d..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_transformer_lowdim_workspace.py +++ /dev/null @@ -1,303 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - os.chdir(ROOT_DIR) - -import os -import hydra -import torch -from omegaconf import OmegaConf -import pathlib -from torch.utils.data import DataLoader -import copy -import random -import wandb -import tqdm -import numpy as np -import shutil - -from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.diffusion_transformer_lowdim_policy import DiffusionTransformerLowdimPolicy -from diffusion_policy.dataset.base_dataset import BaseLowdimDataset -from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner -from diffusion_policy.common.checkpoint_util import TopKCheckpointManager -from diffusion_policy.common.json_logger import JsonLogger -from diffusion_policy.model.common.lr_scheduler import get_scheduler -from diffusers.training_utils import EMAModel - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -# %% -class TrainDiffusionTransformerLowdimWorkspace(BaseWorkspace): - include_keys = ['global_step', 'epoch'] - - def __init__(self, cfg: OmegaConf): - super().__init__(cfg) - - # set seed - seed = cfg.training.seed - torch.manual_seed(seed) - np.random.seed(seed) - random.seed(seed) - - # configure model - self.model: DiffusionTransformerLowdimPolicy - self.model = hydra.utils.instantiate(cfg.policy) - - self.ema_model: DiffusionTransformerLowdimPolicy = None - if cfg.training.use_ema: - self.ema_model = copy.deepcopy(self.model) - - # configure training state - self.optimizer = self.model.get_optimizer(**cfg.optimizer) - - self.global_step = 0 - self.epoch = 0 - - def run(self): - cfg = copy.deepcopy(self.cfg) - - # resume training - if cfg.training.resume: - lastest_ckpt_path = self.get_checkpoint_path() - if lastest_ckpt_path.is_file(): - print(f"Resuming from checkpoint {lastest_ckpt_path}") - self.load_checkpoint(path=lastest_ckpt_path) - - # configure dataset - dataset: BaseLowdimDataset - dataset = hydra.utils.instantiate(cfg.task.dataset) - assert isinstance(dataset, BaseLowdimDataset) - train_dataloader = DataLoader(dataset, **cfg.dataloader) - normalizer = dataset.get_normalizer() - - # configure validation dataset - val_dataset = dataset.get_validation_dataset() - val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader) - - self.model.set_normalizer(normalizer) - if cfg.training.use_ema: - self.ema_model.set_normalizer(normalizer) - - # configure lr scheduler - lr_scheduler = get_scheduler( - cfg.training.lr_scheduler, - optimizer=self.optimizer, - num_warmup_steps=cfg.training.lr_warmup_steps, - num_training_steps=( - len(train_dataloader) * cfg.training.num_epochs) \ - // cfg.training.gradient_accumulate_every, - # pytorch assumes stepping LRScheduler every epoch - # however huggingface diffusers steps it every batch - last_epoch=self.global_step-1 - ) - - # configure ema - ema: EMAModel = None - if cfg.training.use_ema: - ema = hydra.utils.instantiate( - cfg.ema, - model=self.ema_model) - - # configure env runner - env_runner: BaseLowdimRunner - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=self.output_dir) - assert isinstance(env_runner, BaseLowdimRunner) - - # configure logging - wandb_run = wandb.init( - dir=str(self.output_dir), - config=OmegaConf.to_container(cfg, resolve=True), - **cfg.logging - ) - wandb.config.update( - { - "output_dir": self.output_dir, - } - ) - - # configure checkpoint - topk_manager = TopKCheckpointManager( - save_dir=os.path.join(self.output_dir, 'checkpoints'), - **cfg.checkpoint.topk - ) - - # device transfer - device = torch.device(cfg.training.device) - self.model.to(device) - if self.ema_model is not None: - self.ema_model.to(device) - optimizer_to(self.optimizer, device) - - # save batch for sampling - train_sampling_batch = None - - if cfg.training.debug: - cfg.training.num_epochs = 2 - cfg.training.max_train_steps = 3 - cfg.training.max_val_steps = 3 - cfg.training.rollout_every = 1 - cfg.training.checkpoint_every = 1 - cfg.training.val_every = 1 - cfg.training.sample_every = 1 - - # training loop - log_path = os.path.join(self.output_dir, 'logs.json.txt') - with JsonLogger(log_path) as json_logger: - for local_epoch_idx in range(cfg.training.num_epochs): - step_log = dict() - # ========= train for this epoch ========== - train_losses = list() - with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - # device transfer - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - if train_sampling_batch is None: - train_sampling_batch = batch - - # compute loss - raw_loss = self.model.compute_loss(batch) - loss = raw_loss / cfg.training.gradient_accumulate_every - loss.backward() - - # step optimizer - if self.global_step % cfg.training.gradient_accumulate_every == 0: - self.optimizer.step() - self.optimizer.zero_grad() - lr_scheduler.step() - - # update ema - if cfg.training.use_ema: - ema.step(self.model) - - # logging - raw_loss_cpu = raw_loss.item() - tepoch.set_postfix(loss=raw_loss_cpu, refresh=False) - train_losses.append(raw_loss_cpu) - step_log = { - 'train_loss': raw_loss_cpu, - 'global_step': self.global_step, - 'epoch': self.epoch, - 'lr': lr_scheduler.get_last_lr()[0] - } - - is_last_batch = (batch_idx == (len(train_dataloader)-1)) - if not is_last_batch: - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - - if (cfg.training.max_train_steps is not None) \ - and batch_idx >= (cfg.training.max_train_steps-1): - break - - # at the end of each epoch - # replace train_loss with epoch average - train_loss = np.mean(train_losses) - step_log['train_loss'] = train_loss - - # ========= eval for this epoch ========== - policy = self.model - if cfg.training.use_ema: - policy = self.ema_model - policy.eval() - - # run rollout - if (self.epoch % cfg.training.rollout_every) == 0: - runner_log = env_runner.run(policy) - # log all - step_log.update(runner_log) - - # run validation - if (self.epoch % cfg.training.val_every) == 0: - with torch.no_grad(): - val_losses = list() - with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - loss = self.model.compute_loss(batch) - val_losses.append(loss) - if (cfg.training.max_val_steps is not None) \ - and batch_idx >= (cfg.training.max_val_steps-1): - break - if len(val_losses) > 0: - val_loss = torch.mean(torch.tensor(val_losses)).item() - # log epoch average validation loss - step_log['val_loss'] = val_loss - - # run diffusion sampling on a training batch - if (self.epoch % cfg.training.sample_every) == 0: - with torch.no_grad(): - # sample trajectory from training set, and evaluate difference - batch = dict_apply(train_sampling_batch, lambda x: x.to(device, non_blocking=True)) - obs_dict = {'obs': batch['obs']} - gt_action = batch['action'] - - result = policy.predict_action(obs_dict) - if cfg.pred_action_steps_only: - pred_action = result['action'] - start = cfg.n_obs_steps - 1 - end = start + cfg.n_action_steps - gt_action = gt_action[:,start:end] - else: - pred_action = result['action_pred'] - mse = torch.nn.functional.mse_loss(pred_action, gt_action) - step_log['train_action_mse_error'] = mse.item() - del batch - del obs_dict - del gt_action - del result - del pred_action - del mse - - # checkpoint - if (self.epoch % cfg.training.checkpoint_every) == 0: - # checkpointing - if cfg.checkpoint.save_last_ckpt: - self.save_checkpoint() - if cfg.checkpoint.save_last_snapshot: - self.save_snapshot() - - # sanitize metric names - metric_dict = dict() - for key, value in step_log.items(): - new_key = key.replace('/', '_') - metric_dict[new_key] = value - - # We can't copy the last checkpoint here - # since save_checkpoint uses threads. - # therefore at this point the file might have been empty! - topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict) - - if topk_ckpt_path is not None: - self.save_checkpoint(path=topk_ckpt_path) - # ========= eval end for this epoch ========== - policy.train() - - # end of epoch - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - self.epoch += 1 - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), - config_name=pathlib.Path(__file__).stem) -def main(cfg): - workspace = TrainDiffusionTransformerLowdimWorkspace(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_hybrid_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_hybrid_workspace.py deleted file mode 100644 index 9219427ca..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_hybrid_workspace.py +++ /dev/null @@ -1,296 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - os.chdir(ROOT_DIR) - -import os -import hydra -import torch -from omegaconf import OmegaConf -import pathlib -from torch.utils.data import DataLoader -import copy -import random -import wandb -import tqdm -import numpy as np -import shutil -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.diffusion_unet_hybrid_image_policy import DiffusionUnetHybridImagePolicy -from diffusion_policy.dataset.base_dataset import BaseImageDataset -from diffusion_policy.env_runner.base_image_runner import BaseImageRunner -from diffusion_policy.common.checkpoint_util import TopKCheckpointManager -from diffusion_policy.common.json_logger import JsonLogger -from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to -from diffusion_policy.model.diffusion.ema_model import EMAModel -from diffusion_policy.model.common.lr_scheduler import get_scheduler - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -class TrainDiffusionUnetHybridWorkspace(BaseWorkspace): - include_keys = ['global_step', 'epoch'] - - def __init__(self, cfg: OmegaConf, output_dir=None): - super().__init__(cfg, output_dir=output_dir) - - # set seed - seed = cfg.training.seed - torch.manual_seed(seed) - np.random.seed(seed) - random.seed(seed) - - # configure model - self.model: DiffusionUnetHybridImagePolicy = hydra.utils.instantiate(cfg.policy) - - self.ema_model: DiffusionUnetHybridImagePolicy = None - if cfg.training.use_ema: - self.ema_model = copy.deepcopy(self.model) - - # configure training state - self.optimizer = hydra.utils.instantiate( - cfg.optimizer, params=self.model.parameters()) - - # configure training state - self.global_step = 0 - self.epoch = 0 - - def run(self): - cfg = copy.deepcopy(self.cfg) - - # resume training - if cfg.training.resume: - lastest_ckpt_path = self.get_checkpoint_path() - if lastest_ckpt_path.is_file(): - print(f"Resuming from checkpoint {lastest_ckpt_path}") - self.load_checkpoint(path=lastest_ckpt_path) - - # configure dataset - dataset: BaseImageDataset - dataset = hydra.utils.instantiate(cfg.task.dataset) - assert isinstance(dataset, BaseImageDataset) - train_dataloader = DataLoader(dataset, **cfg.dataloader) - normalizer = dataset.get_normalizer() - - # configure validation dataset - val_dataset = dataset.get_validation_dataset() - val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader) - - self.model.set_normalizer(normalizer) - if cfg.training.use_ema: - self.ema_model.set_normalizer(normalizer) - - # configure lr scheduler - lr_scheduler = get_scheduler( - cfg.training.lr_scheduler, - optimizer=self.optimizer, - num_warmup_steps=cfg.training.lr_warmup_steps, - num_training_steps=( - len(train_dataloader) * cfg.training.num_epochs) \ - // cfg.training.gradient_accumulate_every, - # pytorch assumes stepping LRScheduler every epoch - # however huggingface diffusers steps it every batch - last_epoch=self.global_step-1 - ) - - # configure ema - ema: EMAModel = None - if cfg.training.use_ema: - ema = hydra.utils.instantiate( - cfg.ema, - model=self.ema_model) - - # configure env - env_runner: BaseImageRunner - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=self.output_dir) - assert isinstance(env_runner, BaseImageRunner) - - # configure logging - wandb_run = wandb.init( - dir=str(self.output_dir), - config=OmegaConf.to_container(cfg, resolve=True), - **cfg.logging - ) - wandb.config.update( - { - "output_dir": self.output_dir, - } - ) - - # configure checkpoint - topk_manager = TopKCheckpointManager( - save_dir=os.path.join(self.output_dir, 'checkpoints'), - **cfg.checkpoint.topk - ) - - # device transfer - device = torch.device(cfg.training.device) - self.model.to(device) - if self.ema_model is not None: - self.ema_model.to(device) - optimizer_to(self.optimizer, device) - - # save batch for sampling - train_sampling_batch = None - - if cfg.training.debug: - cfg.training.num_epochs = 2 - cfg.training.max_train_steps = 3 - cfg.training.max_val_steps = 3 - cfg.training.rollout_every = 1 - cfg.training.checkpoint_every = 1 - cfg.training.val_every = 1 - cfg.training.sample_every = 1 - - # training loop - log_path = os.path.join(self.output_dir, 'logs.json.txt') - with JsonLogger(log_path) as json_logger: - for local_epoch_idx in range(cfg.training.num_epochs): - step_log = dict() - # ========= train for this epoch ========== - train_losses = list() - with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - # device transfer - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - if train_sampling_batch is None: - train_sampling_batch = batch - - # compute loss - raw_loss = self.model.compute_loss(batch) - loss = raw_loss / cfg.training.gradient_accumulate_every - loss.backward() - - # step optimizer - if self.global_step % cfg.training.gradient_accumulate_every == 0: - self.optimizer.step() - self.optimizer.zero_grad() - lr_scheduler.step() - - # update ema - if cfg.training.use_ema: - ema.step(self.model) - - # logging - raw_loss_cpu = raw_loss.item() - tepoch.set_postfix(loss=raw_loss_cpu, refresh=False) - train_losses.append(raw_loss_cpu) - step_log = { - 'train_loss': raw_loss_cpu, - 'global_step': self.global_step, - 'epoch': self.epoch, - 'lr': lr_scheduler.get_last_lr()[0] - } - - is_last_batch = (batch_idx == (len(train_dataloader)-1)) - if not is_last_batch: - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - - if (cfg.training.max_train_steps is not None) \ - and batch_idx >= (cfg.training.max_train_steps-1): - break - - # at the end of each epoch - # replace train_loss with epoch average - train_loss = np.mean(train_losses) - step_log['train_loss'] = train_loss - - # ========= eval for this epoch ========== - policy = self.model - if cfg.training.use_ema: - policy = self.ema_model - policy.eval() - - # run rollout - if (self.epoch % cfg.training.rollout_every) == 0: - runner_log = env_runner.run(policy) - # log all - step_log.update(runner_log) - - # run validation - if (self.epoch % cfg.training.val_every) == 0: - with torch.no_grad(): - val_losses = list() - with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - loss = self.model.compute_loss(batch) - val_losses.append(loss) - if (cfg.training.max_val_steps is not None) \ - and batch_idx >= (cfg.training.max_val_steps-1): - break - if len(val_losses) > 0: - val_loss = torch.mean(torch.tensor(val_losses)).item() - # log epoch average validation loss - step_log['val_loss'] = val_loss - - # run diffusion sampling on a training batch - if (self.epoch % cfg.training.sample_every) == 0: - with torch.no_grad(): - # sample trajectory from training set, and evaluate difference - batch = dict_apply(train_sampling_batch, lambda x: x.to(device, non_blocking=True)) - obs_dict = batch['obs'] - gt_action = batch['action'] - - result = policy.predict_action(obs_dict) - pred_action = result['action_pred'] - mse = torch.nn.functional.mse_loss(pred_action, gt_action) - step_log['train_action_mse_error'] = mse.item() - del batch - del obs_dict - del gt_action - del result - del pred_action - del mse - - # checkpoint - if (self.epoch % cfg.training.checkpoint_every) == 0: - # checkpointing - if cfg.checkpoint.save_last_ckpt: - self.save_checkpoint() - if cfg.checkpoint.save_last_snapshot: - self.save_snapshot() - - # sanitize metric names - metric_dict = dict() - for key, value in step_log.items(): - new_key = key.replace('/', '_') - metric_dict[new_key] = value - - # We can't copy the last checkpoint here - # since save_checkpoint uses threads. - # therefore at this point the file might have been empty! - topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict) - - if topk_ckpt_path is not None: - self.save_checkpoint(path=topk_ckpt_path) - # ========= eval end for this epoch ========== - policy.train() - - # end of epoch - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - self.epoch += 1 - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), - config_name=pathlib.Path(__file__).stem) -def main(cfg): - workspace = TrainDiffusionUnetHybridWorkspace(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_image_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_image_workspace.py deleted file mode 100644 index d8302351c..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_image_workspace.py +++ /dev/null @@ -1,300 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - os.chdir(ROOT_DIR) - -import os -import hydra -import torch -from omegaconf import OmegaConf -import pathlib -from torch.utils.data import DataLoader -import copy -import random -import wandb -import tqdm -import numpy as np -import shutil -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.diffusion_unet_image_policy import DiffusionUnetImagePolicy -from diffusion_policy.dataset.base_dataset import BaseImageDataset -from diffusion_policy.env_runner.base_image_runner import BaseImageRunner -from diffusion_policy.common.checkpoint_util import TopKCheckpointManager -from diffusion_policy.common.json_logger import JsonLogger -from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to -from diffusion_policy.model.diffusion.ema_model import EMAModel -from diffusion_policy.model.common.lr_scheduler import get_scheduler - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -class TrainDiffusionUnetImageWorkspace(BaseWorkspace): - include_keys = ['global_step', 'epoch'] - - def __init__(self, cfg: OmegaConf, output_dir=None): - super().__init__(cfg, output_dir=output_dir) - - # set seed - seed = cfg.training.seed - torch.manual_seed(seed) - np.random.seed(seed) - random.seed(seed) - - # configure model - self.model: DiffusionUnetImagePolicy = hydra.utils.instantiate(cfg.policy) - - self.ema_model: DiffusionUnetImagePolicy = None - if cfg.training.use_ema: - self.ema_model = copy.deepcopy(self.model) - - # configure training state - self.optimizer = hydra.utils.instantiate( - cfg.optimizer, params=self.model.parameters()) - - # configure training state - self.global_step = 0 - self.epoch = 0 - - def run(self): - cfg = copy.deepcopy(self.cfg) - - # resume training - if cfg.training.resume: - lastest_ckpt_path = self.get_checkpoint_path() - if lastest_ckpt_path.is_file(): - print(f"Resuming from checkpoint {lastest_ckpt_path}") - self.load_checkpoint(path=lastest_ckpt_path) - - # configure dataset - dataset: BaseImageDataset - dataset = hydra.utils.instantiate(cfg.task.dataset) - assert isinstance(dataset, BaseImageDataset) - train_dataloader = DataLoader(dataset, **cfg.dataloader) - normalizer = dataset.get_normalizer() - - # configure validation dataset - val_dataset = dataset.get_validation_dataset() - val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader) - - self.model.set_normalizer(normalizer) - if cfg.training.use_ema: - self.ema_model.set_normalizer(normalizer) - - # configure lr scheduler - lr_scheduler = get_scheduler( - cfg.training.lr_scheduler, - optimizer=self.optimizer, - num_warmup_steps=cfg.training.lr_warmup_steps, - num_training_steps=( - len(train_dataloader) * cfg.training.num_epochs) \ - // cfg.training.gradient_accumulate_every, - # pytorch assumes stepping LRScheduler every epoch - # however huggingface diffusers steps it every batch - last_epoch=self.global_step-1 - ) - - # configure ema - ema: EMAModel = None - if cfg.training.use_ema: - ema = hydra.utils.instantiate( - cfg.ema, - model=self.ema_model) - - # configure env - env_runner: BaseImageRunner - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=self.output_dir) - assert isinstance(env_runner, BaseImageRunner) - - # configure logging - wandb_run = wandb.init( - dir=str(self.output_dir), - config=OmegaConf.to_container(cfg, resolve=True), - **cfg.logging - ) - wandb.config.update( - { - "output_dir": self.output_dir, - } - ) - - # configure checkpoint - topk_manager = TopKCheckpointManager( - save_dir=os.path.join(self.output_dir, 'checkpoints'), - **cfg.checkpoint.topk - ) - - # device transfer - device = torch.device(cfg.training.device) - self.model.to(device) - if self.ema_model is not None: - self.ema_model.to(device) - optimizer_to(self.optimizer, device) - - # save batch for sampling - train_sampling_batch = None - - if cfg.training.debug: - cfg.training.num_epochs = 2 - cfg.training.max_train_steps = 3 - cfg.training.max_val_steps = 3 - cfg.training.rollout_every = 1 - cfg.training.checkpoint_every = 1 - cfg.training.val_every = 1 - cfg.training.sample_every = 1 - - # training loop - log_path = os.path.join(self.output_dir, 'logs.json.txt') - with JsonLogger(log_path) as json_logger: - for local_epoch_idx in range(cfg.training.num_epochs): - step_log = dict() - # ========= train for this epoch ========== - if cfg.training.freeze_encoder: - self.model.obs_encoder.eval() - self.model.obs_encoder.requires_grad_(False) - - train_losses = list() - with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - # device transfer - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - if train_sampling_batch is None: - train_sampling_batch = batch - - # compute loss - raw_loss = self.model.compute_loss(batch) - loss = raw_loss / cfg.training.gradient_accumulate_every - loss.backward() - - # step optimizer - if self.global_step % cfg.training.gradient_accumulate_every == 0: - self.optimizer.step() - self.optimizer.zero_grad() - lr_scheduler.step() - - # update ema - if cfg.training.use_ema: - ema.step(self.model) - - # logging - raw_loss_cpu = raw_loss.item() - tepoch.set_postfix(loss=raw_loss_cpu, refresh=False) - train_losses.append(raw_loss_cpu) - step_log = { - 'train_loss': raw_loss_cpu, - 'global_step': self.global_step, - 'epoch': self.epoch, - 'lr': lr_scheduler.get_last_lr()[0] - } - - is_last_batch = (batch_idx == (len(train_dataloader)-1)) - if not is_last_batch: - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - - if (cfg.training.max_train_steps is not None) \ - and batch_idx >= (cfg.training.max_train_steps-1): - break - - # at the end of each epoch - # replace train_loss with epoch average - train_loss = np.mean(train_losses) - step_log['train_loss'] = train_loss - - # ========= eval for this epoch ========== - policy = self.model - if cfg.training.use_ema: - policy = self.ema_model - policy.eval() - - # run rollout - if (self.epoch % cfg.training.rollout_every) == 0: - runner_log = env_runner.run(policy) - # log all - step_log.update(runner_log) - - # run validation - if (self.epoch % cfg.training.val_every) == 0: - with torch.no_grad(): - val_losses = list() - with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - loss = self.model.compute_loss(batch) - val_losses.append(loss) - if (cfg.training.max_val_steps is not None) \ - and batch_idx >= (cfg.training.max_val_steps-1): - break - if len(val_losses) > 0: - val_loss = torch.mean(torch.tensor(val_losses)).item() - # log epoch average validation loss - step_log['val_loss'] = val_loss - - # run diffusion sampling on a training batch - if (self.epoch % cfg.training.sample_every) == 0: - with torch.no_grad(): - # sample trajectory from training set, and evaluate difference - batch = dict_apply(train_sampling_batch, lambda x: x.to(device, non_blocking=True)) - obs_dict = batch['obs'] - gt_action = batch['action'] - - result = policy.predict_action(obs_dict) - pred_action = result['action_pred'] - mse = torch.nn.functional.mse_loss(pred_action, gt_action) - step_log['train_action_mse_error'] = mse.item() - del batch - del obs_dict - del gt_action - del result - del pred_action - del mse - - # checkpoint - if (self.epoch % cfg.training.checkpoint_every) == 0: - # checkpointing - if cfg.checkpoint.save_last_ckpt: - self.save_checkpoint() - if cfg.checkpoint.save_last_snapshot: - self.save_snapshot() - - # sanitize metric names - metric_dict = dict() - for key, value in step_log.items(): - new_key = key.replace('/', '_') - metric_dict[new_key] = value - - # We can't copy the last checkpoint here - # since save_checkpoint uses threads. - # therefore at this point the file might have been empty! - topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict) - - if topk_ckpt_path is not None: - self.save_checkpoint(path=topk_ckpt_path) - # ========= eval end for this epoch ========== - policy.train() - - # end of epoch - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - self.epoch += 1 - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), - config_name=pathlib.Path(__file__).stem) -def main(cfg): - workspace = TrainDiffusionUnetImageWorkspace(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_lowdim_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_lowdim_workspace.py deleted file mode 100644 index 617b8aeb2..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_lowdim_workspace.py +++ /dev/null @@ -1,306 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - os.chdir(ROOT_DIR) - -import os -import hydra -import torch -from omegaconf import OmegaConf -import pathlib -from torch.utils.data import DataLoader -import copy -import numpy as np -import random -import wandb -import tqdm -import shutil - -from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.diffusion_unet_lowdim_policy import DiffusionUnetLowdimPolicy -from diffusion_policy.dataset.base_dataset import BaseLowdimDataset -from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner -from diffusion_policy.common.checkpoint_util import TopKCheckpointManager -from diffusion_policy.common.json_logger import JsonLogger -from diffusion_policy.model.common.lr_scheduler import get_scheduler -from diffusers.training_utils import EMAModel - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -# %% -class TrainDiffusionUnetLowdimWorkspace(BaseWorkspace): - include_keys = ['global_step', 'epoch'] - - def __init__(self, cfg: OmegaConf, output_dir=None): - super().__init__(cfg, output_dir=output_dir) - - # set seed - seed = cfg.training.seed - torch.manual_seed(seed) - np.random.seed(seed) - random.seed(seed) - - # configure model - self.model: DiffusionUnetLowdimPolicy - self.model = hydra.utils.instantiate(cfg.policy) - - self.ema_model: DiffusionUnetLowdimPolicy = None - if cfg.training.use_ema: - self.ema_model = copy.deepcopy(self.model) - - # configure training state - self.optimizer = hydra.utils.instantiate( - cfg.optimizer, params=self.model.parameters()) - - self.global_step = 0 - self.epoch = 0 - - def run(self): - cfg = copy.deepcopy(self.cfg) - - # resume training - if cfg.training.resume: - lastest_ckpt_path = self.get_checkpoint_path() - if lastest_ckpt_path.is_file(): - print(f"Resuming from checkpoint {lastest_ckpt_path}") - self.load_checkpoint(path=lastest_ckpt_path) - - # configure dataset - dataset: BaseLowdimDataset - dataset = hydra.utils.instantiate(cfg.task.dataset) - assert isinstance(dataset, BaseLowdimDataset) - train_dataloader = DataLoader(dataset, **cfg.dataloader) - normalizer = dataset.get_normalizer() - - # configure validation dataset - val_dataset = dataset.get_validation_dataset() - val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader) - - self.model.set_normalizer(normalizer) - if cfg.training.use_ema: - self.ema_model.set_normalizer(normalizer) - - # configure lr scheduler - lr_scheduler = get_scheduler( - cfg.training.lr_scheduler, - optimizer=self.optimizer, - num_warmup_steps=cfg.training.lr_warmup_steps, - num_training_steps=( - len(train_dataloader) * cfg.training.num_epochs) \ - // cfg.training.gradient_accumulate_every, - # pytorch assumes stepping LRScheduler every epoch - # however huggingface diffusers steps it every batch - last_epoch=self.global_step-1 - ) - - # configure ema - ema: EMAModel = None - if cfg.training.use_ema: - ema = hydra.utils.instantiate( - cfg.ema, - model=self.ema_model) - - # configure env runner - env_runner: BaseLowdimRunner - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=self.output_dir) - assert isinstance(env_runner, BaseLowdimRunner) - - # configure logging - wandb_run = wandb.init( - dir=str(self.output_dir), - config=OmegaConf.to_container(cfg, resolve=True), - **cfg.logging - ) - wandb.config.update( - { - "output_dir": self.output_dir, - } - ) - - # configure checkpoint - topk_manager = TopKCheckpointManager( - save_dir=os.path.join(self.output_dir, 'checkpoints'), - **cfg.checkpoint.topk - ) - - # device transfer - device = torch.device(cfg.training.device) - self.model.to(device) - if self.ema_model is not None: - self.ema_model.to(device) - optimizer_to(self.optimizer, device) - - # save batch for sampling - train_sampling_batch = None - - if cfg.training.debug: - cfg.training.num_epochs = 2 - cfg.training.max_train_steps = 3 - cfg.training.max_val_steps = 3 - cfg.training.rollout_every = 1 - cfg.training.checkpoint_every = 1 - cfg.training.val_every = 1 - cfg.training.sample_every = 1 - - # training loop - log_path = os.path.join(self.output_dir, 'logs.json.txt') - with JsonLogger(log_path) as json_logger: - for local_epoch_idx in range(cfg.training.num_epochs): - step_log = dict() - # ========= train for this epoch ========== - train_losses = list() - with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - # device transfer - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - if train_sampling_batch is None: - train_sampling_batch = batch - - # compute loss - raw_loss = self.model.compute_loss(batch) - loss = raw_loss / cfg.training.gradient_accumulate_every - loss.backward() - - # step optimizer - if self.global_step % cfg.training.gradient_accumulate_every == 0: - self.optimizer.step() - self.optimizer.zero_grad() - lr_scheduler.step() - - # update ema - if cfg.training.use_ema: - ema.step(self.model) - - # logging - raw_loss_cpu = raw_loss.item() - tepoch.set_postfix(loss=raw_loss_cpu, refresh=False) - train_losses.append(raw_loss_cpu) - step_log = { - 'train_loss': raw_loss_cpu, - 'global_step': self.global_step, - 'epoch': self.epoch, - 'lr': lr_scheduler.get_last_lr()[0] - } - - is_last_batch = (batch_idx == (len(train_dataloader)-1)) - if not is_last_batch: - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - - if (cfg.training.max_train_steps is not None) \ - and batch_idx >= (cfg.training.max_train_steps-1): - break - - # at the end of each epoch - # replace train_loss with epoch average - train_loss = np.mean(train_losses) - step_log['train_loss'] = train_loss - - # ========= eval for this epoch ========== - policy = self.model - if cfg.training.use_ema: - policy = self.ema_model - policy.eval() - - # run rollout - if (self.epoch % cfg.training.rollout_every) == 0: - runner_log = env_runner.run(policy) - # log all - step_log.update(runner_log) - - # run validation - if (self.epoch % cfg.training.val_every) == 0: - with torch.no_grad(): - val_losses = list() - with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - loss = self.model.compute_loss(batch) - val_losses.append(loss) - if (cfg.training.max_val_steps is not None) \ - and batch_idx >= (cfg.training.max_val_steps-1): - break - if len(val_losses) > 0: - val_loss = torch.mean(torch.tensor(val_losses)).item() - # log epoch average validation loss - step_log['val_loss'] = val_loss - - # run diffusion sampling on a training batch - if (self.epoch % cfg.training.sample_every) == 0: - with torch.no_grad(): - # sample trajectory from training set, and evaluate difference - batch = train_sampling_batch - obs_dict = {'obs': batch['obs']} - gt_action = batch['action'] - - result = policy.predict_action(obs_dict) - if cfg.pred_action_steps_only: - pred_action = result['action'] - start = cfg.n_obs_steps - 1 - end = start + cfg.n_action_steps - gt_action = gt_action[:,start:end] - else: - pred_action = result['action_pred'] - mse = torch.nn.functional.mse_loss(pred_action, gt_action) - # log - step_log['train_action_mse_error'] = mse.item() - # release RAM - del batch - del obs_dict - del gt_action - del result - del pred_action - del mse - - # checkpoint - if (self.epoch % cfg.training.checkpoint_every) == 0: - # checkpointing - if cfg.checkpoint.save_last_ckpt: - self.save_checkpoint() - if cfg.checkpoint.save_last_snapshot: - self.save_snapshot() - - # sanitize metric names - metric_dict = dict() - for key, value in step_log.items(): - new_key = key.replace('/', '_') - metric_dict[new_key] = value - - # We can't copy the last checkpoint here - # since save_checkpoint uses threads. - # therefore at this point the file might have been empty! - topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict) - - if topk_ckpt_path is not None: - self.save_checkpoint(path=topk_ckpt_path) - # ========= eval end for this epoch ========== - policy.train() - - # end of epoch - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - self.epoch += 1 - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), - config_name=pathlib.Path(__file__).stem) -def main(cfg): - workspace = TrainDiffusionUnetLowdimWorkspace(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_video_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_video_workspace.py deleted file mode 100644 index 1c1309236..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_video_workspace.py +++ /dev/null @@ -1,235 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - os.chdir(ROOT_DIR) - -import os -import hydra -import torch -from omegaconf import OmegaConf -import pathlib -from torch.utils.data import DataLoader -import copy -import random -import wandb -import tqdm -import numpy as np - -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.diffusion_unet_video_policy import DiffusionUnetVideoPolicy -from diffusion_policy.dataset.base_dataset import BaseImageDataset -from diffusion_policy.env_runner.base_image_runner import BaseImageRunner -from diffusion_policy.common.checkpoint_util import TopKCheckpointManager -from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to -from diffusion_policy.model.diffusion.ema_model import EMAModel -from diffusion_policy.model.common.lr_scheduler import get_scheduler - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -class TrainDiffusionUnetVideoWorkspace(BaseWorkspace): - include_keys = ['global_step', 'epoch'] - - def __init__(self, cfg: OmegaConf, output_dir=None): - super().__init__(cfg, output_dir=output_dir) - - # set seed - seed = cfg.training.seed - torch.manual_seed(seed) - np.random.seed(seed) - random.seed(seed) - - # configure model - self.model: DiffusionUnetVideoPolicy = hydra.utils.instantiate(cfg.policy) - - self.ema_model: DiffusionUnetVideoPolicy = None - if cfg.training.use_ema: - self.ema_model = copy.deepcopy(self.model) - - # configure training state - self.optimizer = hydra.utils.instantiate( - cfg.optimizer, params=self.model.parameters()) - - # configure training state - self.global_step = 0 - self.epoch = 0 - - def run(self): - cfg = copy.deepcopy(self.cfg) - - # resume training - if cfg.training.resume: - lastest_ckpt_path = self.get_checkpoint_path() - if lastest_ckpt_path.is_file(): - print(f"Resuming from checkpoint {lastest_ckpt_path}") - self.load_checkpoint(path=lastest_ckpt_path) - - # configure dataset - dataset: BaseImageDataset - dataset = hydra.utils.instantiate(cfg.task.dataset) - assert isinstance(dataset, BaseImageDataset) - train_dataloader = DataLoader(dataset, **cfg.dataloader) - normalizer = dataset.get_normalizer() - - self.model.set_normalizer(normalizer) - if cfg.training.use_ema: - self.ema_model.set_normalizer(normalizer) - - # configure lr scheduler - lr_scheduler = get_scheduler( - cfg.training.lr_scheduler, - optimizer=self.optimizer, - num_warmup_steps=cfg.training.lr_warmup_steps, - num_training_steps=( - len(train_dataloader) * cfg.training.num_epochs) \ - // cfg.training.gradient_accumulate_every, - # pytorch assumes stepping LRScheduler every epoch - # however huggingface diffusers steps it every batch - last_epoch=self.global_step-1 - ) - - # configure ema - ema: EMAModel = None - if cfg.training.use_ema: - ema = hydra.utils.instantiate( - cfg.ema, - model=self.ema_model) - - # configure env - env_runner: BaseImageRunner - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=self.output_dir) - assert isinstance(env_runner, BaseImageRunner) - - # configure logging - wandb_run = wandb.init( - dir=str(self.output_dir), - config=OmegaConf.to_container(cfg, resolve=True), - **cfg.logging - ) - wandb.config.update( - { - "output_dir": self.output_dir, - } - ) - - # configure checkpoint - topk_manager = TopKCheckpointManager( - save_dir=os.path.join(self.output_dir, 'checkpoints'), - **cfg.checkpoint.topk - ) - - # device transfer - device = torch.device(cfg.training.device) - self.model.to(device) - if self.ema_model is not None: - self.ema_model.to(device) - optimizer_to(self.optimizer, device) - - # save batch for validation - val_batch = next(iter(train_dataloader)) - - # training loop - for _ in range(cfg.training.num_epochs): - with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch in tepoch: - # device transfer - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - - # compute loss - raw_loss = self.model.compute_loss(batch) - loss = raw_loss / cfg.training.gradient_accumulate_every - loss.backward() - - # step optimizer - if self.global_step % cfg.training.gradient_accumulate_every == 0: - self.optimizer.step() - self.optimizer.zero_grad() - lr_scheduler.step() - - # update ema - if cfg.training.use_ema: - ema.step(self.model) - - # logging - raw_loss_cpu = raw_loss.item() - # tepoch.set_postfix(loss=raw_loss_cpu, refresh=False) - step_log = { - 'train_loss': raw_loss_cpu, - 'global_step': self.global_step, - 'epoch': self.epoch, - 'lr': lr_scheduler.get_last_lr()[0] - } - - # eval - should_eval = self.global_step % cfg.training.eval_every == 0 - if (not cfg.training.eval_first) and (self.global_step == 0): - should_eval = False - if should_eval: - policy = self.model - if cfg.training.use_ema: - policy = self.ema_model - policy.eval() - runner_log = env_runner.run(policy) - policy.train() - step_log.update(runner_log) - - # checkpointing - if cfg.checkpoint.save_last_ckpt: - self.save_checkpoint() - if cfg.checkpoint.save_last_snapshot: - self.save_snapshot() - - save_path = topk_manager.get_ckpt_path({ - 'epoch': self.epoch, - 'test_score': runner_log['test/mean_score'] - }) - if save_path is not None: - self.save_checkpoint(path=save_path) - - # validation - should_val = self.global_step % cfg.training.val_every == 0 - if should_val: - policy = self.model - if cfg.training.use_ema: - policy = self.ema_model - policy.eval() - with torch.no_grad(): - # sample trajectory from training set, and evaluate difference - batch = dict_apply(val_batch, lambda x: x.to(device, non_blocking=True)) - obs_dict = batch['obs'] - gt_action = batch['action'] - - result = policy.predict_action(obs_dict) - pred_action = result['action_pred'] - mse = torch.nn.functional.mse_loss(pred_action, gt_action) - step_log['train_action_mse_error'] = mse.item() - del batch - del obs_dict - del gt_action - del result - del pred_action - del mse - policy.train() - - # log - wandb_run.log(step_log, step=self.global_step) - - self.global_step += 1 - self.epoch += 1 - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), - config_name=pathlib.Path(__file__).stem) -def main(cfg): - workspace = TrainDiffusionUnetVideoWorkspace(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_ibc_dfo_hybrid_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_ibc_dfo_hybrid_workspace.py deleted file mode 100644 index 1370897ed..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_ibc_dfo_hybrid_workspace.py +++ /dev/null @@ -1,283 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - os.chdir(ROOT_DIR) - -import os -import hydra -import torch -from omegaconf import OmegaConf -import pathlib -from torch.utils.data import DataLoader -import copy -import random -import wandb -import tqdm -import numpy as np -import shutil -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.ibc_dfo_hybrid_image_policy import IbcDfoHybridImagePolicy -from diffusion_policy.dataset.base_dataset import BaseImageDataset -from diffusion_policy.env_runner.base_image_runner import BaseImageRunner -from diffusion_policy.common.checkpoint_util import TopKCheckpointManager -from diffusion_policy.common.json_logger import JsonLogger -from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to -from diffusion_policy.model.diffusion.ema_model import EMAModel -from diffusion_policy.model.common.lr_scheduler import get_scheduler - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -class TrainIbcDfoHybridWorkspace(BaseWorkspace): - include_keys = ['global_step', 'epoch'] - - def __init__(self, cfg: OmegaConf, output_dir=None): - super().__init__(cfg, output_dir=output_dir) - - # set seed - seed = cfg.training.seed - torch.manual_seed(seed) - np.random.seed(seed) - random.seed(seed) - - # configure model - self.model: IbcDfoHybridImagePolicy= hydra.utils.instantiate(cfg.policy) - - # configure training state - self.optimizer = hydra.utils.instantiate( - cfg.optimizer, params=self.model.parameters()) - - # configure training state - self.global_step = 0 - self.epoch = 0 - - def run(self): - cfg = copy.deepcopy(self.cfg) - - # resume training - if cfg.training.resume: - lastest_ckpt_path = self.get_checkpoint_path() - if lastest_ckpt_path.is_file(): - print(f"Resuming from checkpoint {lastest_ckpt_path}") - self.load_checkpoint(path=lastest_ckpt_path) - - # configure dataset - dataset: BaseImageDataset - dataset = hydra.utils.instantiate(cfg.task.dataset) - assert isinstance(dataset, BaseImageDataset) - train_dataloader = DataLoader(dataset, **cfg.dataloader) - normalizer = dataset.get_normalizer() - - # configure validation dataset - val_dataset = dataset.get_validation_dataset() - val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader) - - self.model.set_normalizer(normalizer) - - # configure lr scheduler - lr_scheduler = get_scheduler( - cfg.training.lr_scheduler, - optimizer=self.optimizer, - num_warmup_steps=cfg.training.lr_warmup_steps, - num_training_steps=( - len(train_dataloader) * cfg.training.num_epochs) \ - // cfg.training.gradient_accumulate_every, - # pytorch assumes stepping LRScheduler every epoch - # however huggingface diffusers steps it every batch - last_epoch=self.global_step-1 - ) - - # configure env - env_runner: BaseImageRunner - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=self.output_dir) - assert isinstance(env_runner, BaseImageRunner) - - # configure logging - wandb_run = wandb.init( - dir=str(self.output_dir), - config=OmegaConf.to_container(cfg, resolve=True), - **cfg.logging - ) - wandb.config.update( - { - "output_dir": self.output_dir, - } - ) - - # configure checkpoint - topk_manager = TopKCheckpointManager( - save_dir=os.path.join(self.output_dir, 'checkpoints'), - **cfg.checkpoint.topk - ) - - # device transfer - device = torch.device(cfg.training.device) - self.model.to(device) - optimizer_to(self.optimizer, device) - - # save batch for sampling - train_sampling_batch = None - - if cfg.training.debug: - cfg.training.num_epochs = 2 - cfg.training.max_train_steps = 3 - cfg.training.max_val_steps = 3 - cfg.training.rollout_every = 1 - cfg.training.checkpoint_every = 1 - cfg.training.val_every = 1 - cfg.training.sample_every = 1 - - # training loop - log_path = os.path.join(self.output_dir, 'logs.json.txt') - with JsonLogger(log_path) as json_logger: - for local_epoch_idx in range(cfg.training.num_epochs): - step_log = dict() - # ========= train for this epoch ========== - train_losses = list() - with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - # device transfer - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - if train_sampling_batch is None: - train_sampling_batch = batch - - # compute loss - raw_loss = self.model.compute_loss(batch) - loss = raw_loss / cfg.training.gradient_accumulate_every - loss.backward() - - # step optimizer - if self.global_step % cfg.training.gradient_accumulate_every == 0: - self.optimizer.step() - self.optimizer.zero_grad() - lr_scheduler.step() - - # logging - raw_loss_cpu = raw_loss.item() - tepoch.set_postfix(loss=raw_loss_cpu, refresh=False) - train_losses.append(raw_loss_cpu) - step_log = { - 'train_loss': raw_loss_cpu, - 'global_step': self.global_step, - 'epoch': self.epoch, - 'lr': lr_scheduler.get_last_lr()[0] - } - - is_last_batch = (batch_idx == (len(train_dataloader)-1)) - if not is_last_batch: - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - - if (cfg.training.max_train_steps is not None) \ - and batch_idx >= (cfg.training.max_train_steps-1): - break - - # at the end of each epoch - # replace train_loss with epoch average - train_loss = np.mean(train_losses) - step_log['train_loss'] = train_loss - - # ========= eval for this epoch ========== - policy = self.model - policy.eval() - - # run rollout - if (self.epoch % cfg.training.rollout_every) == 0: - runner_log = env_runner.run(policy) - # log all - step_log.update(runner_log) - - # run validation - if (self.epoch % cfg.training.val_every) == 0: - with torch.no_grad(): - val_losses = list() - with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - loss = self.model.compute_loss(batch) - val_losses.append(loss) - if (cfg.training.max_val_steps is not None) \ - and batch_idx >= (cfg.training.max_val_steps-1): - break - if len(val_losses) > 0: - val_loss = torch.mean(torch.tensor(val_losses)).item() - # log epoch average validation loss - step_log['val_loss'] = val_loss - - # run diffusion sampling on a training batch - if (self.epoch % cfg.training.sample_every) == 0: - with torch.no_grad(): - # sample trajectory from training set, and evaluate difference - batch = train_sampling_batch - n_samples = cfg.training.sample_max_batch - batch = dict_apply(train_sampling_batch, - lambda x: x.to(device, non_blocking=True)) - obs_dict = dict_apply(batch['obs'], lambda x: x[:n_samples]) - gt_action = batch['action'] - - result = policy.predict_action(obs_dict) - pred_action = result['action'] - start = cfg.n_obs_steps - 1 - end = start + cfg.n_action_steps - gt_action = gt_action[:,start:end] - mse = torch.nn.functional.mse_loss(pred_action, gt_action) - # log - step_log['train_action_mse_error'] = mse.item() - # release RAM - del batch - del obs_dict - del gt_action - del result - del pred_action - del mse - - # checkpoint - if (self.epoch % cfg.training.checkpoint_every) == 0: - # checkpointing - if cfg.checkpoint.save_last_ckpt: - self.save_checkpoint() - if cfg.checkpoint.save_last_snapshot: - self.save_snapshot() - - # sanitize metric names - metric_dict = dict() - for key, value in step_log.items(): - new_key = key.replace('/', '_') - metric_dict[new_key] = value - - # We can't copy the last checkpoint here - # since save_checkpoint uses threads. - # therefore at this point the file might have been empty! - topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict) - - if topk_ckpt_path is not None: - self.save_checkpoint(path=topk_ckpt_path) - # ========= eval end for this epoch ========== - policy.train() - - # end of epoch - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - self.epoch += 1 - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), - config_name=pathlib.Path(__file__).stem) -def main(cfg): - workspace = TrainIbcDfoHybridWorkspace(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_ibc_dfo_lowdim_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_ibc_dfo_lowdim_workspace.py deleted file mode 100644 index fad9c6759..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_ibc_dfo_lowdim_workspace.py +++ /dev/null @@ -1,282 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - os.chdir(ROOT_DIR) - -import os -import hydra -import torch -from omegaconf import OmegaConf -import pathlib -from torch.utils.data import DataLoader -import copy -import numpy as np -import random -import wandb -import tqdm -import shutil - -from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.ibc_dfo_lowdim_policy import IbcDfoLowdimPolicy -from diffusion_policy.dataset.base_dataset import BaseLowdimDataset -from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner -from diffusion_policy.common.checkpoint_util import TopKCheckpointManager -from diffusion_policy.common.json_logger import JsonLogger -from diffusion_policy.model.common.lr_scheduler import get_scheduler - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -# %% -class TrainIbcDfoLowdimWorkspace(BaseWorkspace): - include_keys = ['global_step', 'epoch'] - - def __init__(self, cfg: OmegaConf, output_dir=None): - super().__init__(cfg, output_dir=output_dir) - - # set seed - seed = cfg.training.seed - torch.manual_seed(seed) - np.random.seed(seed) - random.seed(seed) - - # configure model - self.model: IbcDfoLowdimPolicy - self.model = hydra.utils.instantiate(cfg.policy) - - # configure training state - self.optimizer = hydra.utils.instantiate( - cfg.optimizer, params=self.model.parameters()) - - self.global_step = 0 - self.epoch = 0 - - def run(self): - cfg = copy.deepcopy(self.cfg) - - # resume training - if cfg.training.resume: - lastest_ckpt_path = self.get_checkpoint_path() - if lastest_ckpt_path.is_file(): - print(f"Resuming from checkpoint {lastest_ckpt_path}") - self.load_checkpoint(path=lastest_ckpt_path) - - # configure dataset - dataset: BaseLowdimDataset - dataset = hydra.utils.instantiate(cfg.task.dataset) - assert isinstance(dataset, BaseLowdimDataset) - train_dataloader = DataLoader(dataset, **cfg.dataloader) - normalizer = dataset.get_normalizer() - - # configure validation dataset - val_dataset = dataset.get_validation_dataset() - val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader) - - self.model.set_normalizer(normalizer) - - # configure lr scheduler - lr_scheduler = get_scheduler( - cfg.training.lr_scheduler, - optimizer=self.optimizer, - num_warmup_steps=cfg.training.lr_warmup_steps, - num_training_steps=( - len(train_dataloader) * cfg.training.num_epochs) \ - // cfg.training.gradient_accumulate_every, - # pytorch assumes stepping LRScheduler every epoch - # however huggingface diffusers steps it every batch - last_epoch=self.global_step-1 - ) - - # configure env runner - env_runner: BaseLowdimRunner - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=self.output_dir) - assert isinstance(env_runner, BaseLowdimRunner) - - # configure logging - wandb_run = wandb.init( - dir=str(self.output_dir), - config=OmegaConf.to_container(cfg, resolve=True), - **cfg.logging - ) - wandb.config.update( - { - "output_dir": self.output_dir, - } - ) - - # configure checkpoint - topk_manager = TopKCheckpointManager( - save_dir=os.path.join(self.output_dir, 'checkpoints'), - **cfg.checkpoint.topk - ) - - # device transfer - device = torch.device(cfg.training.device) - self.model.to(device) - optimizer_to(self.optimizer, device) - - # save batch for sampling - train_sampling_batch = None - - if cfg.training.debug: - cfg.training.num_epochs = 2 - cfg.training.max_train_steps = 3 - cfg.training.max_val_steps = 3 - cfg.training.rollout_every = 1 - cfg.training.checkpoint_every = 1 - cfg.training.val_every = 1 - cfg.training.sample_every = 1 - - # training loop - log_path = os.path.join(self.output_dir, 'logs.json.txt') - with JsonLogger(log_path) as json_logger: - for local_epoch_idx in range(cfg.training.num_epochs): - step_log = dict() - # ========= train for this epoch ========== - train_losses = list() - with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - # device transfer - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - if train_sampling_batch is None: - train_sampling_batch = batch - - # compute loss - raw_loss = self.model.compute_loss(batch) - loss = raw_loss / cfg.training.gradient_accumulate_every - loss.backward() - - # step optimizer - if self.global_step % cfg.training.gradient_accumulate_every == 0: - self.optimizer.step() - self.optimizer.zero_grad() - lr_scheduler.step() - - # logging - raw_loss_cpu = raw_loss.item() - tepoch.set_postfix(loss=raw_loss_cpu, refresh=False) - train_losses.append(raw_loss_cpu) - step_log = { - 'train_loss': raw_loss_cpu, - 'global_step': self.global_step, - 'epoch': self.epoch, - 'lr': lr_scheduler.get_last_lr()[0] - } - - is_last_batch = (batch_idx == (len(train_dataloader)-1)) - if not is_last_batch: - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - - if (cfg.training.max_train_steps is not None) \ - and batch_idx >= (cfg.training.max_train_steps-1): - break - - # at the end of each epoch - # replace train_loss with epoch average - train_loss = np.mean(train_losses) - step_log['train_loss'] = train_loss - - # ========= eval for this epoch ========== - policy = self.model - policy.eval() - - # run rollout - if (self.epoch % cfg.training.rollout_every) == 0: - runner_log = env_runner.run(policy) - # log all - step_log.update(runner_log) - - # run validation - if (self.epoch % cfg.training.val_every) == 0: - with torch.no_grad(): - val_losses = list() - with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - loss = self.model.compute_loss(batch) - val_losses.append(loss) - if (cfg.training.max_val_steps is not None) \ - and batch_idx >= (cfg.training.max_val_steps-1): - break - if len(val_losses) > 0: - val_loss = torch.mean(torch.tensor(val_losses)).item() - # log epoch average validation loss - step_log['val_loss'] = val_loss - - # run diffusion sampling on a training batch - if (self.epoch % cfg.training.sample_every) == 0: - with torch.no_grad(): - # sample trajectory from training set, and evaluate difference - batch = train_sampling_batch - n_samples = cfg.training.sample_max_batch - obs_dict = {'obs': batch['obs'][:n_samples]} - gt_action = batch['action'][:n_samples] - - result = policy.predict_action(obs_dict) - pred_action = result['action'] - start = cfg.n_obs_steps - 1 - end = start + cfg.n_action_steps - gt_action = gt_action[:,start:end] - mse = torch.nn.functional.mse_loss(pred_action, gt_action) - # log - step_log['train_action_mse_error'] = mse.item() - # release RAM - del batch - del obs_dict - del gt_action - del result - del pred_action - del mse - - # checkpoint - if (self.epoch % cfg.training.checkpoint_every) == 0: - # checkpointing - if cfg.checkpoint.save_last_ckpt: - self.save_checkpoint() - if cfg.checkpoint.save_last_snapshot: - self.save_snapshot() - - # sanitize metric names - metric_dict = dict() - for key, value in step_log.items(): - new_key = key.replace('/', '_') - metric_dict[new_key] = value - - # We can't copy the last checkpoint here - # since save_checkpoint uses threads. - # therefore at this point the file might have been empty! - topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict) - - if topk_ckpt_path is not None: - self.save_checkpoint(path=topk_ckpt_path) - # ========= eval end for this epoch ========== - policy.train() - - # end of epoch - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - self.epoch += 1 - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), - config_name=pathlib.Path(__file__).stem) -def main(cfg): - workspace = TrainIbcDfoLowdimWorkspace(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_robomimic_image_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_robomimic_image_workspace.py deleted file mode 100644 index 2217aee48..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_robomimic_image_workspace.py +++ /dev/null @@ -1,254 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - os.chdir(ROOT_DIR) - -import os -import hydra -import torch -from omegaconf import OmegaConf -import pathlib -from torch.utils.data import DataLoader -import copy -import random -import wandb -import tqdm -import numpy as np -import shutil -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.robomimic_image_policy import RobomimicImagePolicy -from diffusion_policy.dataset.base_dataset import BaseImageDataset -from diffusion_policy.env_runner.base_image_runner import BaseImageRunner -from diffusion_policy.common.checkpoint_util import TopKCheckpointManager -from diffusion_policy.common.json_logger import JsonLogger -from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to - - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -class TrainRobomimicImageWorkspace(BaseWorkspace): - include_keys = ['global_step', 'epoch'] - - def __init__(self, cfg: OmegaConf, output_dir=None): - super().__init__(cfg, output_dir=output_dir) - - # set seed - seed = cfg.training.seed - torch.manual_seed(seed) - np.random.seed(seed) - random.seed(seed) - - # configure model - self.model: RobomimicImagePolicy = hydra.utils.instantiate(cfg.policy) - - # configure training state - self.global_step = 0 - self.epoch = 0 - - def run(self): - cfg = copy.deepcopy(self.cfg) - - # resume training - if cfg.training.resume: - lastest_ckpt_path = self.get_checkpoint_path() - if lastest_ckpt_path.is_file(): - print(f"Resuming from checkpoint {lastest_ckpt_path}") - self.load_checkpoint(path=lastest_ckpt_path) - - # configure dataset - dataset: BaseImageDataset - dataset = hydra.utils.instantiate(cfg.task.dataset) - assert isinstance(dataset, BaseImageDataset) - train_dataloader = DataLoader(dataset, **cfg.dataloader) - normalizer = dataset.get_normalizer() - - # configure validation dataset - val_dataset = dataset.get_validation_dataset() - val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader) - - self.model.set_normalizer(normalizer) - - # configure env - env_runner: BaseImageRunner - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=self.output_dir) - assert isinstance(env_runner, BaseImageRunner) - - # configure logging - wandb_run = wandb.init( - dir=str(self.output_dir), - config=OmegaConf.to_container(cfg, resolve=True), - **cfg.logging - ) - wandb.config.update( - { - "output_dir": self.output_dir, - } - ) - - # configure checkpoint - topk_manager = TopKCheckpointManager( - save_dir=os.path.join(self.output_dir, 'checkpoints'), - **cfg.checkpoint.topk - ) - - # device transfer - device = torch.device(cfg.training.device) - self.model.to(device) - - # save batch for sampling - train_sampling_batch = None - - if cfg.training.debug: - cfg.training.num_epochs = 2 - cfg.training.max_train_steps = 3 - cfg.training.max_val_steps = 3 - cfg.training.rollout_every = 1 - cfg.training.checkpoint_every = 1 - cfg.training.val_every = 1 - cfg.training.sample_every = 1 - - # training loop - log_path = os.path.join(self.output_dir, 'logs.json.txt') - with JsonLogger(log_path) as json_logger: - for local_epoch_idx in range(cfg.training.num_epochs): - step_log = dict() - # ========= train for this epoch ========== - train_losses = list() - with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - # device transfer - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - if train_sampling_batch is None: - train_sampling_batch = batch - - info = self.model.train_on_batch(batch, epoch=self.epoch) - - # logging - loss_cpu = info['losses']['action_loss'].item() - tepoch.set_postfix(loss=loss_cpu, refresh=False) - train_losses.append(loss_cpu) - step_log = { - 'train_loss': loss_cpu, - 'global_step': self.global_step, - 'epoch': self.epoch - } - - is_last_batch = (batch_idx == (len(train_dataloader)-1)) - if not is_last_batch: - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - - if (cfg.training.max_train_steps is not None) \ - and batch_idx >= (cfg.training.max_train_steps-1): - break - - # at the end of each epoch - # replace train_loss with epoch average - train_loss = np.mean(train_losses) - step_log['train_loss'] = train_loss - - # ========= eval for this epoch ========== - self.model.eval() - - # run rollout - if (self.epoch % cfg.training.rollout_every) == 0: - runner_log = env_runner.run(self.model) - # log all - step_log.update(runner_log) - - # run validation - if (self.epoch % cfg.training.val_every) == 0: - with torch.no_grad(): - val_losses = list() - with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - info = self.model.train_on_batch(batch, epoch=self.epoch, validate=True) - loss = info['losses']['action_loss'] - val_losses.append(loss) - if (cfg.training.max_val_steps is not None) \ - and batch_idx >= (cfg.training.max_val_steps-1): - break - if len(val_losses) > 0: - val_loss = torch.mean(torch.tensor(val_losses)).item() - # log epoch average validation loss - step_log['val_loss'] = val_loss - - # run diffusion sampling on a training batch - if (self.epoch % cfg.training.sample_every) == 0: - with torch.no_grad(): - # sample trajectory from training set, and evaluate difference - batch = dict_apply(train_sampling_batch, lambda x: x.to(device, non_blocking=True)) - obs_dict = batch['obs'] - gt_action = batch['action'] - T = gt_action.shape[1] - - pred_actions = list() - self.model.reset() - for i in range(T): - result = self.model.predict_action( - dict_apply(obs_dict, lambda x: x[:,[i]]) - ) - pred_actions.append(result['action']) - pred_actions = torch.cat(pred_actions, dim=1) - mse = torch.nn.functional.mse_loss(pred_actions, gt_action) - step_log['train_action_mse_error'] = mse.item() - del batch - del obs_dict - del gt_action - del result - del pred_actions - del mse - - # checkpoint - if (self.epoch % cfg.training.checkpoint_every) == 0: - # checkpointing - if cfg.checkpoint.save_last_ckpt: - self.save_checkpoint() - if cfg.checkpoint.save_last_snapshot: - self.save_snapshot() - - # sanitize metric names - metric_dict = dict() - for key, value in step_log.items(): - new_key = key.replace('/', '_') - metric_dict[new_key] = value - - # We can't copy the last checkpoint here - # since save_checkpoint uses threads. - # therefore at this point the file might have been empty! - topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict) - - if topk_ckpt_path is not None: - self.save_checkpoint(path=topk_ckpt_path) - # ========= eval end for this epoch ========== - self.model.train() - - # end of epoch - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - self.epoch += 1 - - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), - config_name=pathlib.Path(__file__).stem) -def main(cfg): - workspace = TrainRobomimicImageWorkspace(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_robomimic_lowdim_workspace.py b/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_robomimic_lowdim_workspace.py deleted file mode 100644 index d39f1c13d..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/diffusion_policy/workspace/train_robomimic_lowdim_workspace.py +++ /dev/null @@ -1,221 +0,0 @@ -if __name__ == "__main__": - import sys - import os - import pathlib - - ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) - sys.path.append(ROOT_DIR) - os.chdir(ROOT_DIR) - -import os -import hydra -import torch -from omegaconf import OmegaConf -import pathlib -from torch.utils.data import DataLoader -import copy -import random -import wandb -import tqdm -import numpy as np -import shutil -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.robomimic_lowdim_policy import RobomimicLowdimPolicy -from diffusion_policy.dataset.base_dataset import BaseLowdimDataset -from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner -from diffusion_policy.common.checkpoint_util import TopKCheckpointManager -from diffusion_policy.common.json_logger import JsonLogger -from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to - - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -class TrainRobomimicLowdimWorkspace(BaseWorkspace): - include_keys = ['global_step', 'epoch'] - - def __init__(self, cfg: OmegaConf): - super().__init__(cfg) - - # set seed - seed = cfg.training.seed - torch.manual_seed(seed) - np.random.seed(seed) - random.seed(seed) - - # configure model - self.model: RobomimicLowdimPolicy = hydra.utils.instantiate(cfg.policy) - - # configure training state - self.global_step = 0 - self.epoch = 0 - - def run(self): - cfg = copy.deepcopy(self.cfg) - - # resume training - if cfg.training.resume: - lastest_ckpt_path = self.get_checkpoint_path() - if lastest_ckpt_path.is_file(): - print(f"Resuming from checkpoint {lastest_ckpt_path}") - self.load_checkpoint(path=lastest_ckpt_path) - - # configure dataset - dataset: BaseLowdimDataset - dataset = hydra.utils.instantiate(cfg.task.dataset) - assert isinstance(dataset, BaseLowdimDataset) - train_dataloader = DataLoader(dataset, **cfg.dataloader) - normalizer = dataset.get_normalizer() - - # configure validation dataset - val_dataset = dataset.get_validation_dataset() - val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader) - - self.model.set_normalizer(normalizer) - - # configure env - env_runner: BaseLowdimRunner - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=self.output_dir) - assert isinstance(env_runner, BaseLowdimRunner) - - # configure logging - wandb_run = wandb.init( - dir=str(self.output_dir), - config=OmegaConf.to_container(cfg, resolve=True), - **cfg.logging - ) - wandb.config.update( - { - "output_dir": self.output_dir, - } - ) - - # configure checkpoint - topk_manager = TopKCheckpointManager( - save_dir=os.path.join(self.output_dir, 'checkpoints'), - **cfg.checkpoint.topk - ) - - # device transfer - device = torch.device(cfg.training.device) - self.model.to(device) - - if cfg.training.debug: - cfg.training.num_epochs = 2 - cfg.training.max_train_steps = 3 - cfg.training.max_val_steps = 3 - cfg.training.rollout_every = 1 - cfg.training.checkpoint_every = 1 - cfg.training.val_every = 1 - - # training loop - log_path = os.path.join(self.output_dir, 'logs.json.txt') - with JsonLogger(log_path) as json_logger: - for local_epoch_idx in range(cfg.training.num_epochs): - step_log = dict() - # ========= train for this epoch ========== - train_losses = list() - with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - # device transfer - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - info = self.model.train_on_batch(batch, epoch=self.epoch) - - # logging - loss_cpu = info['losses']['action_loss'].item() - tepoch.set_postfix(loss=loss_cpu, refresh=False) - train_losses.append(loss_cpu) - step_log = { - 'train_loss': loss_cpu, - 'global_step': self.global_step, - 'epoch': self.epoch - } - - is_last_batch = (batch_idx == (len(train_dataloader)-1)) - if not is_last_batch: - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - - if (cfg.training.max_train_steps is not None) \ - and batch_idx >= (cfg.training.max_train_steps-1): - break - - # at the end of each epoch - # replace train_loss with epoch average - train_loss = np.mean(train_losses) - step_log['train_loss'] = train_loss - - # ========= eval for this epoch ========== - self.model.eval() - - # run rollout - if (self.epoch % cfg.training.rollout_every) == 0: - runner_log = env_runner.run(self.model) - # log all - step_log.update(runner_log) - - # run validation - if (self.epoch % cfg.training.val_every) == 0: - with torch.no_grad(): - val_losses = list() - with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", - leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch: - for batch_idx, batch in enumerate(tepoch): - batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True)) - info = self.model.train_on_batch(batch, epoch=self.epoch, validate=True) - loss = info['losses']['action_loss'] - val_losses.append(loss) - if (cfg.training.max_val_steps is not None) \ - and batch_idx >= (cfg.training.max_val_steps-1): - break - if len(val_losses) > 0: - val_loss = torch.mean(torch.tensor(val_losses)).item() - # log epoch average validation loss - step_log['val_loss'] = val_loss - - # checkpoint - if (self.epoch % cfg.training.checkpoint_every) == 0: - # checkpointing - if cfg.checkpoint.save_last_ckpt: - self.save_checkpoint() - if cfg.checkpoint.save_last_snapshot: - self.save_snapshot() - - # sanitize metric names - metric_dict = dict() - for key, value in step_log.items(): - new_key = key.replace('/', '_') - metric_dict[new_key] = value - - # We can't copy the last checkpoint here - # since save_checkpoint uses threads. - # therefore at this point the file might have been empty! - topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict) - - if topk_ckpt_path is not None: - self.save_checkpoint(path=topk_ckpt_path) - # ========= eval end for this epoch ========== - self.model.train() - - # end of epoch - # log of last step is combined with validation and rollout - wandb_run.log(step_log, step=self.global_step) - json_logger.log(step_log) - self.global_step += 1 - self.epoch += 1 - - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), - config_name=pathlib.Path(__file__).stem) -def main(cfg): - workspace = TrainRobomimicLowdimWorkspace(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/env.yaml b/notebooks_py/diffusion_/diffusion_policy/env.yaml deleted file mode 100644 index c3f55ede5..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/env.yaml +++ /dev/null @@ -1,227 +0,0 @@ -name: sam -channels: -- conda-forge -- nvidia -- pytorch -dependencies: -- _libgcc_mutex=0.1=conda_forge -- _openmp_mutex=4.5=2_kmp_llvm -- anyio=4.4.0=pyhd8ed1ab_0 -- argon2-cffi=23.1.0=pyhd8ed1ab_0 -- argon2-cffi-bindings=21.2.0=py310ha75aee5_5 -- arrow=1.3.0=pyhd8ed1ab_0 -- asttokens=2.4.1=pyhd8ed1ab_0 -- async-lru=2.0.4=pyhd8ed1ab_0 -- attrs=24.2.0=pyh71513ae_0 -- babel=2.14.0=pyhd8ed1ab_0 -- beautifulsoup4=4.12.3=pyha770c72_0 -- blas=2.116=mkl -- blas-devel=3.9.0=16_linux64_mkl -- bleach=6.1.0=pyhd8ed1ab_0 -- brotli-python=1.1.0=py310hf71b8c6_2 -- bzip2=1.0.8=h4bc722e_7 -- ca-certificates=2024.8.30=hbcca054_0 -- cached-property=1.5.2=hd8ed1ab_1 -- cached_property=1.5.2=pyha770c72_1 -- certifi=2024.8.30=pyhd8ed1ab_0 -- cffi=1.17.0=py310h8deb56e_1 -- charset-normalizer=3.3.2=pyhd8ed1ab_0 -- comm=0.2.2=pyhd8ed1ab_0 -- cuda-cudart=11.8.89=0 -- cuda-cupti=11.8.87=0 -- cuda-libraries=11.8.0=0 -- cuda-nvrtc=11.8.89=0 -- cuda-nvtx=11.8.86=0 -- cuda-runtime=11.8.0=0 -- cuda-version=12.6=3 -- debugpy=1.8.5=py310hf71b8c6_1 -- decorator=5.1.1=pyhd8ed1ab_0 -- defusedxml=0.7.1=pyhd8ed1ab_0 -- entrypoints=0.4=pyhd8ed1ab_0 -- exceptiongroup=1.2.2=pyhd8ed1ab_0 -- executing=2.1.0=pyhd8ed1ab_0 -- ffmpeg=4.3=hf484d3e_0 -- filelock=3.15.4=pyhd8ed1ab_0 -- fqdn=1.5.1=pyhd8ed1ab_0 -- freetype=2.12.1=h267a509_2 -- gmp=6.3.0=hac33072_2 -- gmpy2=2.1.5=py310he8512ff_2 -- gnutls=3.6.13=h85f3911_1 -- h11=0.14.0=pyhd8ed1ab_0 -- h2=4.1.0=pyhd8ed1ab_0 -- hpack=4.0.0=pyh9f0ad1d_0 -- httpcore=1.0.5=pyhd8ed1ab_0 -- httpx=0.27.2=pyhd8ed1ab_0 -- hyperframe=6.0.1=pyhd8ed1ab_0 -- icu=73.2=h59595ed_0 -- idna=3.8=pyhd8ed1ab_0 -- importlib-metadata=8.4.0=pyha770c72_0 -- importlib_metadata=8.4.0=hd8ed1ab_0 -- importlib_resources=6.4.4=pyhd8ed1ab_0 -- ipykernel=6.29.5=pyh3099207_0 -- ipython=8.27.0=pyh707e725_0 -- ipywidgets=8.1.5=pyhd8ed1ab_0 -- isoduration=20.11.0=pyhd8ed1ab_0 -- jedi=0.19.1=pyhd8ed1ab_0 -- jinja2=3.1.4=pyhd8ed1ab_0 -- jpeg=9e=h166bdaf_2 -- json5=0.9.25=pyhd8ed1ab_0 -- jsonpointer=3.0.0=py310hff52083_1 -- jsonschema=4.23.0=pyhd8ed1ab_0 -- jsonschema-specifications=2023.12.1=pyhd8ed1ab_0 -- jsonschema-with-format-nongpl=4.23.0=hd8ed1ab_0 -- jupyter=1.1.1=pyhd8ed1ab_0 -- jupyter-lsp=2.2.5=pyhd8ed1ab_0 -- jupyter_client=8.6.2=pyhd8ed1ab_0 -- jupyter_console=6.6.3=pyhd8ed1ab_0 -- jupyter_core=5.7.2=py310hff52083_0 -- jupyter_events=0.10.0=pyhd8ed1ab_0 -- jupyter_server=2.14.2=pyhd8ed1ab_0 -- jupyter_server_terminals=0.5.3=pyhd8ed1ab_0 -- jupyterlab=4.2.5=pyhd8ed1ab_0 -- jupyterlab_pygments=0.3.0=pyhd8ed1ab_1 -- jupyterlab_server=2.27.3=pyhd8ed1ab_0 -- jupyterlab_widgets=3.0.13=pyhd8ed1ab_0 -- keyutils=1.6.1=h166bdaf_0 -- krb5=1.21.3=h659f571_0 -- lame=3.100=h166bdaf_1003 -- lcms2=2.15=hfd0df8a_0 -- ld_impl_linux-64=2.40=hf3520f5_7 -- lerc=4.0.0=h27087fc_0 -- libblas=3.9.0=16_linux64_mkl -- libcblas=3.9.0=16_linux64_mkl -- libcublas=11.11.3.6=0 -- libcufft=10.9.0.58=0 -- libcufile=1.11.1.6=0 -- libcurand=10.3.7.68=0 -- libcusolver=11.4.1.48=0 -- libcusparse=11.7.5.86=0 -- libdeflate=1.17=h0b41bf4_0 -- libedit=3.1.20191231=he28a2e2_2 -- libffi=3.4.2=h7f98852_5 -- libgcc=14.1.0=h77fa898_1 -- libgcc-ng=14.1.0=h69a702a_1 -- libgfortran=14.1.0=h69a702a_1 -- libgfortran-ng=14.1.0=h69a702a_1 -- libgfortran5=14.1.0=hc5f4f2c_1 -- libgomp=14.1.0=h77fa898_1 -- libhwloc=2.11.1=default_hecaa2ac_1000 -- libiconv=1.17=hd590300_2 -- libjpeg-turbo=2.0.0=h9bf148f_0 -- liblapack=3.9.0=16_linux64_mkl -- liblapacke=3.9.0=16_linux64_mkl -- libnpp=11.8.0.86=0 -- libnsl=2.0.1=hd590300_0 -- libnvjpeg=11.9.0.86=0 -- libpng=1.6.43=h2797004_0 -- libsodium=1.0.18=h36c2ea0_1 -- libsqlite=3.46.0=hde9e2c9_0 -- libstdcxx=14.1.0=hc0a3c3a_1 -- libstdcxx-ng=14.1.0=h4852527_1 -- libtiff=4.5.0=h6adf6a1_2 -- libuuid=2.38.1=h0b41bf4_0 -- libwebp-base=1.4.0=hd590300_0 -- libxcb=1.13=h7f98852_1004 -- libxcrypt=4.4.36=hd590300_1 -- libxml2=2.12.7=hc051c1a_1 -- libzlib=1.2.13=h4ab18f5_6 -- llvm-openmp=15.0.7=h0cdce71_0 -- markupsafe=2.1.5=py310ha75aee5_1 -- matplotlib-inline=0.1.7=pyhd8ed1ab_0 -- mistune=3.0.2=pyhd8ed1ab_0 -- mkl=2022.1.0=h84fe81f_915 -- mkl-devel=2022.1.0=ha770c72_916 -- mkl-include=2022.1.0=h84fe81f_915 -- mpc=1.3.1=h24ddda3_0 -- mpfr=4.2.1=h38ae2d0_2 -- mpmath=1.3.0=pyhd8ed1ab_0 -- nbclient=0.10.0=pyhd8ed1ab_0 -- nbconvert-core=7.16.4=pyhd8ed1ab_1 -- nbformat=5.10.4=pyhd8ed1ab_0 -- ncurses=6.5=he02047a_1 -- nest-asyncio=1.6.0=pyhd8ed1ab_0 -- nettle=3.6=he412f7d_0 -- networkx=3.3=pyhd8ed1ab_1 -- notebook=7.2.2=pyhd8ed1ab_0 -- notebook-shim=0.2.4=pyhd8ed1ab_0 -- numpy=2.1.0=py310hd6e36ab_1 -- openh264=2.1.1=h780b84a_0 -- openjpeg=2.5.0=hfec8fc6_2 -- openssl=3.3.1=hb9d3cd8_3 -- overrides=7.7.0=pyhd8ed1ab_0 -- packaging=24.1=pyhd8ed1ab_0 -- pandocfilters=1.5.0=pyhd8ed1ab_0 -- parso=0.8.4=pyhd8ed1ab_0 -- pexpect=4.9.0=pyhd8ed1ab_0 -- pickleshare=0.7.5=py_1003 -- pillow=9.4.0=py310h023d228_1 -- pip=24.2=pyh8b19718_1 -- pkgutil-resolve-name=1.3.10=pyhd8ed1ab_1 -- platformdirs=4.2.2=pyhd8ed1ab_0 -- prometheus_client=0.20.0=pyhd8ed1ab_0 -- prompt-toolkit=3.0.47=pyha770c72_0 -- prompt_toolkit=3.0.47=hd8ed1ab_0 -- psutil=6.0.0=py310hc51659f_0 -- pthread-stubs=0.4=h36c2ea0_1001 -- ptyprocess=0.7.0=pyhd3deb0d_0 -- pure_eval=0.2.3=pyhd8ed1ab_0 -- pycparser=2.22=pyhd8ed1ab_0 -- pygments=2.18.0=pyhd8ed1ab_0 -- pysocks=1.7.1=pyha2e5f31_6 -- python=3.10.14=hd12c33a_0_cpython -- python-dateutil=2.9.0=pyhd8ed1ab_0 -- python-fastjsonschema=2.20.0=pyhd8ed1ab_0 -- python-json-logger=2.0.7=pyhd8ed1ab_0 -- python_abi=3.10=5_cp310 -- pytorch=2.4.1=py3.10_cuda11.8_cudnn9.1.0_0 -- pytorch-cuda=11.8=h7e8668a_5 -- pytorch-mutex=1.0=cuda -- pytz=2024.1=pyhd8ed1ab_0 -- pyyaml=6.0.2=py310h5b4e0ec_0 -- pyzmq=26.2.0=py310h71f11fc_0 -- readline=8.2=h8228510_1 -- referencing=0.35.1=pyhd8ed1ab_0 -- requests=2.32.3=pyhd8ed1ab_0 -- rfc3339-validator=0.1.4=pyhd8ed1ab_0 -- rfc3986-validator=0.1.1=pyh9f0ad1d_0 -- rpds-py=0.20.0=py310h505e2c1_1 -- send2trash=1.8.3=pyh0d859eb_0 -- setuptools=73.0.1=pyhd8ed1ab_0 -- six=1.16.0=pyh6c4a22f_0 -- sniffio=1.3.1=pyhd8ed1ab_0 -- soupsieve=2.5=pyhd8ed1ab_1 -- stack_data=0.6.2=pyhd8ed1ab_0 -- sympy=1.13.2=pypyh2585a3b_103 -- tbb=2021.12.0=h84d6215_4 -- terminado=0.18.1=pyh0d859eb_0 -- tinycss2=1.3.0=pyhd8ed1ab_0 -- tk=8.6.13=noxft_h4845f30_101 -- tomli=2.0.1=pyhd8ed1ab_0 -- torchaudio=2.4.1=py310_cu118 -- torchtriton=3.0.0=py310 -- torchvision=0.19.1=py310_cu118 -- tornado=6.4.1=py310ha75aee5_1 -- traitlets=5.14.3=pyhd8ed1ab_0 -- types-python-dateutil=2.9.0.20240821=pyhd8ed1ab_0 -- typing-extensions=4.12.2=hd8ed1ab_0 -- typing_extensions=4.12.2=pyha770c72_0 -- typing_utils=0.1.0=pyhd8ed1ab_0 -- tzdata=2024a=h8827d51_1 -- uri-template=1.3.0=pyhd8ed1ab_0 -- urllib3=2.2.2=pyhd8ed1ab_1 -- wcwidth=0.2.13=pyhd8ed1ab_0 -- webcolors=24.8.0=pyhd8ed1ab_0 -- webencodings=0.5.1=pyhd8ed1ab_2 -- websocket-client=1.8.0=pyhd8ed1ab_0 -- wheel=0.44.0=pyhd8ed1ab_0 -- widgetsnbextension=4.0.13=pyhd8ed1ab_0 -- xorg-libxau=1.0.11=hd590300_0 -- xorg-libxdmcp=1.1.3=h7f98852_0 -- xz=5.2.6=h166bdaf_0 -- yaml=0.2.5=h7f98852_2 -- zeromq=4.3.5=h75354e8_4 -- zipp=3.20.1=pyhd8ed1ab_0 -- zlib=1.2.13=h4ab18f5_6 -- zstandard=0.23.0=py310ha39cb0e_1 -- zstd=1.5.6=ha6fb4c9_0 - diff --git a/notebooks_py/diffusion_/diffusion_policy/eval.py b/notebooks_py/diffusion_/diffusion_policy/eval.py deleted file mode 100644 index 06003f97d..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/eval.py +++ /dev/null @@ -1,64 +0,0 @@ -""" -Usage: -python eval.py --checkpoint data/image/pusht/diffusion_policy_cnn/train_0/checkpoints/latest.ckpt -o data/pusht_eval_output -""" - -import sys -# use line-buffering for both stdout and stderr -sys.stdout = open(sys.stdout.fileno(), mode='w', buffering=1) -sys.stderr = open(sys.stderr.fileno(), mode='w', buffering=1) - -import os -import pathlib -import click -import hydra -import torch -import dill -import wandb -import json -from diffusion_policy.workspace.base_workspace import BaseWorkspace - -@click.command() -@click.option('-c', '--checkpoint', required=True) -@click.option('-o', '--output_dir', required=True) -@click.option('-d', '--device', default='cuda:0') -def main(checkpoint, output_dir, device): - if os.path.exists(output_dir): - click.confirm(f"Output path {output_dir} already exists! Overwrite?", abort=True) - pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True) - - # load checkpoint - payload = torch.load(open(checkpoint, 'rb'), pickle_module=dill) - cfg = payload['cfg'] - cls = hydra.utils.get_class(cfg._target_) - workspace = cls(cfg, output_dir=output_dir) - workspace: BaseWorkspace - workspace.load_payload(payload, exclude_keys=None, include_keys=None) - - # get policy from workspace - policy = workspace.model - if cfg.training.use_ema: - policy = workspace.ema_model - - device = torch.device(device) - policy.to(device) - policy.eval() - - # run eval - env_runner = hydra.utils.instantiate( - cfg.task.env_runner, - output_dir=output_dir) - runner_log = env_runner.run(policy) - - # dump log to json - json_log = dict() - for key, value in runner_log.items(): - if isinstance(value, wandb.sdk.data_types.video.Video): - json_log[key] = value._path - else: - json_log[key] = value - out_path = os.path.join(output_dir, 'eval_log.json') - json.dump(json_log, open(out_path, 'w'), indent=2, sort_keys=True) - -if __name__ == '__main__': - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/eval_real_robot.py b/notebooks_py/diffusion_/diffusion_policy/eval_real_robot.py deleted file mode 100644 index cb3cdd17e..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/eval_real_robot.py +++ /dev/null @@ -1,418 +0,0 @@ -""" -Usage: -(robodiff)$ python eval_real_robot.py -i -o --robot_ip - -================ Human in control ============== -Robot movement: -Move your SpaceMouse to move the robot EEF (locked in xy plane). -Press SpaceMouse right button to unlock z axis. -Press SpaceMouse left button to enable rotation axes. - -Recording control: -Click the opencv window (make sure it's in focus). -Press "C" to start evaluation (hand control over to policy). -Press "Q" to exit program. - -================ Policy in control ============== -Make sure you can hit the robot hardware emergency-stop button quickly! - -Recording control: -Press "S" to stop evaluation and gain control back. -""" - -# %% -import time -from multiprocessing.managers import SharedMemoryManager -import click -import cv2 -import numpy as np -import torch -import dill -import hydra -import pathlib -import skvideo.io -from omegaconf import OmegaConf -import scipy.spatial.transform as st -from diffusion_policy.real_world.real_env import RealEnv -from diffusion_policy.real_world.spacemouse_shared_memory import Spacemouse -from diffusion_policy.common.precise_sleep import precise_wait -from diffusion_policy.real_world.real_inference_util import ( - get_real_obs_resolution, - get_real_obs_dict) -from diffusion_policy.common.pytorch_util import dict_apply -from diffusion_policy.workspace.base_workspace import BaseWorkspace -from diffusion_policy.policy.base_image_policy import BaseImagePolicy -from diffusion_policy.common.cv2_util import get_image_transform - - -OmegaConf.register_new_resolver("eval", eval, replace=True) - -@click.command() -@click.option('--input', '-i', required=True, help='Path to checkpoint') -@click.option('--output', '-o', required=True, help='Directory to save recording') -@click.option('--robot_ip', '-ri', required=True, help="UR5's IP address e.g. 192.168.0.204") -@click.option('--match_dataset', '-m', default=None, help='Dataset used to overlay and adjust initial condition') -@click.option('--match_episode', '-me', default=None, type=int, help='Match specific episode from the match dataset') -@click.option('--vis_camera_idx', default=0, type=int, help="Which RealSense camera to visualize.") -@click.option('--init_joints', '-j', is_flag=True, default=False, help="Whether to initialize robot joint configuration in the beginning.") -@click.option('--steps_per_inference', '-si', default=6, type=int, help="Action horizon for inference.") -@click.option('--max_duration', '-md', default=60, help='Max duration for each epoch in seconds.') -@click.option('--frequency', '-f', default=10, type=float, help="Control frequency in Hz.") -@click.option('--command_latency', '-cl', default=0.01, type=float, help="Latency between receiving SapceMouse command to executing on Robot in Sec.") -def main(input, output, robot_ip, match_dataset, match_episode, - vis_camera_idx, init_joints, - steps_per_inference, max_duration, - frequency, command_latency): - # load match_dataset - match_camera_idx = 0 - episode_first_frame_map = dict() - if match_dataset is not None: - match_dir = pathlib.Path(match_dataset) - match_video_dir = match_dir.joinpath('videos') - for vid_dir in match_video_dir.glob("*/"): - episode_idx = int(vid_dir.stem) - match_video_path = vid_dir.joinpath(f'{match_camera_idx}.mp4') - if match_video_path.exists(): - frames = skvideo.io.vread( - str(match_video_path), num_frames=1) - episode_first_frame_map[episode_idx] = frames[0] - print(f"Loaded initial frame for {len(episode_first_frame_map)} episodes") - - # load checkpoint - ckpt_path = input - payload = torch.load(open(ckpt_path, 'rb'), pickle_module=dill) - cfg = payload['cfg'] - cls = hydra.utils.get_class(cfg._target_) - workspace = cls(cfg) - workspace: BaseWorkspace - workspace.load_payload(payload, exclude_keys=None, include_keys=None) - - # hacks for method-specific setup. - action_offset = 0 - delta_action = False - if 'diffusion' in cfg.name: - # diffusion model - policy: BaseImagePolicy - policy = workspace.model - if cfg.training.use_ema: - policy = workspace.ema_model - - device = torch.device('cuda') - policy.eval().to(device) - - # set inference params - policy.num_inference_steps = 16 # DDIM inference iterations - policy.n_action_steps = policy.horizon - policy.n_obs_steps + 1 - - elif 'robomimic' in cfg.name: - # BCRNN model - policy: BaseImagePolicy - policy = workspace.model - - device = torch.device('cuda') - policy.eval().to(device) - - # BCRNN always has action horizon of 1 - steps_per_inference = 1 - action_offset = cfg.n_latency_steps - delta_action = cfg.task.dataset.get('delta_action', False) - - elif 'ibc' in cfg.name: - policy: BaseImagePolicy - policy = workspace.model - policy.pred_n_iter = 5 - policy.pred_n_samples = 4096 - - device = torch.device('cuda') - policy.eval().to(device) - steps_per_inference = 1 - action_offset = 1 - delta_action = cfg.task.dataset.get('delta_action', False) - else: - raise RuntimeError("Unsupported policy type: ", cfg.name) - - # setup experiment - dt = 1/frequency - - obs_res = get_real_obs_resolution(cfg.task.shape_meta) - n_obs_steps = cfg.n_obs_steps - print("n_obs_steps: ", n_obs_steps) - print("steps_per_inference:", steps_per_inference) - print("action_offset:", action_offset) - - with SharedMemoryManager() as shm_manager: - with Spacemouse(shm_manager=shm_manager) as sm, RealEnv( - output_dir=output, - robot_ip=robot_ip, - frequency=frequency, - n_obs_steps=n_obs_steps, - obs_image_resolution=obs_res, - obs_float32=True, - init_joints=init_joints, - enable_multi_cam_vis=True, - record_raw_video=True, - # number of threads per camera view for video recording (H.264) - thread_per_video=3, - # video recording quality, lower is better (but slower). - video_crf=21, - shm_manager=shm_manager) as env: - cv2.setNumThreads(1) - - # Should be the same as demo - # realsense exposure - env.realsense.set_exposure(exposure=120, gain=0) - # realsense white balance - env.realsense.set_white_balance(white_balance=5900) - - print("Waiting for realsense") - time.sleep(1.0) - - print("Warming up policy inference") - obs = env.get_obs() - with torch.no_grad(): - policy.reset() - obs_dict_np = get_real_obs_dict( - env_obs=obs, shape_meta=cfg.task.shape_meta) - obs_dict = dict_apply(obs_dict_np, - lambda x: torch.from_numpy(x).unsqueeze(0).to(device)) - result = policy.predict_action(obs_dict) - action = result['action'][0].detach().to('cpu').numpy() - assert action.shape[-1] == 2 - del result - - print('Ready!') - while True: - # ========= human control loop ========== - print("Human in control!") - state = env.get_robot_state() - target_pose = state['TargetTCPPose'] - t_start = time.monotonic() - iter_idx = 0 - while True: - # calculate timing - t_cycle_end = t_start + (iter_idx + 1) * dt - t_sample = t_cycle_end - command_latency - t_command_target = t_cycle_end + dt - - # pump obs - obs = env.get_obs() - - # visualize - episode_id = env.replay_buffer.n_episodes - vis_img = obs[f'camera_{vis_camera_idx}'][-1] - match_episode_id = episode_id - if match_episode is not None: - match_episode_id = match_episode - if match_episode_id in episode_first_frame_map: - match_img = episode_first_frame_map[match_episode_id] - ih, iw, _ = match_img.shape - oh, ow, _ = vis_img.shape - tf = get_image_transform( - input_res=(iw, ih), - output_res=(ow, oh), - bgr_to_rgb=False) - match_img = tf(match_img).astype(np.float32) / 255 - vis_img = np.minimum(vis_img, match_img) - - text = f'Episode: {episode_id}' - cv2.putText( - vis_img, - text, - (10,20), - fontFace=cv2.FONT_HERSHEY_SIMPLEX, - fontScale=0.5, - thickness=1, - color=(255,255,255) - ) - cv2.imshow('default', vis_img[...,::-1]) - key_stroke = cv2.pollKey() - if key_stroke == ord('q'): - # Exit program - env.end_episode() - exit(0) - elif key_stroke == ord('c'): - # Exit human control loop - # hand control over to the policy - break - - precise_wait(t_sample) - # get teleop command - sm_state = sm.get_motion_state_transformed() - # print(sm_state) - dpos = sm_state[:3] * (env.max_pos_speed / frequency) - drot_xyz = sm_state[3:] * (env.max_rot_speed / frequency) - - if not sm.is_button_pressed(0): - # translation mode - drot_xyz[:] = 0 - else: - dpos[:] = 0 - if not sm.is_button_pressed(1): - # 2D translation mode - dpos[2] = 0 - - drot = st.Rotation.from_euler('xyz', drot_xyz) - target_pose[:3] += dpos - target_pose[3:] = (drot * st.Rotation.from_rotvec( - target_pose[3:])).as_rotvec() - # clip target pose - target_pose[:2] = np.clip(target_pose[:2], [0.25, -0.45], [0.77, 0.40]) - - # execute teleop command - env.exec_actions( - actions=[target_pose], - timestamps=[t_command_target-time.monotonic()+time.time()]) - precise_wait(t_cycle_end) - iter_idx += 1 - - # ========== policy control loop ============== - try: - # start episode - policy.reset() - start_delay = 1.0 - eval_t_start = time.time() + start_delay - t_start = time.monotonic() + start_delay - env.start_episode(eval_t_start) - # wait for 1/30 sec to get the closest frame actually - # reduces overall latency - frame_latency = 1/30 - precise_wait(eval_t_start - frame_latency, time_func=time.time) - print("Started!") - iter_idx = 0 - term_area_start_timestamp = float('inf') - perv_target_pose = None - while True: - # calculate timing - t_cycle_end = t_start + (iter_idx + steps_per_inference) * dt - - # get obs - print('get_obs') - obs = env.get_obs() - obs_timestamps = obs['timestamp'] - print(f'Obs latency {time.time() - obs_timestamps[-1]}') - - # run inference - with torch.no_grad(): - s = time.time() - obs_dict_np = get_real_obs_dict( - env_obs=obs, shape_meta=cfg.task.shape_meta) - obs_dict = dict_apply(obs_dict_np, - lambda x: torch.from_numpy(x).unsqueeze(0).to(device)) - result = policy.predict_action(obs_dict) - # this action starts from the first obs step - action = result['action'][0].detach().to('cpu').numpy() - print('Inference latency:', time.time() - s) - - # convert policy action to env actions - if delta_action: - assert len(action) == 1 - if perv_target_pose is None: - perv_target_pose = obs['robot_eef_pose'][-1] - this_target_pose = perv_target_pose.copy() - this_target_pose[[0,1]] += action[-1] - perv_target_pose = this_target_pose - this_target_poses = np.expand_dims(this_target_pose, axis=0) - else: - this_target_poses = np.zeros((len(action), len(target_pose)), dtype=np.float64) - this_target_poses[:] = target_pose - this_target_poses[:,[0,1]] = action - - # deal with timing - # the same step actions are always the target for - action_timestamps = (np.arange(len(action), dtype=np.float64) + action_offset - ) * dt + obs_timestamps[-1] - action_exec_latency = 0.01 - curr_time = time.time() - is_new = action_timestamps > (curr_time + action_exec_latency) - if np.sum(is_new) == 0: - # exceeded time budget, still do something - this_target_poses = this_target_poses[[-1]] - # schedule on next available step - next_step_idx = int(np.ceil((curr_time - eval_t_start) / dt)) - action_timestamp = eval_t_start + (next_step_idx) * dt - print('Over budget', action_timestamp - curr_time) - action_timestamps = np.array([action_timestamp]) - else: - this_target_poses = this_target_poses[is_new] - action_timestamps = action_timestamps[is_new] - - # clip actions - this_target_poses[:,:2] = np.clip( - this_target_poses[:,:2], [0.25, -0.45], [0.77, 0.40]) - - # execute actions - env.exec_actions( - actions=this_target_poses, - timestamps=action_timestamps - ) - print(f"Submitted {len(this_target_poses)} steps of actions.") - - # visualize - episode_id = env.replay_buffer.n_episodes - vis_img = obs[f'camera_{vis_camera_idx}'][-1] - text = 'Episode: {}, Time: {:.1f}'.format( - episode_id, time.monotonic() - t_start - ) - cv2.putText( - vis_img, - text, - (10,20), - fontFace=cv2.FONT_HERSHEY_SIMPLEX, - fontScale=0.5, - thickness=1, - color=(255,255,255) - ) - cv2.imshow('default', vis_img[...,::-1]) - - - key_stroke = cv2.pollKey() - if key_stroke == ord('s'): - # Stop episode - # Hand control back to human - env.end_episode() - print('Stopped.') - break - - # auto termination - terminate = False - if time.monotonic() - t_start > max_duration: - terminate = True - print('Terminated by the timeout!') - - term_pose = np.array([ 3.40948500e-01, 2.17721816e-01, 4.59076878e-02, 2.22014183e+00, -2.22184883e+00, -4.07186655e-04]) - curr_pose = obs['robot_eef_pose'][-1] - dist = np.linalg.norm((curr_pose - term_pose)[:2], axis=-1) - if dist < 0.03: - # in termination area - curr_timestamp = obs['timestamp'][-1] - if term_area_start_timestamp > curr_timestamp: - term_area_start_timestamp = curr_timestamp - else: - term_area_time = curr_timestamp - term_area_start_timestamp - if term_area_time > 0.5: - terminate = True - print('Terminated by the policy!') - else: - # out of the area - term_area_start_timestamp = float('inf') - - if terminate: - env.end_episode() - break - - # wait for execution - precise_wait(t_cycle_end - frame_latency) - iter_idx += steps_per_inference - - except KeyboardInterrupt: - print("Interrupted!") - # stop robot. - env.end_episode() - - print("Stopped.") - - - -# %% -if __name__ == '__main__': - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/help.txt b/notebooks_py/diffusion_/diffusion_policy/help.txt deleted file mode 100644 index e69de29bb..000000000 diff --git a/notebooks_py/diffusion_/diffusion_policy/media/multimodal_sim.png b/notebooks_py/diffusion_/diffusion_policy/media/multimodal_sim.png deleted file mode 100644 index 1df8ed1a9..000000000 Binary files a/notebooks_py/diffusion_/diffusion_policy/media/multimodal_sim.png and /dev/null differ diff --git a/notebooks_py/diffusion_/diffusion_policy/media/teaser.png b/notebooks_py/diffusion_/diffusion_policy/media/teaser.png deleted file mode 100644 index 2e14aa04b..000000000 Binary files a/notebooks_py/diffusion_/diffusion_policy/media/teaser.png and /dev/null differ diff --git a/notebooks_py/diffusion_/diffusion_policy/multirun_metrics.py b/notebooks_py/diffusion_/diffusion_policy/multirun_metrics.py deleted file mode 100644 index 1254a6558..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/multirun_metrics.py +++ /dev/null @@ -1,267 +0,0 @@ -from typing import List, Optional -import pathlib -import pandas as pd -import numpy as np -import numba -import click -import time -import collections -import json -import wandb -import yaml -import numbers -import scipy.ndimage as sn -from diffusion_policy.common.json_logger import read_json_log, JsonLogger -import logging - -@numba.jit(nopython=True) -def get_indexed_window_average( - arr: np.ndarray, idxs: np.ndarray, window_size: int): - result = np.zeros(idxs.shape, dtype=arr.dtype) - length = arr.shape[0] - for i in range(len(idxs)): - idx = idxs[i] - start = max(idx - window_size, 0) - end = min(start + window_size, length) - result[i] = np.mean(arr[start:end]) - return result - - -def compute_metrics(log_df: pd.DataFrame, key: str, - end_step: Optional[int]=None, - k_min_loss: int=10, - k_around_max: int=10, - max_k_window: int=10, - replace_slash: int=True, - ): - if key not in log_df: - return dict() - - # prepare data - if end_step is not None: - log_df = log_df.iloc[:end_step] - is_key = ~pd.isnull(log_df[key]) - is_key_idxs = is_key.index[is_key].to_numpy() - if len(is_key_idxs) == 0: - return dict() - - key_data = log_df[key][is_key].to_numpy() - # after adding validation to workspace - # rollout happens at the last step of each epoch - # where the reported train_loss and val_loss - # are already the average for that epoch - train_loss = log_df['train_loss'][is_key].to_numpy() - val_loss = log_df['val_loss'][is_key].to_numpy() - - result = dict() - - log_key = key - if replace_slash: - log_key = key.replace('/', '_') - # max - max_value = np.max(key_data) - result['max/'+log_key] = max_value - - # k_around_max - max_idx = np.argmax(key_data) - end = min(max_idx + k_around_max // 2, len(key_data)) - start = max(end - k_around_max, 0) - k_around_max_value = np.mean(key_data[start:end]) - result['k_around_max/'+log_key] = k_around_max_value - - # max_k_window - k_window_value = sn.uniform_filter1d(key_data, size=max_k_window, axis=0, mode='nearest') - max_k_window_value = np.max(k_window_value) - result['max_k_window/'+log_key] = max_k_window_value - - # min_train_loss - min_idx = np.argmin(train_loss) - min_train_loss_value = key_data[min_idx] - result['min_train_loss/'+log_key] = min_train_loss_value - - # min_val_loss - min_idx = np.argmin(val_loss) - min_val_loss_value = key_data[min_idx] - result['min_val_loss/'+log_key] = min_val_loss_value - - # k_min_train_loss - min_loss_idxs = np.argsort(train_loss)[:k_min_loss] - k_min_train_loss_value = np.mean(key_data[min_loss_idxs]) - result['k_min_train_loss/'+log_key] = k_min_train_loss_value - - # k_min_val_loss - min_loss_idxs = np.argsort(val_loss)[:k_min_loss] - k_min_val_loss_value = np.mean(key_data[min_loss_idxs]) - result['k_min_val_loss/'+log_key] = k_min_val_loss_value - - # last - result['last/'+log_key] = key_data[-1] - - # global step for visualization - result['metric_global_step/'+log_key] = is_key_idxs[-1] - return result - - -def compute_metrics_agg( - log_dfs: List[pd.DataFrame], - key: str, end_step:int, - **kwargs): - - # compute metrics - results = collections.defaultdict(list) - for log_df in log_dfs: - result = compute_metrics(log_df, key=key, end_step=end_step, **kwargs) - for k, v in result.items(): - results[k].append(v) - # agg - agg_result = dict() - for k, v in results.items(): - value = np.mean(v) - if k.startswith('metric_global_step'): - value = int(value) - agg_result[k] = value - return agg_result - - -@click.command() -@click.option('--input', '-i', required=True, help='Root logging dir, contains train_* dirs') -@click.option('--key', '-k', multiple=True, default=['test/mean_score']) -@click.option('--interval', default=10, type=float) -@click.option('--replace_slash', default=True, type=bool) -@click.option('--index_key', '-ik', default='global_step') -@click.option('--use_wandb', '-w', is_flag=True, default=False) -@click.option('--project', default=None) -@click.option('--name', default=None) -@click.option('--id', default=None) -@click.option('--group', default=None) -def main( - input, - key, - interval, - replace_slash, - index_key, - use_wandb, - # wandb args - project, - name, - id, - group): - root_dir = pathlib.Path(input) - assert root_dir.is_dir() - metrics_dir = root_dir.joinpath('metrics') - metrics_dir.mkdir(exist_ok=True) - - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(levelname)s] %(message)s", - handlers=[ - logging.FileHandler(str(metrics_dir.joinpath("metrics.log"))), - logging.StreamHandler() - ] - ) - - train_dirs = list(root_dir.glob('train_*')) - log_files = [x.joinpath('logs.json.txt') for x in train_dirs] - logging.info("Monitor waiting for log files!") - while True: - # wait for files to show up - files_exist = True - for log_file in log_files: - if not log_file.is_file(): - files_exist = False - if files_exist: - break - time.sleep(1.0) - logging.info("All log files ready!") - - # init path - metric_log_path = metrics_dir.joinpath('logs.json.txt') - metric_path = metrics_dir.joinpath('metrics.json') - config_path = root_dir.joinpath('config.yaml') - - # load config - config = yaml.safe_load(config_path.open('r')) - - # init wandb - wandb_run = None - if use_wandb: - wandb_kwargs = config['logging'] - if project is not None: - wandb_kwargs['project'] = project - if id is not None: - wandb_kwargs['id'] = id - if name is not None: - wandb_kwargs['name'] = name - if group is not None: - wandb_kwargs['group'] = group - wandb_kwargs['resume'] = True - wandb_run = wandb.init( - dir=str(metrics_dir), - config=config, - # auto-resume run, automatically load id - # as long as using the same dir. - # https://docs.wandb.ai/guides/track/advanced/resuming#resuming-guidance - **wandb_kwargs - ) - wandb.config.update( - { - "output_dir": str(root_dir), - } - ) - - with JsonLogger(metric_log_path) as json_logger: - last_log = json_logger.get_last_log() - while True: - # read json files - log_dfs = [read_json_log(str(x), required_keys=key) for x in log_files] - - # previously logged data point - last_log_idx = -1 - if last_log is not None: - last_log_idx = log_dfs[0].index[log_dfs[0][index_key] <= last_log[index_key]][-1] - - start_idx = last_log_idx + 1 - # last idx where we have a data point from all logs - end_idx = min(*[len(x) for x in log_dfs]) - - # log every position - for this_idx in range(start_idx, end_idx): - # compute metrics - all_metrics = dict() - global_step = log_dfs[0]['global_step'][this_idx] - epoch = log_dfs[0]['epoch'][this_idx] - all_metrics['global_step'] = global_step - all_metrics['epoch'] = epoch - for k in key: - metrics = compute_metrics_agg( - log_dfs=log_dfs, key=k, end_step=this_idx+1, - replace_slash=replace_slash) - all_metrics.update(metrics) - - # sanitize metrics - old_metrics = all_metrics - all_metrics = dict() - for k, v in old_metrics.items(): - if isinstance(v, numbers.Integral): - all_metrics[k] = int(v) - elif isinstance(v, numbers.Number): - all_metrics[k] = float(v) - - has_update = all_metrics != last_log - if has_update: - last_log = all_metrics - json_logger.log(all_metrics) - - with metric_path.open('w') as f: - json.dump(all_metrics, f, sort_keys=True, indent=2) - - if wandb_run is not None: - wandb_run.log(all_metrics, step=all_metrics[index_key]) - - logging.info(f"Metrics logged at step {all_metrics[index_key]}") - - time.sleep(interval) - - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/.zattrs b/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/.zattrs deleted file mode 100644 index 9e26dfeeb..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/.zattrs +++ /dev/null @@ -1 +0,0 @@ -{} \ No newline at end of file diff --git a/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/.zgroup b/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/.zgroup deleted file mode 100644 index 3b7daf227..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/.zgroup +++ /dev/null @@ -1,3 +0,0 @@ -{ - "zarr_format": 2 -} \ No newline at end of file diff --git a/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/meta/.zattrs b/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/meta/.zattrs deleted file mode 100644 index 9e26dfeeb..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/meta/.zattrs +++ /dev/null @@ -1 +0,0 @@ -{} \ No newline at end of file diff --git a/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/meta/.zgroup b/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/meta/.zgroup deleted file mode 100644 index 3b7daf227..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/meta/.zgroup +++ /dev/null @@ -1,3 +0,0 @@ -{ - "zarr_format": 2 -} \ No newline at end of file diff --git a/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/meta/episode_ends/.zarray b/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/meta/episode_ends/.zarray deleted file mode 100644 index 35219f38b..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/pusht/pusht_cchi_v7_replay.zarr/meta/episode_ends/.zarray +++ /dev/null @@ -1,20 +0,0 @@ -{ - "chunks": [ - 208 - ], - "compressor": { - "blocksize": 0, - "clevel": 5, - "cname": "zstd", - "id": "blosc", - "shuffle": 2 - }, - "dtype": " (rec_start_time + 20.0): - break - - -if __name__ == "__main__": - test() diff --git a/notebooks_py/diffusion_/diffusion_policy/tests/test_pose_trajectory_interpolator.py b/notebooks_py/diffusion_/diffusion_policy/tests/test_pose_trajectory_interpolator.py deleted file mode 100644 index ba94abe7f..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/tests/test_pose_trajectory_interpolator.py +++ /dev/null @@ -1,126 +0,0 @@ -import sys -import os - -ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) -sys.path.append(ROOT_DIR) -os.chdir(ROOT_DIR) - -from tqdm import tqdm -import numpy as np -import scipy.interpolate as si -import scipy.spatial.transform as st -from diffusion_policy.common.pose_trajectory_interpolator import ( - rotation_distance, - pose_distance, - PoseTrajectoryInterpolator) - - -def test_rotation_distance(): - def rotation_distance_align(a: st.Rotation, b: st.Rotation) -> float: - return st.Rotation.align_vectors(b.as_matrix().T, a.as_matrix().T)[0].magnitude() - - for i in range(10000): - a = st.Rotation.from_euler('xyz', np.random.uniform(-7,7,size=3)) - b = st.Rotation.from_euler('xyz', np.random.uniform(-7,7,size=3)) - x = rotation_distance(a, b) - y = rotation_distance_align(a, b) - assert abs(x-y) < 1e-7 - -def test_pose_trajectory_interpolator(): - t = np.linspace(-1,5,100) - interp = PoseTrajectoryInterpolator( - [0,1,3], - np.zeros((3,6)) - ) - times = interp.times - poses = interp.poses - - trimmed_interp = interp.trim(-1,4) - assert len(trimmed_interp.times) == 5 - trimmed_interp(t) - - trimmed_interp = interp.trim(-1,4) - assert len(trimmed_interp.times) == 5 - trimmed_interp(t) - - trimmed_interp = interp.trim(0.5, 3.5) - assert len(trimmed_interp.times) == 4 - trimmed_interp(t) - - trimmed_interp = interp.trim(0.5, 2.5) - assert len(trimmed_interp.times) == 3 - trimmed_interp(t) - - trimmed_interp = interp.trim(0.5, 1.5) - assert len(trimmed_interp.times) == 3 - trimmed_interp(t) - - trimmed_interp = interp.trim(1.2, 1.5) - assert len(trimmed_interp.times) == 2 - trimmed_interp(t) - - trimmed_interp = interp.trim(1.3, 1.3) - assert len(trimmed_interp.times) == 1 - trimmed_interp(t) - - # import pdb; pdb.set_trace() - -def test_add_waypoint(): - # fuzz testing - for i in tqdm(range(10000)): - rng = np.random.default_rng(i) - n_waypoints = rng.integers(1, 5) - waypoint_times = np.sort(rng.uniform(0, 1, size=n_waypoints)) - last_waypoint_time = waypoint_times[-1] - insert_time = rng.uniform(-0.1, 1.1) - curr_time = rng.uniform(-0.1, 1.1) - max_pos_speed = rng.poisson(3) + 1e-3 - max_rot_speed = rng.poisson(3) + 1e-3 - waypoint_poses = rng.normal(0, 3, size=(n_waypoints, 6)) - new_pose = rng.normal(0, 3, size=6) - - if rng.random() < 0.1: - last_waypoint_time = None - if rng.random() < 0.1: - curr_time = None - - interp = PoseTrajectoryInterpolator( - times=waypoint_times, - poses=waypoint_poses) - new_interp = interp.add_waypoint( - pose=new_pose, - time=insert_time, - max_pos_speed=max_pos_speed, - max_rot_speed=max_rot_speed, - curr_time=curr_time, - last_waypoint_time=last_waypoint_time - ) - -def test_drive_to_waypoint(): - # fuzz testing - for i in tqdm(range(10000)): - rng = np.random.default_rng(i) - n_waypoints = rng.integers(1, 5) - waypoint_times = np.sort(rng.uniform(0, 1, size=n_waypoints)) - insert_time = rng.uniform(-0.1, 1.1) - curr_time = rng.uniform(-0.1, 1.1) - max_pos_speed = rng.poisson(3) + 1e-3 - max_rot_speed = rng.poisson(3) + 1e-3 - waypoint_poses = rng.normal(0, 3, size=(n_waypoints, 6)) - new_pose = rng.normal(0, 3, size=6) - - interp = PoseTrajectoryInterpolator( - times=waypoint_times, - poses=waypoint_poses) - new_interp = interp.drive_to_waypoint( - pose=new_pose, - time=insert_time, - curr_time=curr_time, - max_pos_speed=max_pos_speed, - max_rot_speed=max_rot_speed - ) - - - -if __name__ == '__main__': - test_drive_to_waypoint() diff --git a/notebooks_py/diffusion_/diffusion_policy/tests/test_precise_sleep.py b/notebooks_py/diffusion_/diffusion_policy/tests/test_precise_sleep.py deleted file mode 100644 index 85c0d94cb..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/tests/test_precise_sleep.py +++ /dev/null @@ -1,56 +0,0 @@ -import sys -import os - -ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) -sys.path.append(ROOT_DIR) -os.chdir(ROOT_DIR) - -import time -import numpy as np -from diffusion_policy.common.precise_sleep import precise_sleep, precise_wait - - -def test_sleep(): - dt = 0.1 - tol = 1e-3 - time_samples = list() - for i in range(100): - precise_sleep(dt) - # time.sleep(dt) - time_samples.append(time.monotonic()) - time_deltas = np.diff(time_samples) - - from matplotlib import pyplot as plt - plt.plot(time_deltas) - plt.ylim((dt-tol,dt+tol)) - - -def test_wait(): - dt = 0.1 - tol = 1e-3 - errors = list() - t_start = time.monotonic() - for i in range(1,100): - t_end_desired = t_start + i * dt - time.sleep(t_end_desired - time.monotonic()) - t_end = time.monotonic() - errors.append(t_end - t_end_desired) - - new_errors = list() - t_start = time.monotonic() - for i in range(1,100): - t_end_desired = t_start + i * dt - precise_wait(t_end_desired) - t_end = time.monotonic() - new_errors.append(t_end - t_end_desired) - - from matplotlib import pyplot as plt - plt.plot(errors, label='time.sleep') - plt.plot(new_errors, label='sleep/spin hybrid') - plt.ylim((-tol,+tol)) - plt.title('0.1 sec sleep error') - plt.legend() - - -if __name__ == '__main__': - test_sleep() diff --git a/notebooks_py/diffusion_/diffusion_policy/tests/test_replay_buffer.py b/notebooks_py/diffusion_/diffusion_policy/tests/test_replay_buffer.py deleted file mode 100644 index c048366b6..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/tests/test_replay_buffer.py +++ /dev/null @@ -1,62 +0,0 @@ -import sys -import os - -ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) -sys.path.append(ROOT_DIR) -os.chdir(ROOT_DIR) - -import zarr -from diffusion_policy.common.replay_buffer import ReplayBuffer - -def test(): - import numpy as np - buff = ReplayBuffer.create_empty_numpy() - buff.add_episode({ - 'obs': np.zeros((100,10), dtype=np.float16) - }) - buff.add_episode({ - 'obs': np.ones((50,10)), - 'action': np.ones((50,2)) - }) - # buff.rechunk(256) - obs = buff.get_episode(0) - - import numpy as np - buff = ReplayBuffer.create_empty_zarr() - buff.add_episode({ - 'obs': np.zeros((100,10), dtype=np.float16) - }) - buff.add_episode({ - 'obs': np.ones((50,10)), - 'action': np.ones((50,2)) - }) - obs = buff.get_episode(0) - buff.set_chunks({ - 'obs': (100,10), - 'action': (100,2) - }) - - -def test_real(): - import os - dist_group = zarr.open( - os.path.expanduser('~/dev/diffusion_policy/data/pusht/pusht_cchi_v2.zarr'), 'r') - - buff = ReplayBuffer.create_empty_numpy() - key, group = next(iter(dist_group.items())) - for key, group in dist_group.items(): - buff.add_episode(group) - - # out_path = os.path.expanduser('~/dev/diffusion_policy/data/pusht_cchi2_v2_replay.zarr') - out_path = os.path.expanduser('~/dev/diffusion_policy/data/test.zarr') - out_store = zarr.DirectoryStore(out_path) - buff.save_to_store(out_store) - - buff = ReplayBuffer.copy_from_path(out_path, store=zarr.MemoryStore()) - buff.pop_episode() - - -def test_pop(): - buff = ReplayBuffer.create_from_path( - '/home/chengchi/dev/diffusion_policy/data/pusht_cchi_v3_replay.zarr', - mode='rw') diff --git a/notebooks_py/diffusion_/diffusion_policy/tests/test_ring_buffer.py b/notebooks_py/diffusion_/diffusion_policy/tests/test_ring_buffer.py deleted file mode 100644 index 18d03bc81..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/tests/test_ring_buffer.py +++ /dev/null @@ -1,188 +0,0 @@ -import sys -import os - -ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) -sys.path.append(ROOT_DIR) -os.chdir(ROOT_DIR) - -import time -import numpy as np -import multiprocessing as mp -from multiprocessing.managers import SharedMemoryManager -from diffusion_policy.shared_memory.shared_memory_ring_buffer import ( - SharedMemoryRingBuffer, - SharedAtomicCounter) - - -def test(): - shm_manager = SharedMemoryManager() - shm_manager.start() - ring_buffer = SharedMemoryRingBuffer.create_from_examples( - shm_manager, - {'timestamp': np.array(0, dtype=np.float64)}, - buffer_size=128 - ) - for i in range(30): - ring_buffer.put({ - 'timestamp': np.array( - time.perf_counter(), - dtype=np.float64) - }) - print(ring_buffer.get()) - - -def _timestamp_worker(ring_buffer, start_event, stop_event): - while not stop_event.is_set(): - start_event.set() - ring_buffer.put({ - 'timestamp': np.array( - time.time(), - dtype=np.float64) - }) - - -def test_mp(): - shm_manager = SharedMemoryManager() - shm_manager.start() - ring_buffer = SharedMemoryRingBuffer.create_from_examples( - shm_manager, - {'timestamp': np.array(0, dtype=np.float64)}, - get_max_k=1, - get_time_budget=0.01, - put_desired_frequency=1000 - ) - start_event = mp.Event() - stop_event = mp.Event() - worker = mp.Process(target=_timestamp_worker, args=( - ring_buffer, start_event, stop_event)) - worker.start() - start_event.wait() - for i in range(1000): - t = float(ring_buffer.get()['timestamp']) - curr_t = time.time() - print('latency', curr_t - t) - stop_event.set() - worker.join() - - -def test_get_last_k(): - shm_manager = SharedMemoryManager() - shm_manager.start() - ring_buffer = SharedMemoryRingBuffer.create_from_examples( - shm_manager, - {'counter': np.array(0, dtype=np.int64)}, - buffer_size=8 - ) - - from collections import deque - k = 4 - last_k = deque(maxlen=k) - for i in range(100): - ring_buffer.put({ - 'counter': np.array(i, dtype=np.int64) - }) - last_k.append(i) - if i > k: - result = ring_buffer.get_last_k(k)['counter'] - assert np.allclose(result, last_k) - - print(ring_buffer.shared_arrays['counter'].get()) - result = ring_buffer.get_last_k(4) - print(result) - - -def test_timing(): - shm_manager = SharedMemoryManager() - shm_manager.start() - ring_buffer = SharedMemoryRingBuffer.create_from_examples( - shm_manager, - {'counter': np.array(0, dtype=np.int64)}, - get_max_k=8, - get_time_budget=0.1, - put_desired_frequency=100 - ) - # print(ring_buffer.timestamp_array.get()) - print('buffer_size', ring_buffer.buffer_size) - - dt = 1 / 150 - t_init = time.monotonic() - for i in range(1000): - t_start = time.monotonic() - ring_buffer.put({ - 'counter': np.array(i, dtype=np.int64) - }, wait=False) - if (i % 10 == 0) and (i > 0): - result = ring_buffer.get_last_k(8) - - t_end =time.monotonic() - desired_t = (i+1) * dt + t_init - if desired_t > t_end: - time.sleep(desired_t - t_end) - hz = 1 / (time.monotonic() - t_start) - print(f'{hz}Hz') - - -def _timestamp_image_worker(ring_buffer, img_shape, dt, start_event, stop_event): - i = 0 - t_init = time.monotonic() - image = np.ones(img_shape, dtype=np.uint8) - while not stop_event.is_set(): - t_start = time.monotonic() - start_event.set() - ring_buffer.put({ - 'img': image, - 'timestamp': time.time(), - 'counter': i - }) - t_end = time.monotonic() - desired_t = (i+1) * dt + t_init - # print('alive') - if desired_t > t_end: - time.sleep(desired_t - t_end) - # hz = 1 / (time.monotonic() - t_start) - i += 1 - - -def test_timing_mp(): - shm_manager = SharedMemoryManager() - shm_manager.start() - - hz = 200 - img_shape = (1920,1080,3) - ring_buffer = SharedMemoryRingBuffer.create_from_examples( - shm_manager, - examples={ - 'img': np.zeros(img_shape, dtype=np.uint8), - 'timestamp': time.time(), - 'counter': 0 - }, - get_max_k=60, - get_time_budget=0.02, - put_desired_frequency=hz - ) - start_event = mp.Event() - stop_event = mp.Event() - worker = mp.Process(target=_timestamp_image_worker, args=( - ring_buffer, img_shape, 1/hz, start_event, stop_event)) - worker.start() - start_event.wait() - out = None - t_start = time.monotonic() - k = 1 - for i in range(1000): - if ring_buffer.count < k: - time.sleep(0) - continue - out = ring_buffer.get_last_k(k=k, out=out) - t = float(out['timestamp'][-1]) - curr_t = time.time() - print('latency', curr_t - t) - t_end = time.monotonic() - print('Get Hz', 1/(t_end-t_start)*1000) - stop_event.set() - worker.join() - - -if __name__ == '__main__': - # test_mp() - test_timing_mp() diff --git a/notebooks_py/diffusion_/diffusion_policy/tests/test_robomimic_image_runner.py b/notebooks_py/diffusion_/diffusion_policy/tests/test_robomimic_image_runner.py deleted file mode 100644 index 629e6c776..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/tests/test_robomimic_image_runner.py +++ /dev/null @@ -1,38 +0,0 @@ -import sys -import os - -ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) -sys.path.append(ROOT_DIR) -os.chdir(ROOT_DIR) - -from diffusion_policy.env_runner.robomimic_image_runner import RobomimicImageRunner - -def test(): - import os - from omegaconf import OmegaConf - cfg_path = os.path.expanduser('~/dev/diffusion_policy/diffusion_policy/config/task/lift_image.yaml') - cfg = OmegaConf.load(cfg_path) - cfg['n_obs_steps'] = 1 - cfg['n_action_steps'] = 1 - cfg['past_action_visible'] = False - runner_cfg = cfg['env_runner'] - runner_cfg['n_train'] = 1 - runner_cfg['n_test'] = 1 - del runner_cfg['_target_'] - runner = RobomimicImageRunner( - **runner_cfg, - output_dir='/tmp/test') - - # import pdb; pdb.set_trace() - - self = runner - env = self.env - env.seed(seeds=self.env_seeds) - obs = env.reset() - for i in range(10): - _ = env.step(env.action_space.sample()) - - imgs = env.render() - -if __name__ == '__main__': - test() diff --git a/notebooks_py/diffusion_/diffusion_policy/tests/test_robomimic_lowdim_runner.py b/notebooks_py/diffusion_/diffusion_policy/tests/test_robomimic_lowdim_runner.py deleted file mode 100644 index 83390fda0..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/tests/test_robomimic_lowdim_runner.py +++ /dev/null @@ -1,34 +0,0 @@ -import sys -import os - -ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) -sys.path.append(ROOT_DIR) -os.chdir(ROOT_DIR) - -from diffusion_policy.env_runner.robomimic_lowdim_runner import RobomimicLowdimRunner - -def test(): - import os - from omegaconf import OmegaConf - cfg_path = os.path.expanduser('~/dev/diffusion_policy/diffusion_policy/config/task/lift_lowdim.yaml') - cfg = OmegaConf.load(cfg_path) - cfg['n_obs_steps'] = 1 - cfg['n_action_steps'] = 1 - cfg['past_action_visible'] = False - runner_cfg = cfg['env_runner'] - runner_cfg['n_train'] = 1 - runner_cfg['n_test'] = 0 - del runner_cfg['_target_'] - runner = RobomimicLowdimRunner( - **runner_cfg, - output_dir='/tmp/test') - - # import pdb; pdb.set_trace() - - self = runner - env = self.env - env.seed(seeds=self.env_seeds) - obs = env.reset() - -if __name__ == '__main__': - test() diff --git a/notebooks_py/diffusion_/diffusion_policy/tests/test_shared_queue.py b/notebooks_py/diffusion_/diffusion_policy/tests/test_shared_queue.py deleted file mode 100644 index 4e6f9d9cd..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/tests/test_shared_queue.py +++ /dev/null @@ -1,67 +0,0 @@ -import sys -import os - -ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) -sys.path.append(ROOT_DIR) -os.chdir(ROOT_DIR) - -import numpy as np -from multiprocessing.managers import SharedMemoryManager -from diffusion_policy.shared_memory.shared_memory_queue import SharedMemoryQueue, Full, Empty - - -def test(): - shm_manager = SharedMemoryManager() - shm_manager.start() - example = { - 'cmd': 0, - 'pose': np.zeros((6,)) - } - queue = SharedMemoryQueue.create_from_examples( - shm_manager=shm_manager, - examples=example, - buffer_size=3 - ) - raised = False - try: - queue.get() - except Empty: - raised = True - assert raised - - data = { - 'cmd': 1, - 'pose': np.ones((6,)) - } - queue.put(data) - result = queue.get() - assert result['cmd'] == data['cmd'] - assert np.allclose(result['pose'], data['pose']) - - queue.put(data) - queue.put(data) - queue.put(data) - assert queue.qsize() == 3 - raised = False - try: - queue.put(data) - except Full: - raised = True - assert raised - - result = queue.get_all() - assert np.allclose(result['cmd'], [1,1,1]) - - queue.put({'cmd': 0}) - queue.put({'cmd': 1}) - queue.put({'cmd': 2}) - queue.get() - queue.put({'cmd': 3}) - - result = queue.get_k(3) - assert np.allclose(result['cmd'], [1,2,3]) - - queue.clear() - -if __name__ == "__main__": - test() diff --git a/notebooks_py/diffusion_/diffusion_policy/tests/test_single_realsense.py b/notebooks_py/diffusion_/diffusion_policy/tests/test_single_realsense.py deleted file mode 100644 index 836ad6b2c..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/tests/test_single_realsense.py +++ /dev/null @@ -1,87 +0,0 @@ -import sys -import os - -ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) -sys.path.append(ROOT_DIR) -os.chdir(ROOT_DIR) - -import cv2 -import json -import time -from multiprocessing.managers import SharedMemoryManager -from diffusion_policy.real_world.single_realsense import SingleRealsense - -def test(): - - serials = SingleRealsense.get_connected_devices_serial() - # import pdb; pdb.set_trace() - serial = serials[0] - config = json.load(open('/home/cchi/dev/diffusion_policy/diffusion_policy/real_world/realsense_config/415_high_accuracy_mode.json', 'r')) - - def transform(data): - color = data['color'] - h,w,_ = color.shape - factor = 2 - color = cv2.resize(color, (w//factor,h//factor), interpolation=cv2.INTER_AREA) - # color = color[:,140:500] - data['color'] = color - return data - - # at 960x540 with //3, 60fps and 30fps are indistinguishable - - with SharedMemoryManager() as shm_manager: - with SingleRealsense( - shm_manager=shm_manager, - serial_number=serial, - resolution=(1280,720), - # resolution=(960,540), - # resolution=(640,480), - capture_fps=30, - enable_color=True, - # enable_depth=True, - # enable_infrared=True, - # advanced_mode_config=config, - # transform=transform, - # recording_transform=transform - # verbose=True - ) as realsense: - cv2.setNumThreads(1) - realsense.set_exposure(exposure=150, gain=5) - intr = realsense.get_intrinsics() - print(intr) - - - video_path = 'data_local/test.mp4' - rec_start_time = time.time() + 2 - realsense.start_recording(video_path, start_time=rec_start_time) - - data = None - while True: - data = realsense.get(out=data) - t = time.time() - # print('capture_latency', data['receive_timestamp']-data['capture_timestamp'], 'receive_latency', t - data['receive_timestamp']) - # print('receive', t - data['receive_timestamp']) - - # dt = time.time() - data['timestamp'] - # print(dt) - # print(data['capture_timestamp'] - rec_start_time) - - bgr = data['color'] - # print(bgr.shape) - cv2.imshow('default', bgr) - key = cv2.pollKey() - # if key == ord('q'): - # break - # elif key == ord('r'): - # video_path = 'data_local/test.mp4' - # realsense.start_recording(video_path) - # elif key == ord('s'): - # realsense.stop_recording() - - time.sleep(1/60) - if time.time() > (rec_start_time + 20.0): - break - - -if __name__ == "__main__": - test() diff --git a/notebooks_py/diffusion_/diffusion_policy/tests/test_timestamp_accumulator.py b/notebooks_py/diffusion_/diffusion_policy/tests/test_timestamp_accumulator.py deleted file mode 100644 index 144826bd1..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/tests/test_timestamp_accumulator.py +++ /dev/null @@ -1,151 +0,0 @@ -import sys -import os - -ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) -sys.path.append(ROOT_DIR) -os.chdir(ROOT_DIR) - -import numpy as np -import time -from diffusion_policy.common.timestamp_accumulator import ( - get_accumulate_timestamp_idxs, - TimestampObsAccumulator, - TimestampActionAccumulator -) - - -def test_index(): - buffer = np.zeros(16) - start_time = 0.0 - dt = 1/10 - - timestamps = np.linspace(0,1,100) - gi = list() - next_global_idx = 0 - - local_idxs, global_idxs, next_global_idx = get_accumulate_timestamp_idxs(timestamps, - start_time=start_time, dt=dt, next_global_idx=next_global_idx) - assert local_idxs[0] == 0 - assert global_idxs[0] == 0 - # print(local_idxs) - # print(global_idxs) - # print(timestamps[local_idxs]) - buffer[global_idxs] = timestamps[local_idxs] - gi.extend(global_idxs) - - timestamps = np.linspace(0.5,1.5,100) - local_idxs, global_idxs, next_global_idx = get_accumulate_timestamp_idxs(timestamps, - start_time=start_time, dt=dt, next_global_idx = next_global_idx) - # print(local_idxs) - # print(global_idxs) - # print(timestamps[local_idxs]) - # import pdb; pdb.set_trace() - buffer[global_idxs] = timestamps[local_idxs] - gi.extend(global_idxs) - - assert np.all(buffer[1:] > buffer[:-1]) - assert np.all(np.array(gi) == np.array(list(range(len(gi))))) - # print(buffer) - - # start over - next_global_idx = 0 - timestamps = np.linspace(0,1,3) - local_idxs, global_idxs, next_global_idx = get_accumulate_timestamp_idxs(timestamps, - start_time=start_time, dt=dt, next_global_idx = next_global_idx) - assert local_idxs[0] == 0 - assert local_idxs[-1] == 2 - # print(local_idxs) - # print(global_idxs) - # print(timestamps[local_idxs]) - - # test numerical error issue - # this becomes a problem when eps <= 1e-7 - start_time = time.time() - next_global_idx = 0 - timestamps = np.arange(100000) * dt + start_time - local_idxs, global_idxs, next_global_idx = get_accumulate_timestamp_idxs(timestamps, - start_time=start_time, dt=dt, next_global_idx = next_global_idx) - assert local_idxs == global_idxs - # print(local_idxs) - # print(global_idxs) - # print(timestamps[local_idxs]) - - -def test_obs_accumulator(): - dt = 1/10 - ddt = 1/100 - n = 100 - d = 6 - start_time = time.time() - toa = TimestampObsAccumulator(start_time, dt) - poses = np.arange(n).reshape((n,1)) - poses = np.repeat(poses, d, axis=1) - timestamps = np.arange(n) * ddt + start_time - - toa.put({ - 'pose': poses, - 'timestamp': timestamps - }, timestamps) - assert np.all(toa.data['pose'][:,0] == np.arange(10)*10) - assert len(toa) == 10 - - # add the same thing, result shouldn't change - toa.put({ - 'pose': poses, - 'timestamp': timestamps - }, timestamps) - assert np.all(toa.data['pose'][:,0] == np.arange(10)*10) - assert len(toa) == 10 - - # add lower than desired freuquency to test fill_in - dt = 1/10 - ddt = 1/5 - n = 10 - d = 6 - start_time = time.time() - toa = TimestampObsAccumulator(start_time, dt) - poses = np.arange(n).reshape((n,1)) - poses = np.repeat(poses, d, axis=1) - timestamps = np.arange(n) * ddt + start_time - - toa.put({ - 'pose': poses, - 'timestamp': timestamps - }, timestamps) - assert len(toa) == 1 + (n-1) * 2 - - timestamps = (np.arange(n) + 2) * ddt + start_time - toa.put({ - 'pose': poses, - 'timestamp': timestamps - }, timestamps) - assert len(toa) == 1 + (n-1) * 2 + 4 - - -def test_action_accumulator(): - dt = 1/10 - n = 10 - d = 6 - start_time = time.time() - taa = TimestampActionAccumulator(start_time, dt) - actions = np.arange(n).reshape((n,1)) - actions = np.repeat(actions, d, axis=1) - - timestamps = np.arange(n) * dt + start_time - taa.put(actions, timestamps) - assert np.all(taa.actions == actions) - assert np.all(taa.timestamps == timestamps) - - # add another round - taa.put(actions-5, timestamps-0.5) - assert np.allclose(taa.timestamps, timestamps) - - # add another round - taa.put(actions+5, timestamps+0.5) - assert len(taa) == 15 - assert np.all(taa.actions[:,0] == np.arange(15)) - - - -if __name__ == '__main__': - test_action_accumulator() diff --git a/notebooks_py/diffusion_/diffusion_policy/train.py b/notebooks_py/diffusion_/diffusion_policy/train.py deleted file mode 100644 index 51c564e6f..000000000 --- a/notebooks_py/diffusion_/diffusion_policy/train.py +++ /dev/null @@ -1,35 +0,0 @@ -""" -Usage: -Training: -python train.py --config-name=train_diffusion_lowdim_workspace -""" - -import sys -# use line-buffering for both stdout and stderr -sys.stdout = open(sys.stdout.fileno(), mode='w', buffering=1) -sys.stderr = open(sys.stderr.fileno(), mode='w', buffering=1) - -import hydra -from omegaconf import OmegaConf -import pathlib -from diffusion_policy.workspace.base_workspace import BaseWorkspace - -# allows arbitrary python code execution in configs using the ${eval:''} resolver -OmegaConf.register_new_resolver("eval", eval, replace=True) - -@hydra.main( - version_base=None, - config_path=str(pathlib.Path(__file__).parent.joinpath( - 'diffusion_policy','config')) -) -def main(cfg: OmegaConf): - # resolve immediately so all the ${now:} resolvers - # will use the same time. - OmegaConf.resolve(cfg) - - cls = hydra.utils.get_class(cfg._target_) - workspace: BaseWorkspace = cls(cfg) - workspace.run() - -if __name__ == "__main__": - main() diff --git a/notebooks_py/diffusion_/image_pusht_diffusion_policy_cnn.yaml b/notebooks_py/diffusion_/image_pusht_diffusion_policy_cnn.yaml deleted file mode 100644 index dd884081c..000000000 --- a/notebooks_py/diffusion_/image_pusht_diffusion_policy_cnn.yaml +++ /dev/null @@ -1,182 +0,0 @@ -_target_: diffusion_policy.workspace.train_diffusion_unet_hybrid_workspace.TrainDiffusionUnetHybridWorkspace -checkpoint: - save_last_ckpt: true - save_last_snapshot: false - topk: - format_str: epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt - k: 5 - mode: max - monitor_key: test_mean_score -dataloader: - batch_size: 64 - num_workers: 8 - persistent_workers: false - pin_memory: true - shuffle: true -dataset_obs_steps: 2 -ema: - _target_: diffusion_policy.model.diffusion.ema_model.EMAModel - inv_gamma: 1.0 - max_value: 0.9999 - min_value: 0.0 - power: 0.75 - update_after_step: 0 -exp_name: default -horizon: 16 -keypoint_visible_rate: 1.0 -logging: - group: null - id: null - mode: online - name: 2023.01.16-20.20.06_train_diffusion_unet_hybrid_pusht_image - project: diffusion_policy_debug - resume: true - tags: - - train_diffusion_unet_hybrid - - pusht_image - - default -multi_run: - run_dir: data/outputs/2023.01.16/20.20.06_train_diffusion_unet_hybrid_pusht_image - wandb_name_base: 2023.01.16-20.20.06_train_diffusion_unet_hybrid_pusht_image -n_action_steps: 8 -n_latency_steps: 0 -n_obs_steps: 2 -name: train_diffusion_unet_hybrid -obs_as_global_cond: true -optimizer: - _target_: torch.optim.AdamW - betas: - - 0.95 - - 0.999 - eps: 1.0e-08 - lr: 0.0001 - weight_decay: 1.0e-06 -past_action_visible: false -policy: - _target_: diffusion_policy.policy.diffusion_unet_hybrid_image_policy.DiffusionUnetHybridImagePolicy - cond_predict_scale: true - crop_shape: - - 84 - - 84 - diffusion_step_embed_dim: 128 - down_dims: - - 512 - - 1024 - - 2048 - eval_fixed_crop: true - horizon: 16 - kernel_size: 5 - n_action_steps: 8 - n_groups: 8 - n_obs_steps: 2 - noise_scheduler: - _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler - beta_end: 0.02 - beta_schedule: squaredcos_cap_v2 - beta_start: 0.0001 - clip_sample: true - num_train_timesteps: 100 - prediction_type: epsilon - variance_type: fixed_small - num_inference_steps: 100 - obs_as_global_cond: true - obs_encoder_group_norm: true - shape_meta: - action: - shape: - - 2 - obs: - agent_pos: - shape: - - 2 - type: low_dim - image: - shape: - - 3 - - 96 - - 96 - type: rgb -shape_meta: - action: - shape: - - 2 - obs: - agent_pos: - shape: - - 2 - type: low_dim - image: - shape: - - 3 - - 96 - - 96 - type: rgb -task: - dataset: - _target_: diffusion_policy.dataset.pusht_image_dataset.PushTImageDataset - horizon: 16 - max_train_episodes: 90 - pad_after: 7 - pad_before: 1 - seed: 42 - val_ratio: 0.02 - zarr_path: data/pusht/pusht_cchi_v7_replay.zarr - env_runner: - _target_: diffusion_policy.env_runner.pusht_image_runner.PushTImageRunner - fps: 10 - legacy_test: true - max_steps: 300 - n_action_steps: 8 - n_envs: null - n_obs_steps: 2 - n_test: 50 - n_test_vis: 4 - n_train: 6 - n_train_vis: 2 - past_action: false - test_start_seed: 100000 - train_start_seed: 0 - image_shape: - - 3 - - 96 - - 96 - name: pusht_image - shape_meta: - action: - shape: - - 2 - obs: - agent_pos: - shape: - - 2 - type: low_dim - image: - shape: - - 3 - - 96 - - 96 - type: rgb -task_name: pusht_image -training: - checkpoint_every: 50 - debug: false - device: cuda:0 - gradient_accumulate_every: 1 - lr_scheduler: cosine - lr_warmup_steps: 500 - max_train_steps: null - max_val_steps: null - num_epochs: 3050 - resume: true - rollout_every: 50 - sample_every: 5 - seed: 42 - tqdm_interval_sec: 1.0 - use_ema: true - val_every: 1 -val_dataloader: - batch_size: 64 - num_workers: 8 - persistent_workers: false - pin_memory: true - shuffle: false diff --git a/notebooks_py/flow/public/company/1.jpg b/notebooks_py/flow/public/company/1.jpg deleted file mode 100644 index 8f9b7aca1..000000000 Binary files a/notebooks_py/flow/public/company/1.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/company/2.jpg b/notebooks_py/flow/public/company/2.jpg deleted file mode 100644 index bbc1cfb96..000000000 Binary files a/notebooks_py/flow/public/company/2.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/company/3.jpg b/notebooks_py/flow/public/company/3.jpg deleted file mode 100644 index 808d177d6..000000000 Binary files a/notebooks_py/flow/public/company/3.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/company/4.jpg b/notebooks_py/flow/public/company/4.jpg deleted file mode 100644 index dec347649..000000000 Binary files a/notebooks_py/flow/public/company/4.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/company/5.jpg b/notebooks_py/flow/public/company/5.jpg deleted file mode 100644 index ce279077c..000000000 Binary files a/notebooks_py/flow/public/company/5.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/dot-texture.svg b/notebooks_py/flow/public/dot-texture.svg deleted file mode 100644 index 32a5a4a87..000000000 --- a/notebooks_py/flow/public/dot-texture.svg +++ /dev/null @@ -1,4 +0,0 @@ - - - - diff --git a/notebooks_py/flow/public/individual-investors/alicia-bell.jpg b/notebooks_py/flow/public/individual-investors/alicia-bell.jpg deleted file mode 100644 index f27b1478e..000000000 Binary files a/notebooks_py/flow/public/individual-investors/alicia-bell.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/individual-investors/anna-roberts.jpg b/notebooks_py/flow/public/individual-investors/anna-roberts.jpg deleted file mode 100644 index 7d1f2b38f..000000000 Binary files a/notebooks_py/flow/public/individual-investors/anna-roberts.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/individual-investors/benjamin-russel.jpg b/notebooks_py/flow/public/individual-investors/benjamin-russel.jpg deleted file mode 100644 index bed7f45e4..000000000 Binary files a/notebooks_py/flow/public/individual-investors/benjamin-russel.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/individual-investors/emma-dorsey.jpg b/notebooks_py/flow/public/individual-investors/emma-dorsey.jpg deleted file mode 100644 index e40a41c86..000000000 Binary files a/notebooks_py/flow/public/individual-investors/emma-dorsey.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/individual-investors/jenny-wilson.jpg b/notebooks_py/flow/public/individual-investors/jenny-wilson.jpg deleted file mode 100644 index eb0f0d4ff..000000000 Binary files a/notebooks_py/flow/public/individual-investors/jenny-wilson.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/individual-investors/kristin-watson.jpg b/notebooks_py/flow/public/individual-investors/kristin-watson.jpg deleted file mode 100644 index f7b82139f..000000000 Binary files a/notebooks_py/flow/public/individual-investors/kristin-watson.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/investors/deccel.svg b/notebooks_py/flow/public/investors/deccel.svg deleted file mode 100644 index 062cff0b5..000000000 --- a/notebooks_py/flow/public/investors/deccel.svg +++ /dev/null @@ -1,14 +0,0 @@ - - - - - - - - - - - - - - diff --git a/notebooks_py/flow/public/investors/remington-schwartz.svg b/notebooks_py/flow/public/investors/remington-schwartz.svg deleted file mode 100644 index 3e198e5d4..000000000 --- a/notebooks_py/flow/public/investors/remington-schwartz.svg +++ /dev/null @@ -1,27 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/notebooks_py/flow/public/linked-avatars/customer.jpg b/notebooks_py/flow/public/linked-avatars/customer.jpg deleted file mode 100644 index 4bc056eba..000000000 Binary files a/notebooks_py/flow/public/linked-avatars/customer.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/linked-avatars/manager.jpg b/notebooks_py/flow/public/linked-avatars/manager.jpg deleted file mode 100644 index bd3b263c1..000000000 Binary files a/notebooks_py/flow/public/linked-avatars/manager.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/logo-cloud/laravel.svg b/notebooks_py/flow/public/logo-cloud/laravel.svg deleted file mode 100644 index 21acbb4eb..000000000 --- a/notebooks_py/flow/public/logo-cloud/laravel.svg +++ /dev/null @@ -1 +0,0 @@ - diff --git a/notebooks_py/flow/public/logo-cloud/savvycal.svg b/notebooks_py/flow/public/logo-cloud/savvycal.svg deleted file mode 100644 index 4bf08c822..000000000 --- a/notebooks_py/flow/public/logo-cloud/savvycal.svg +++ /dev/null @@ -1 +0,0 @@ - diff --git a/notebooks_py/flow/public/logo-cloud/statamic.svg b/notebooks_py/flow/public/logo-cloud/statamic.svg deleted file mode 100644 index 38b03b61c..000000000 --- a/notebooks_py/flow/public/logo-cloud/statamic.svg +++ /dev/null @@ -1 +0,0 @@ - diff --git a/notebooks_py/flow/public/logo-cloud/transistor.svg b/notebooks_py/flow/public/logo-cloud/transistor.svg deleted file mode 100644 index cd31bf0eb..000000000 --- a/notebooks_py/flow/public/logo-cloud/transistor.svg +++ /dev/null @@ -1 +0,0 @@ - diff --git a/notebooks_py/flow/public/logo-cloud/tuple.svg b/notebooks_py/flow/public/logo-cloud/tuple.svg deleted file mode 100644 index a3bb48145..000000000 --- a/notebooks_py/flow/public/logo-cloud/tuple.svg +++ /dev/null @@ -1 +0,0 @@ - diff --git a/notebooks_py/flow/public/logo-cluster/career-builder.svg b/notebooks_py/flow/public/logo-cluster/career-builder.svg deleted file mode 100644 index 212f3803c..000000000 --- a/notebooks_py/flow/public/logo-cluster/career-builder.svg +++ /dev/null @@ -1,4 +0,0 @@ - - - - diff --git a/notebooks_py/flow/public/logo-cluster/dribbble.svg b/notebooks_py/flow/public/logo-cluster/dribbble.svg deleted file mode 100644 index 2803f4c4d..000000000 --- a/notebooks_py/flow/public/logo-cluster/dribbble.svg +++ /dev/null @@ -1,10 +0,0 @@ - - - - - - - - - - diff --git a/notebooks_py/flow/public/logo-cluster/glassdoor.svg b/notebooks_py/flow/public/logo-cluster/glassdoor.svg deleted file mode 100644 index 4d2680fd5..000000000 --- a/notebooks_py/flow/public/logo-cluster/glassdoor.svg +++ /dev/null @@ -1,12 +0,0 @@ - - - - - - - - - - - - diff --git a/notebooks_py/flow/public/logo-cluster/linkedin.svg b/notebooks_py/flow/public/logo-cluster/linkedin.svg deleted file mode 100644 index 1c5fd6180..000000000 --- a/notebooks_py/flow/public/logo-cluster/linkedin.svg +++ /dev/null @@ -1,3 +0,0 @@ - - - diff --git a/notebooks_py/flow/public/logo-cluster/upwork.svg b/notebooks_py/flow/public/logo-cluster/upwork.svg deleted file mode 100644 index 3b46beb0e..000000000 --- a/notebooks_py/flow/public/logo-cluster/upwork.svg +++ /dev/null @@ -1,3 +0,0 @@ - - - diff --git a/notebooks_py/flow/public/logo-cluster/we-work-remotely.svg b/notebooks_py/flow/public/logo-cluster/we-work-remotely.svg deleted file mode 100644 index dff6075e8..000000000 --- a/notebooks_py/flow/public/logo-cluster/we-work-remotely.svg +++ /dev/null @@ -1,5 +0,0 @@ - - - - - diff --git a/notebooks_py/flow/public/logo-timeline/adobe-creative-cloud.svg b/notebooks_py/flow/public/logo-timeline/adobe-creative-cloud.svg deleted file mode 100644 index ba3db1698..000000000 --- a/notebooks_py/flow/public/logo-timeline/adobe-creative-cloud.svg +++ /dev/null @@ -1,15 +0,0 @@ - - - - - - - - - - - - - - - diff --git a/notebooks_py/flow/public/logo-timeline/asana.svg b/notebooks_py/flow/public/logo-timeline/asana.svg deleted file mode 100644 index 5b73fc1e0..000000000 --- a/notebooks_py/flow/public/logo-timeline/asana.svg +++ /dev/null @@ -1,10 +0,0 @@ - - - - - - - - - - diff --git a/notebooks_py/flow/public/logo-timeline/basecamp.svg b/notebooks_py/flow/public/logo-timeline/basecamp.svg deleted file mode 100644 index 616f82b28..000000000 --- a/notebooks_py/flow/public/logo-timeline/basecamp.svg +++ /dev/null @@ -1,3 +0,0 @@ - - - diff --git a/notebooks_py/flow/public/logo-timeline/discord.svg b/notebooks_py/flow/public/logo-timeline/discord.svg deleted file mode 100644 index e242d4474..000000000 --- a/notebooks_py/flow/public/logo-timeline/discord.svg +++ /dev/null @@ -1,3 +0,0 @@ - - - diff --git a/notebooks_py/flow/public/logo-timeline/gmail.svg b/notebooks_py/flow/public/logo-timeline/gmail.svg deleted file mode 100644 index ca71704af..000000000 --- a/notebooks_py/flow/public/logo-timeline/gmail.svg +++ /dev/null @@ -1,7 +0,0 @@ - - - - - - - diff --git a/notebooks_py/flow/public/logo-timeline/google-calendar.svg b/notebooks_py/flow/public/logo-timeline/google-calendar.svg deleted file mode 100644 index c60be086c..000000000 --- a/notebooks_py/flow/public/logo-timeline/google-calendar.svg +++ /dev/null @@ -1,11 +0,0 @@ - - - - - - - - - - - diff --git a/notebooks_py/flow/public/logo-timeline/google-drive.svg b/notebooks_py/flow/public/logo-timeline/google-drive.svg deleted file mode 100644 index adfba0c3d..000000000 --- a/notebooks_py/flow/public/logo-timeline/google-drive.svg +++ /dev/null @@ -1,15 +0,0 @@ - - - - - - - - - - - - - - - diff --git a/notebooks_py/flow/public/logo-timeline/hubspot.svg b/notebooks_py/flow/public/logo-timeline/hubspot.svg deleted file mode 100644 index dc304bd1f..000000000 --- a/notebooks_py/flow/public/logo-timeline/hubspot.svg +++ /dev/null @@ -1,10 +0,0 @@ - - - - - - - - - - diff --git a/notebooks_py/flow/public/logo-timeline/loom.svg b/notebooks_py/flow/public/logo-timeline/loom.svg deleted file mode 100644 index 9120023c5..000000000 --- a/notebooks_py/flow/public/logo-timeline/loom.svg +++ /dev/null @@ -1,10 +0,0 @@ - - - - - - - - - - diff --git a/notebooks_py/flow/public/logo-timeline/microsoft-teams.svg b/notebooks_py/flow/public/logo-timeline/microsoft-teams.svg deleted file mode 100644 index 9cc53f938..000000000 --- a/notebooks_py/flow/public/logo-timeline/microsoft-teams.svg +++ /dev/null @@ -1,28 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/notebooks_py/flow/public/logo-timeline/slack.svg b/notebooks_py/flow/public/logo-timeline/slack.svg deleted file mode 100644 index 2eb4915d8..000000000 --- a/notebooks_py/flow/public/logo-timeline/slack.svg +++ /dev/null @@ -1,13 +0,0 @@ - - - - - - - - - - - - - diff --git a/notebooks_py/flow/public/logo-timeline/zoom.svg b/notebooks_py/flow/public/logo-timeline/zoom.svg deleted file mode 100644 index 71b245a67..000000000 --- a/notebooks_py/flow/public/logo-timeline/zoom.svg +++ /dev/null @@ -1,4 +0,0 @@ - - - - diff --git a/notebooks_py/flow/public/map.png b/notebooks_py/flow/public/map.png deleted file mode 100644 index 8e477c8c6..000000000 Binary files a/notebooks_py/flow/public/map.png and /dev/null differ diff --git a/notebooks_py/flow/public/map/1.jpg b/notebooks_py/flow/public/map/1.jpg deleted file mode 100644 index ba0b360cf..000000000 Binary files a/notebooks_py/flow/public/map/1.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/map/2.jpg b/notebooks_py/flow/public/map/2.jpg deleted file mode 100644 index 1b255d73c..000000000 Binary files a/notebooks_py/flow/public/map/2.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/map/3.jpg b/notebooks_py/flow/public/map/3.jpg deleted file mode 100644 index a8e58c598..000000000 Binary files a/notebooks_py/flow/public/map/3.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/map/4.jpg b/notebooks_py/flow/public/map/4.jpg deleted file mode 100644 index ec51108a7..000000000 Binary files a/notebooks_py/flow/public/map/4.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/map/5.jpg b/notebooks_py/flow/public/map/5.jpg deleted file mode 100644 index 1a2b664d5..000000000 Binary files a/notebooks_py/flow/public/map/5.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/screenshots/app.png b/notebooks_py/flow/public/screenshots/app.png deleted file mode 100644 index 1ef25ddb5..000000000 Binary files a/notebooks_py/flow/public/screenshots/app.png and /dev/null differ diff --git a/notebooks_py/flow/public/screenshots/competitors.png b/notebooks_py/flow/public/screenshots/competitors.png deleted file mode 100644 index 22e062e7d..000000000 Binary files a/notebooks_py/flow/public/screenshots/competitors.png and /dev/null differ diff --git a/notebooks_py/flow/public/screenshots/engagement.png b/notebooks_py/flow/public/screenshots/engagement.png deleted file mode 100644 index cbfb183de..000000000 Binary files a/notebooks_py/flow/public/screenshots/engagement.png and /dev/null differ diff --git a/notebooks_py/flow/public/screenshots/networking.png b/notebooks_py/flow/public/screenshots/networking.png deleted file mode 100644 index 910465a99..000000000 Binary files a/notebooks_py/flow/public/screenshots/networking.png and /dev/null differ diff --git a/notebooks_py/flow/public/screenshots/profile.png b/notebooks_py/flow/public/screenshots/profile.png deleted file mode 100644 index 76cd48bfa..000000000 Binary files a/notebooks_py/flow/public/screenshots/profile.png and /dev/null differ diff --git a/notebooks_py/flow/public/team/celeste-vandermark.jpg b/notebooks_py/flow/public/team/celeste-vandermark.jpg deleted file mode 100644 index 6a7b6ac75..000000000 Binary files a/notebooks_py/flow/public/team/celeste-vandermark.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/team/courtney-henry.jpg b/notebooks_py/flow/public/team/courtney-henry.jpg deleted file mode 100644 index 2e5065b19..000000000 Binary files a/notebooks_py/flow/public/team/courtney-henry.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/team/dries-vincent.jpg b/notebooks_py/flow/public/team/dries-vincent.jpg deleted file mode 100644 index b342899aa..000000000 Binary files a/notebooks_py/flow/public/team/dries-vincent.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/team/emily-selman.jpg b/notebooks_py/flow/public/team/emily-selman.jpg deleted file mode 100644 index 90bcb2b4a..000000000 Binary files a/notebooks_py/flow/public/team/emily-selman.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/team/leonard-krasner.jpg b/notebooks_py/flow/public/team/leonard-krasner.jpg deleted file mode 100644 index 39c6993e0..000000000 Binary files a/notebooks_py/flow/public/team/leonard-krasner.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/team/marcus-eldridge.jpg b/notebooks_py/flow/public/team/marcus-eldridge.jpg deleted file mode 100644 index f1c2e0d95..000000000 Binary files a/notebooks_py/flow/public/team/marcus-eldridge.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/team/michael-foster.jpg b/notebooks_py/flow/public/team/michael-foster.jpg deleted file mode 100644 index bd11c2867..000000000 Binary files a/notebooks_py/flow/public/team/michael-foster.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/team/nolan-sheffield.jpg b/notebooks_py/flow/public/team/nolan-sheffield.jpg deleted file mode 100644 index 7d4517ac0..000000000 Binary files a/notebooks_py/flow/public/team/nolan-sheffield.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/team/whitney-francis.jpg b/notebooks_py/flow/public/team/whitney-francis.jpg deleted file mode 100644 index b2818e360..000000000 Binary files a/notebooks_py/flow/public/team/whitney-francis.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/testimonials/amy-chase.jpg b/notebooks_py/flow/public/testimonials/amy-chase.jpg deleted file mode 100644 index ec117dc85..000000000 Binary files a/notebooks_py/flow/public/testimonials/amy-chase.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/testimonials/conor-neville.jpg b/notebooks_py/flow/public/testimonials/conor-neville.jpg deleted file mode 100644 index 5d545015e..000000000 Binary files a/notebooks_py/flow/public/testimonials/conor-neville.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/testimonials/dillon-lenora.jpg b/notebooks_py/flow/public/testimonials/dillon-lenora.jpg deleted file mode 100644 index 61393338b..000000000 Binary files a/notebooks_py/flow/public/testimonials/dillon-lenora.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/testimonials/harriet-arron.jpg b/notebooks_py/flow/public/testimonials/harriet-arron.jpg deleted file mode 100644 index 2b9df4dd1..000000000 Binary files a/notebooks_py/flow/public/testimonials/harriet-arron.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/testimonials/tina-yards.jpg b/notebooks_py/flow/public/testimonials/tina-yards.jpg deleted file mode 100644 index 280edc9fc..000000000 Binary files a/notebooks_py/flow/public/testimonials/tina-yards.jpg and /dev/null differ diff --git a/notebooks_py/flow/public/testimonials/veronica-winton.jpg b/notebooks_py/flow/public/testimonials/veronica-winton.jpg deleted file mode 100644 index 2ce0eb039..000000000 Binary files a/notebooks_py/flow/public/testimonials/veronica-winton.jpg and /dev/null differ diff --git a/readme.md b/readme.md index 631b2e752..556f7f2a9 100644 --- a/readme.md +++ b/readme.md @@ -47,5 +47,9 @@ https://youtu.be/y4ajXJ3nj1Q?si=CMpEXr9DBH86-9jG https://tevisthompson.com/saving-zelda/ -1e14 = 100 trillion -> (100 million) - (millionaires) - who can create at least 10 AI-Proof jobs that create the future - Artist-Scientist Wizards who use dynamicland.org - the best invention of all time. -odysseys purpose is to implment dynamicland.org's research into every neighborhood in the world. \ No newline at end of file +1e14 = 100 trillion -> (100 million) - (millionaires) - who can create at least 10 AI-Proof jobs that create the future - Artist-Scientist Wizards who use dynamicland.org - the best invention of all time. (zoox called me an inventor and gave me an inventor paper weight and a computer vision patent.) + +odysseys purpose is to implment dynamicland.org's research into every neighborhood in the world, turning everyone into a millionaire artist scientist wizard in 5 years. + + +Jennifer Doudna = princess bubblegum \ No newline at end of file diff --git a/replit.nix b/replit.nix new file mode 100644 index 000000000..d927087a7 --- /dev/null +++ b/replit.nix @@ -0,0 +1,5 @@ +{ pkgs }: { + deps = [ + pkgs.bun + ]; +} \ No newline at end of file diff --git a/scripts/homelab_status_page.sh b/scripts/homelab_status_page.sh index e75fc6563..77e72580e 100755 --- a/scripts/homelab_status_page.sh +++ b/scripts/homelab_status_page.sh @@ -11,37 +11,20 @@ if [ -z "$COMMAND" ]; then fi if [ "$COMMAND" = "kentbeck" ] || [ "$COMMAND" = "build.sh" ]; then - #$ROOT/build.sh "$@" - echo "Running run.sh..." bun run ~/homelab_status_page/web-ui/js/helpers/Kent_Beck_robusteness.js - echo "Running run.sh..." -elif [ "$COMMAND" = "run" ] || [ "$COMMAND" = "run.sh" ]; then - #$ROOT/run.sh "$@" - echo "Running run.sh..." -elif [ "$COMMAND" = "list" ]; then - # $ROOT/build.sh --list "$@" - echo "Running run.sh..." - +elif [ "$COMMAND" = "web" ] || [ "$COMMAND" = "web-ui" ]; then + bun run ~/homelab_status_page/web-ui/my-app/src/index.js +elif [ "$COMMAND" = "web-docker" ]; then + docker run -p 3000:3000 homelab_status_page elif [ "$COMMAND" = "show" ]; then #$ROOT/build.sh --show "$@" echo "Running run.sh..." - -elif [ "$COMMAND" = "install" ]; then - #$ROOT/autotag "$@" - bash ~/homelab_status_page/scripts/install.sh - elif [ "$COMMAND" = "update" ]; then - # cd $ROOT - #git pull - echo "Running run.sh..." - -elif [ "$COMMAND" = "root" ]; then - echo $ROOT + echo "Running install.sh..." elif [ "$COMMAND" = "data" ]; then echo $ROOT/data -elif [ "$COMMAND" = "install" ]; then +elif [ "$COMMAND" = "replit" ]; then echo "Running install..." - # Add your install command logic here elif [ "$COMMAND" = "bootstrap" ]; then echo "Running bootstrap..." bash ~/homelab_status_page/scripts/_bootstrap.sh @@ -49,26 +32,23 @@ elif [ "$COMMAND" = "restart-blog" ]; then echo "Running restart-blog..." #bash ~/homelab_status_page/scripts/_bootstrap.sh bash ~/homelab_status_page/scripts/restart-blog.sh - - # Add your restart-blog command logic here else - - #echo 'default PARAM - no toher requests currently - restarting blag for fun !?!?!' - # echo 'homelab_status_page > Invalid command' - # echo '' - # echo ' * build [PACKAGES]' - # echo ' * run OPTIONS [CONTAINER:TAG] CMD' - # echo ' * list [PACKAGES|*' - # echo ' * show [PACKAGES]*' - # echo ' * autotag [CONTAINER]' - # echo ' * update (runs git pull)' - # echo ' * root (prints repo path)' - # echo ' * data (prints data path)' - # echo ' * install' - # echo ' * bootstrap' - # echo ' * restart-blog (default)' - # echo '' - # echo 'Run "jetson-containers --help" for more info.' - + echo 'default PARAM - no toher requests currently - restarting blag for fun !?!?!' + echo 'homelab > Invalid command' + echo '' + echo ' * build [PACKAGES]' + echo ' * run OPTIONS [CONTAINER:TAG] CMD' + echo ' * list [PACKAGES|*' + echo ' * show [PACKAGES]*' + echo ' * autotag [CONTAINER]' + echo ' * update (runs git pull)' + echo ' * root (prints repo path)' + echo ' * data (prints data path)' + echo ' * install' + echo ' * bootstrap' + echo ' * restart-blog (default)' + echo '' + echo 'Run "homelab --help" for more info.' + echo "this command does not exist, should i create a magic spell we can work and cast together oner? " exit 1 fi diff --git a/scripts/init_containers.sh b/scripts/init_containers.sh index 2cde2589c..59f1da278 100644 --- a/scripts/init_containers.sh +++ b/scripts/init_containers.sh @@ -1,4 +1,42 @@ #!/bin/bash +# run this on jetson to install hashirama platform + +#sudo apt-get install -y nvidia-container-toolkit +git clone --depth 1 https://github.com/doomemacs/doomemacs ~/.config/emacs +~/.config/emacs/bin/doom install +#sudo apt-get install -y nvidia-container-toolkit +sudo ln -s /home/adnan/hashirama/infra/caddy/Caddyfile /etc/caddy/Caddyfile +#immich + + +#etc + + +#tailscale + +#nix ?? + + +#mkdir unreal +# wget -> +# unzip +# + + +# bootstraps a server like tailscale, ollama, etc +# +# +# + + + +curl -fsSL https://ollama.com/install.sh | sh + +curl -fsSL https://tailscale.com/install.sh | sh + + +# docker pull ollama/ollama +# docker run -it --gpus all ollama/ollama # List of package paths packages=( diff --git a/scripts/install-observable-server.sh b/scripts/install-observable-server.sh deleted file mode 100755 index a3c1a7c5b..000000000 --- a/scripts/install-observable-server.sh +++ /dev/null @@ -1,38 +0,0 @@ -# run this on jetson to install hashirama platform - -#sudo apt-get install -y nvidia-container-toolkit -git clone --depth 1 https://github.com/doomemacs/doomemacs ~/.config/emacs -~/.config/emacs/bin/doom install -#sudo apt-get install -y nvidia-container-toolkit -sudo ln -s /home/adnan/hashirama/infra/caddy/Caddyfile /etc/caddy/Caddyfile -#immich - - -#etc - - -#tailscale - -#nix ?? - - -#mkdir unreal -# wget -> -# unzip -# - - -# bootstraps a server like tailscale, ollama, etc -# -# -# - - - -curl -fsSL https://ollama.com/install.sh | sh - -curl -fsSL https://tailscale.com/install.sh | sh - - -# docker pull ollama/ollama -# docker run -it --gpus all ollama/ollama diff --git a/scripts/print_markdown_table.sh b/scripts/print_markdown_table.sh deleted file mode 100755 index 0e2612c8c..000000000 --- a/scripts/print_markdown_table.sh +++ /dev/null @@ -1,60 +0,0 @@ -#!/bin/bash - -# Output markdown or CSV format -OUTPUT_TYPE=${1:-"markdown"} # default to markdown if no argument passed - -# Define file name -OUTPUT_FILE="components_specs.$( [ "$OUTPUT_TYPE" == "csv" ] && echo csv || echo md)" - -# Data of components -declare -A COMPONENTS=( - ["FLOPS"]="738 FP16 TFLOPS,991 FP16 TFLOPS,1.36 FP16 PFLOPS" - ["GPU Model"]="6x 7900XTX,6x 4090,8x 4090" - ["GPU RAM"]="144 GB,192 GB," - ["GPU RAM bandwidth"]="5760 GB/s,6050 GB/s,8064 GB/s" - ["GPU link bandwidth"]="Full PCIe 4.0 x16 (64 GB/s),Full PCIe 4.0 x16 (64 GB/s),Full PCIe 4.0 x16 (64 GB/s)" - ["CPU"]="32 core AMD EPYC,2x AMD GENOA," - ["System RAM"]="128 GB,384 GB," - ["System RAM bandwidth"]="204.8 GB/s,921.6 GB/s," - ["Disk size"]="4 TB raid array + 1 TB boot,1 TB boot," - ["Disk read bandwidth"]="28.7 GB/s,Use network," - ["Networking"]="2x 1 GbE + open OCP3.0 slot (up to 200 GbE),2x open PCIe5 x16," - ["Noise"]="< 50 dB, 31 low speed fans,Loud," - ["Power Supply"]="2x 1600W, 100V~240V,4x 2000W, 200V+," - ["BMC"]="AST2500,AST2600," - ["Operating System"]="Ubuntu 22.04,Ubuntu 22.04," - ["Dimensions"]="12U, 16.25\" deep, 90 lbs,4U, 31\" D, 88 lbs," - ["Rack?"]="Freestanding or rack mount,Supermicro rails," - ["Driver Quality"]="Mediocre,Great," - ["SHIPPING"]="BUY NOW FOR $15,000,BUY NOW FOR $25,000,PREORDER $40,000" -) - -# Function to print markdown -function print_markdown { - echo "| Component | Red | Green | Pro |" > $OUTPUT_FILE - echo "|--------------------------|-----------------------|-----------------------|----------------------|" >> $OUTPUT_FILE - - for component in "${!COMPONENTS[@]}"; do - IFS=',' read -ra SPECS <<< "${COMPONENTS[$component]}" - echo "| **$component** | ${SPECS[0]} | ${SPECS[1]} | ${SPECS[2]:-} |" >> $OUTPUT_FILE - done -} - -# Function to print CSV -function print_csv { - echo "Component,Red,Green,Pro" > $OUTPUT_FILE - - for component in "${!COMPONENTS[@]}"; do - IFS=',' read -ra SPECS <<< "${COMPONENTS[$component]}" - echo "$component,${SPECS[0]},${SPECS[1]},${SPECS[2]:-}" >> $OUTPUT_FILE - done -} - -# Choose output format -if [ "$OUTPUT_TYPE" == "csv" ]; then - print_csv -else - print_markdown -fi - -echo "Components and specs saved to $OUTPUT_FILE" diff --git a/scripts/remove_comments.py b/scripts/remove_comments.py deleted file mode 100755 index 180eebc2e..000000000 --- a/scripts/remove_comments.py +++ /dev/null @@ -1,32 +0,0 @@ -import re - -def extract_multiline_comments(file_path): - """Extract multi-line comments from the file.""" - with open(file_path, 'r') as file: - content = file.read() - - # Regex to match multi-line comments in Python (triple quotes) - multiline_comments = re.findall(r'\"\"\"(.*?)\"\"\"|\'\'\'(.*?)\'\'\'', content, re.DOTALL) - - # Combine results from both double and single quote matches - comments = [comment[0] or comment[1] for comment in multiline_comments] - - return comments - -def dump_to_file(comments, output_file): - """Dump extracted comments to a text file.""" - with open(output_file, 'w') as f: - for comment in comments: - f.write(comment.strip() + '\n\n') - -if __name__ == "__main__": - source_file = 'your_source_file.py' # Replace with the source file you want to process - output_file = 'comments.txt' # File where the comments will be saved - - # Extract comments - comments = extract_multiline_comments(source_file) - - # Dump comments to file - dump_to_file(comments, output_file) - - print(f"Extracted {len(comments)} multi-line comments and saved them to {output_file}.") \ No newline at end of file diff --git a/scripts/remove_comments.sh b/scripts/remove_comments.sh deleted file mode 100755 index a2f45ba8b..000000000 --- a/scripts/remove_comments.sh +++ /dev/null @@ -1,26 +0,0 @@ -#!/bin/bash - -# Function to extract multi-line comments from a JS/JSX file -extract_multiline_comments() { - local file="$1" - # Using sed to extract content between /* and */ (JavaScript multi-line comment style) - sed -n '/\/\*/,/\*\//p' "$file" -} - -# Directory or file to process (you can modify this path) -directory="." # or you can specify a file path directly - -# Output file to dump the comments -output_file="comments.txt" - -# Clear the output file -> "$output_file" - -# Find all JS and JSX files and extract comments -find "$directory" -name "*.js" -o -name "*.jsx" | while read -r file; do - echo "Processing file: $file" - extract_multiline_comments "$file" >> "$output_file" - echo -e "\n" >> "$output_file" -done - -echo "All comments have been extracted and saved to $output_file." diff --git a/scripts/setup_sshfs.sh b/scripts/setup_sshfs.sh deleted file mode 100755 index 317ca4afc..000000000 --- a/scripts/setup_sshfs.sh +++ /dev/null @@ -1,33 +0,0 @@ -# Option 4: SSHFS with Tailscale (On-Demand Access over SSH) -# If you are comfortable with SSH, SSHFS over Tailscale allows you to mount the 4TB directory from your desktop to your MacBook, providing on-demand access. - -# Setup: -# Install SSHFS: - - -# Check if the operating system is macOS or Linux -if [[ "$OSTYPE" == "darwin"* ]]; then - brew install sshfs - # macOS specific commands can be added here -elif [[ "$OSTYPE" == "linux-gnu"* ]]; then - sudo apt install sshfs - # Linux specific commands can be added here -else - echo "Unsupported OS" -fi - -#sshfs @100.101.102.103:/path/to/4TB_folder /mnt/remote_data - - # On Linux - # On macOS -# Mount Desktop Directory on MacBook: - -# bash -# Copy code -#Access Files as Needed: Your scripts or programs can access the mounted directory, and only the files requested are transferred. - -# fusermount -u ~/remote_mount # On Linux -# umount ~/remote_mount # On macOS - -# web hooks -# observable, \ No newline at end of file diff --git a/scripts/start.sh b/scripts/start.sh deleted file mode 100755 index a4a83cc65..000000000 --- a/scripts/start.sh +++ /dev/null @@ -1,8 +0,0 @@ - - -go run main.go -bun --hot js/diffusion-democracy-observablehq.js -bun --hot js/bun-livekit-server.js -deno run --access-all js/figma-proxy-or-clone.ts - -echo "voice modulator - index APU " diff --git a/scripts/start_app.sh b/scripts/start_app.sh deleted file mode 100755 index 29e67bdf2..000000000 --- a/scripts/start_app.sh +++ /dev/null @@ -1,27 +0,0 @@ -#!/bin/bash -# "llama-backend": 7095, -# "cgi-backend": 7096, -#"flow-backend": 7097, -# "robotics-odyssey-demo": 7998, -# "main.go": 7999, - -# caddy port 80 -# "air": 8001, -# "jupyter": 8888 -# Start Node.js services -pm2 start js/llama-backend.js --name "server1" -pm2 start js/cgi-backend.js --name "server2" -pm2 start js/robotics-odyssey-demo.ts --name "server3" -# pm2 start js/server4.js --name "server4" -# pm2 start server5.js --name "server5" - -# Start Go server -#nohup go run main.go > go_server.log 2>&1 & -# switch to air later -air - -# Start Jupyter Notebook -#nohup jupyter notebook --no-browser --port=8888 > jupyter.log 2>&1 & - - -# to stop them - pm2 stop all diff --git a/scripts/symlink-all-scripts copy.sh b/scripts/symlink-all-scripts copy.sh deleted file mode 100755 index 79b616062..000000000 --- a/scripts/symlink-all-scripts copy.sh +++ /dev/null @@ -1,58 +0,0 @@ -#!/bin/bash - -# Directory containing the scripts -SCRIPT_DIR=$(dirname "$0") - -# Target directory for symlinks -TARGET_DIR="/usr/local/bin" - -# Iterate over each .sh file in the script directory -for script in "$SCRIPT_DIR"/*.sh; do - # Ensure the script is executable - chmod +x "$script" - - # Get the base name of the script - script_name=$(basename "$script") - - # Create a symlink in the target directory - ln -sf "$script" "$TARGET_DIR/$script_name" - - # Create a symlink in the target directory - if ln -sf "$script" "$TARGET_DIR/$script_name"; then - echo "Symlink created for $script_name" - else - echo "Failed to create symlink for $script_name" >&2 - fi -done - -echo "Symlinks created for all scripts in $SCRIPT_DIR to $TARGET_DIR" - - - -alias yarn="bun" -alias npm="bun" -alias deno="deno --accept-all" -add_alias() { - if [ $# -ne 2 ]; then - echo "Usage: add_alias " - return 1 - fi - - local alias_name=$1 - local command=$2 - - # Check if the alias already exists - if alias "$alias_name" &>/dev/null; then - echo "Alias '$alias_name' already exists." - return 1 - fi - - # Add the alias to the file - echo "alias $alias_name=\"$command\"" >> "$0" - echo "Alias '$alias_name' added successfully." -} - -# Example usage: -# add_alias myalias "mycommand" - - diff --git a/scripts/symlink-all-scripts.sh b/scripts/symlink-all-scripts.sh deleted file mode 100755 index 79b616062..000000000 --- a/scripts/symlink-all-scripts.sh +++ /dev/null @@ -1,58 +0,0 @@ -#!/bin/bash - -# Directory containing the scripts -SCRIPT_DIR=$(dirname "$0") - -# Target directory for symlinks -TARGET_DIR="/usr/local/bin" - -# Iterate over each .sh file in the script directory -for script in "$SCRIPT_DIR"/*.sh; do - # Ensure the script is executable - chmod +x "$script" - - # Get the base name of the script - script_name=$(basename "$script") - - # Create a symlink in the target directory - ln -sf "$script" "$TARGET_DIR/$script_name" - - # Create a symlink in the target directory - if ln -sf "$script" "$TARGET_DIR/$script_name"; then - echo "Symlink created for $script_name" - else - echo "Failed to create symlink for $script_name" >&2 - fi -done - -echo "Symlinks created for all scripts in $SCRIPT_DIR to $TARGET_DIR" - - - -alias yarn="bun" -alias npm="bun" -alias deno="deno --accept-all" -add_alias() { - if [ $# -ne 2 ]; then - echo "Usage: add_alias " - return 1 - fi - - local alias_name=$1 - local command=$2 - - # Check if the alias already exists - if alias "$alias_name" &>/dev/null; then - echo "Alias '$alias_name' already exists." - return 1 - fi - - # Add the alias to the file - echo "alias $alias_name=\"$command\"" >> "$0" - echo "Alias '$alias_name' added successfully." -} - -# Example usage: -# add_alias myalias "mycommand" - -